diff --git a/moto/dynamodb2/models.py b/moto/dynamodb2/models.py index aab9f8397..99ace6da0 100644 --- a/moto/dynamodb2/models.py +++ b/moto/dynamodb2/models.py @@ -500,12 +500,20 @@ class DynamoDBBackend(BaseBackend): def update_item(self, table_name, key, update_expression, attribute_updates): table = self.get_table(table_name) - if table.hash_key_attr in key: - # Sometimes the key is wrapped in a dict with the key name - key = key[table.hash_key_attr] + if all([table.hash_key_attr in key, table.range_key_attr in key]): + # Covers cases where table has hash and range keys, ``key`` param will be a dict + hash_value = DynamoType(key[table.hash_key_attr]) + range_value = DynamoType(key[table.range_key_attr]) + elif table.hash_key_attr in key: + # Covers tables that have a range key where ``key`` param is a dict + hash_value = DynamoType(key[table.hash_key_attr]) + range_value = None + else: + # Covers other cases + hash_value = DynamoType(key) + range_value = None - hash_value = DynamoType(key) - item = table.get_item(hash_value) + item = table.get_item(hash_value, range_value) if update_expression: item.update(update_expression) else: diff --git a/tests/test_dynamodb2/test_dynamodb_table_with_range_key.py b/tests/test_dynamodb2/test_dynamodb_table_with_range_key.py index 74f7615c0..282599bad 100644 --- a/tests/test_dynamodb2/test_dynamodb_table_with_range_key.py +++ b/tests/test_dynamodb2/test_dynamodb_table_with_range_key.py @@ -143,6 +143,36 @@ def test_item_add_and_describe_and_update(): }) +@requires_boto_gte("2.9") +@mock_dynamodb2 +def test_item_partial_save(): + table = create_table() + + data = { + 'forum_name': 'LOLCat Forum', + 'subject': 'The LOLz', + 'Body': 'http://url_to_lolcat.gif', + 'SentBy': 'User A', + } + + table.put_item(data=data) + returned_item = table.get_item(forum_name="LOLCat Forum", subject='The LOLz') + + returned_item['SentBy'] = 'User B' + returned_item.partial_save() + + returned_item = table.get_item( + forum_name='LOLCat Forum', + subject='The LOLz' + ) + dict(returned_item).should.equal({ + 'forum_name': 'LOLCat Forum', + 'subject': 'The LOLz', + 'Body': 'http://url_to_lolcat.gif', + 'SentBy': 'User B', + }) + + @requires_boto_gte("2.9") @mock_dynamodb2 def test_item_put_without_table():