Ensure binary attrs are parsed correctly

Author: garrettheel

pynamodb/connection/base.py

@@ -1,11 +1,7 @@
 """
 Lowest level connection
 """
-import json
 import logging
-import random
-import sys
-import time
 import uuid
 from base64 import b64decode
 from threading import local
@@ -15,7 +11,6 @@
 import botocore.exceptions
 from botocore.awsrequest import AWSPreparedRequest, create_request_object
 from botocore.client import ClientError
-from botocore.hooks import first_non_none_response
 from botocore.exceptions import BotoCoreError
 from botocore.session import get_session
 
@@ -30,18 +25,14 @@
     PUT_ITEM, SELECT, LIMIT, QUERY, SCAN, ITEM, LOCAL_SECONDARY_INDEXES,
     KEYS, KEY, SEGMENT, TOTAL_SEGMENTS, CREATE_TABLE, PROVISIONED_THROUGHPUT, READ_CAPACITY_UNITS,
     WRITE_CAPACITY_UNITS, GLOBAL_SECONDARY_INDEXES, PROJECTION, EXCLUSIVE_START_TABLE_NAME, TOTAL,
-    DELETE_TABLE, UPDATE_TABLE, LIST_TABLES, GLOBAL_SECONDARY_INDEX_UPDATES, ATTRIBUTES,
-    CONSUMED_CAPACITY, CAPACITY_UNITS, ATTRIBUTE_TYPES,
-    ITEMS, DEFAULT_ENCODING, BINARY, BINARY_SET, LAST_EVALUATED_KEY, RESPONSES, UNPROCESSED_KEYS,
-    UNPROCESSED_ITEMS, STREAM_SPECIFICATION, STREAM_VIEW_TYPE, STREAM_ENABLED,
-    EXPRESSION_ATTRIBUTE_NAMES, EXPRESSION_ATTRIBUTE_VALUES,
-    CONDITION_EXPRESSION, FILTER_EXPRESSION,
+    DELETE_TABLE, UPDATE_TABLE, LIST_TABLES, GLOBAL_SECONDARY_INDEX_UPDATES, CONSUMED_CAPACITY, CAPACITY_UNITS,
+    ATTRIBUTE_TYPES, DEFAULT_ENCODING, BINARY, BINARY_SET, STREAM_SPECIFICATION, STREAM_VIEW_TYPE, STREAM_ENABLED,
+    EXPRESSION_ATTRIBUTE_NAMES, EXPRESSION_ATTRIBUTE_VALUES, CONDITION_EXPRESSION, FILTER_EXPRESSION,
     TRANSACT_WRITE_ITEMS, TRANSACT_GET_ITEMS, CLIENT_REQUEST_TOKEN, TRANSACT_ITEMS, TRANSACT_CONDITION_CHECK,
     TRANSACT_GET, TRANSACT_PUT, TRANSACT_DELETE, TRANSACT_UPDATE, UPDATE_EXPRESSION,
     RETURN_VALUES_ON_CONDITION_FAILURE_VALUES, RETURN_VALUES_ON_CONDITION_FAILURE,
     AVAILABLE_BILLING_MODES, DEFAULT_BILLING_MODE, BILLING_MODE, PAY_PER_REQUEST_BILLING_MODE,
-    PROVISIONED_BILLING_MODE,
-    TIME_TO_LIVE_SPECIFICATION, ENABLED, UPDATE_TIME_TO_LIVE, TAGS, VALUE
+    PROVISIONED_BILLING_MODE, TIME_TO_LIVE_SPECIFICATION, ENABLED, UPDATE_TIME_TO_LIVE, TAGS, VALUE
 )
 from pynamodb.exceptions import (
     TableError, QueryError, PutError, DeleteError, UpdateError, GetError, ScanError, TableDoesNotExist,
@@ -351,46 +342,6 @@ def send_pre_boto_callback(self, operation_name, req_uuid, table_name):
     def _make_api_call(self, operation_name: str, operation_kwargs: Dict, settings: OperationSettings = OperationSettings.default) -> Dict:
         return self.client._make_api_call(operation_name, operation_kwargs)
 
-    @staticmethod
-    def _handle_binary_attributes(data):
-        """ Simulate botocore's binary attribute handling """
-        if ITEM in data:
-            for attr in data[ITEM].values():
-                _convert_binary(attr)
-        if ITEMS in data:
-            for item in data[ITEMS]:
-                for attr in item.values():
-                    _convert_binary(attr)
-        if RESPONSES in data:
-            if isinstance(data[RESPONSES], list):
-                for item in data[RESPONSES]:
-                    for attr in item.values():
-                        _convert_binary(attr)
-            else:
-                for item_list in data[RESPONSES].values():
-                    for item in item_list:
-                        for attr in item.values():
-                            _convert_binary(attr)
-        if LAST_EVALUATED_KEY in data:
-            for attr in data[LAST_EVALUATED_KEY].values():
-                _convert_binary(attr)
-        if UNPROCESSED_KEYS in data:
-            for table_data in data[UNPROCESSED_KEYS].values():
-                for item in table_data[KEYS]:
-                    for attr in item.values():
-                        _convert_binary(attr)
-        if UNPROCESSED_ITEMS in data:
-            for table_unprocessed_requests in data[UNPROCESSED_ITEMS].values():
-                for request in table_unprocessed_requests:
-                    for item_mapping in request.values():
-                        for item in item_mapping.values():
-                            for attr in item.values():
-                                _convert_binary(attr)
-        if ATTRIBUTES in data:
-            for attr in data[ATTRIBUTES].values():
-                _convert_binary(attr)
-        return data
-
     @property
     def session(self) -> botocore.session.Session:
         """

tests/test_base_connection.py

@@ -5,10 +5,9 @@
 import json
 from unittest import mock
 from unittest.mock import patch
-from unittest.mock import ANY
 
 import botocore.exceptions
-from botocore.awsrequest import AWSPreparedRequest, AWSRequest, AWSResponse
+from botocore.awsrequest import AWSResponse
 from botocore.client import ClientError
 from botocore.exceptions import BotoCoreError
 
@@ -25,7 +24,6 @@
 from pynamodb.expressions.update import SetAction
 from pynamodb.settings import OperationSettings
 from .data import DESCRIBE_TABLE_DATA, GET_ITEM_DATA, LIST_TABLE_DATA
-from .deep_eq import deep_eq
 
 PATCH_METHOD = 'pynamodb.connection.Connection._make_api_call'
 TEST_TABLE_NAME = 'ci-table'
@@ -1512,74 +1510,41 @@ def test_connection_make_api_call__throws_retry_disabled(send_mock):
     assert len(send_mock.mock_calls) == 1
 
 
-def test_connection_handle_binary_attributes__unprocessed_items():
+@mock.patch('botocore.httpsession.URLLib3Session.send')
+def test_connection_make_api_call__binary_attributes(send_mock):
     binary_blob = b'\x00\xFF\x00\xFF'
-
-    unprocessed_items = []
-    for _ in range(0, 5):
-        unprocessed_items.append({
-            'PutRequest': {
-                'Item': {
-                    'name': {STRING: 'daniel'},
-                    'picture': {BINARY: base64.b64encode(binary_blob).decode(DEFAULT_ENCODING)}
-                }
-            }
-        })
-
-    expected_unprocessed_items = []
-    for _ in range(0, 5):
-        expected_unprocessed_items.append({
-            'PutRequest': {
-                'Item': {
-                    'name': {STRING: 'daniel'},
-                    'picture': {BINARY: binary_blob}
+    resp_text = json.dumps({
+        UNPROCESSED_ITEMS: {
+            'someTable': [{
+                'PutRequest': {
+                    'Item': {
+                        'name': {STRING: 'daniel'},
+                        'picture': {BINARY: base64.b64encode(binary_blob).decode(DEFAULT_ENCODING)},
+                    }
                 }
-            }
-        })
+            }],
+        }
+    })
 
-    assert (
-        Connection._handle_binary_attributes({UNPROCESSED_ITEMS: {'someTable': unprocessed_items}}) ==
-        {UNPROCESSED_ITEMS: {'someTable': expected_unprocessed_items}}
+    resp = mock.Mock(
+        spec=AWSResponse,
+        status_code=200,
+        headers={},
+        content=resp_text.encode(),
     )
 
+    send_mock.return_value = resp
 
-def test_connection_handle_binary_attributes__unprocessed_keys():
-    binary_blob = b'\x00\xFF\x00\xFF'
-    unprocessed_keys = {
-        'UnprocessedKeys': {
-            'MyTable': {
-                'AttributesToGet': ['ForumName'],
-                'Keys': [
-                    {
-                        'ForumName': {'S': 'FooForum'},
-                        'Subject': {'B': base64.b64encode(binary_blob).decode(DEFAULT_ENCODING)}
-                    },
-                    {
-                        'ForumName': {'S': 'FooForum'},
-                        'Subject': {'S': 'thread-1'}
-                    }
-                ],
-                'ConsistentRead': False
-            },
-            'MyOtherTable': {
-                'AttributesToGet': ['ForumName'],
-                'Keys': [
-                    {
-                        'ForumName': {'S': 'FooForum'},
-                        'Subject': {'B': base64.b64encode(binary_blob).decode(DEFAULT_ENCODING)}
-                    },
-                    {
-                        'ForumName': {'S': 'FooForum'},
-                        'Subject': {'S': 'thread-1'}
-                    }
-                ],
-                'ConsistentRead': False
+    resp = Connection()._make_api_call('BatchWriteItem', {})
+
+    assert resp['UnprocessedItems']['someTable'] == [{
+        'PutRequest': {
+            'Item': {
+                'name': {STRING: 'daniel'},
+                'picture': {BINARY: binary_blob}
             }
         }
-    }
-    data = Connection._handle_binary_attributes(unprocessed_keys)
-    assert data['UnprocessedKeys']['MyTable']['Keys'][0]['Subject']['B'] == binary_blob
-    assert data['UnprocessedKeys']['MyOtherTable']['Keys'][0]['Subject']['B'] == binary_blob
+    }]
 
 
 def test_connection_update_time_to_live__fail():