Merge pull request #504 from pcraciunoiu/feature/update-item-types-and-test

Support for other UPDATE DynamoTypes and query Select=COUNT

moto/dynamodb2/models.py

@@ -124,11 +124,20 @@ class Item(object):
     def update_with_attribute_updates(self, attribute_updates):
         for attribute_name, update_action in attribute_updates.items():
             action = update_action['Action']
+            if action == 'DELETE' and not 'Value' in update_action:
+                del self.attrs[attribute_name]
+                continue
             new_value = list(update_action['Value'].values())[0]
             if action == 'PUT':
                 # TODO deal with other types
                 if isinstance(new_value, list) or isinstance(new_value, set):
                     self.attrs[attribute_name] = DynamoType({"SS": new_value})
+                elif isinstance(new_value, dict):
+                    self.attrs[attribute_name] = DynamoType({"M": new_value})
+                elif update_action['Value'].keys() == ['N']:
+                    self.attrs[attribute_name] = DynamoType({"N": new_value})
+                elif update_action['Value'].keys() == ['NULL']:
+                    del self.attrs[attribute_name]
                 else:
                     self.attrs[attribute_name] = DynamoType({"S": new_value})
 

moto/dynamodb2/responses.py

@@ -330,9 +330,10 @@ class DynamoHandler(BaseResponse):
 
         result = {
             "Count": len(items),
-            "Items": [item.attrs for item in items],
             "ConsumedCapacityUnits": 1,
         }
+        if self.body.get('Select', '').upper() != 'COUNT':
+            result["Items"] = [item.attrs for item in items]
 
         # Implement this when we do pagination
         # if not last_page:

tests/test_dynamodb2/test_dynamodb_table_with_range_key.py

@@ -735,9 +735,7 @@ def test_boto3_conditions():
     results['Count'].should.equal(1)
 
 
-
+def _create_table_with_range_key():
-@mock_dynamodb2
-def test_boto3_query_gsi_range_comparison():
     dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
 
     # Create the DynamoDB table.
@@ -788,7 +786,52 @@ def test_boto3_query_gsi_range_comparison():
             'WriteCapacityUnits': 5
         }
     )
-    table = dynamodb.Table('users')
+    return dynamodb.Table('users')
+
+
+@mock_dynamodb2
+def test_update_item_range_key_set():
+    table = _create_table_with_range_key()
+    table.put_item(Item={
+        'forum_name': 'the-key',
+        'subject': '123',
+        'username': 'johndoe',
+        'created': Decimal('3'),
+    })
+
+    item_key = {'forum_name': 'the-key', 'subject': '123'}
+    table.update_item(
+        Key=item_key,
+        AttributeUpdates={
+            'username': {
+                'Action': u'PUT',
+                'Value': 'johndoe2'
+            },
+            'created': {
+                'Action': u'PUT',
+                'Value': Decimal('4'),
+            },
+            'mapfield': {
+                'Action': u'PUT',
+                'Value': {'key': 'value'},
+            }
+        },
+    )
+
+    returned_item = dict((k, str(v) if isinstance(v, Decimal) else v)
+                         for k, v in table.get_item(Key=item_key)['Item'].items())
+    dict(returned_item).should.equal({
+        'username': "johndoe2",
+        'forum_name': 'the-key',
+        'subject': '123',
+        'created': '4',
+        'mapfield': {'key': 'value'},
+    })
+
+
+@mock_dynamodb2
+def test_boto3_query_gsi_range_comparison():
+    table = _create_table_with_range_key()
 
     table.put_item(Item={
         'forum_name': 'the-key',