make update_item respect attr_name differences (#160)

pynamodb · Sep 20, 2016 · 4ab34ca · 4ab34ca
1 parent 5229759
commit 4ab34ca
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 2 deletions.
diff --git a/pynamodb/models.py b/pynamodb/models.py
@@ -349,12 +349,12 @@ def update_item(self, attribute, value=None, action=None, conditional_operator=N
         if len(expected_values):
             kwargs.update(expected=self._build_expected_values(expected_values, UPDATE_FILTER_OPERATOR_MAP))
         kwargs[pythonic(ATTR_UPDATES)] = {
-            attribute: {
+            attribute_cls.attr_name: {
                 ACTION: action.upper() if action else None,
             }
         }
         if action is not None and action.upper() != DELETE:
-            kwargs[pythonic(ATTR_UPDATES)][attribute][VALUE] = {ATTR_TYPE_MAP[attribute_cls.attr_type]: value}
+            kwargs[pythonic(ATTR_UPDATES)][attribute_cls.attr_name][VALUE] = {ATTR_TYPE_MAP[attribute_cls.attr_type]: value}
         kwargs[pythonic(RETURN_VALUES)] = ALL_NEW
         kwargs.update(conditional_operator=conditional_operator)
         data = self._get_connection().update_item(

diff --git a/pynamodb/tests/test_model.py b/pynamodb/tests/test_model.py
@@ -900,6 +900,29 @@ def test_update_item(self):
             }
             deep_eq(args, params, _assert=True)
 
+        with patch(PATCH_METHOD) as req:
+            item.update_item('custom_aliases', set(['lita']), action='add')
+            args = req.call_args[0][1]
+            params = {
+                'TableName': 'SimpleModel',
+                'ReturnValues': 'ALL_NEW',
+                'Key': {
+                    'user_name': {
+                        'S': 'foo'
+                    }
+                },
+                'AttributeUpdates': {
+                    'aliases': {
+                        'Action': 'ADD',
+                        'Value': {
+                            'SS': set(['lita'])
+                        }
+                    }
+                },
+                'ReturnConsumedCapacity': 'TOTAL'
+            }
+            deep_eq(args, params, _assert=True)
+
     def test_save(self):
         """
         Model.save