Refactor dynamodb paginator fix handler

happy-fish · Jan 8, 2020 · 17c11ad · 17c11ad
1 parent 58622cd
commit 17c11ad
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 22 deletions.
diff --git a/awscli/clidriver.py b/awscli/clidriver.py
@@ -506,8 +506,7 @@ def __call__(self, args, parsed_globals):
             event,
             call_parameters=call_parameters,
             parsed_args=parsed_args,
-            parsed_globals=parsed_globals,
-            session=self._session,
+            parsed_globals=parsed_globals
         )
         # There are two possible values for override. It can be some type
         # of exception that will be raised if detected or it can represent

diff --git a/awscli/customizations/dynamodb.py b/awscli/customizations/dynamodb.py
@@ -20,16 +20,7 @@
 
 
 def register_dynamodb_paginator_fix(event_emitter):
-    event_emitter.register(
-        'calling-command.dynamodb.*', _maybe_register_pagination_fix
-    )
-
-
-def _maybe_register_pagination_fix(parsed_globals, session, **kwargs):
-    if parsed_globals.paginate:
-        session.register(
-            'after-call.dynamodb.*', parse_last_evaluated_key_binary
-        )
+    DynamoDBPaginatorFix(event_emitter).register_events()
 
 
 def parse_last_evaluated_key_binary(parsed, **kwargs):
@@ -43,7 +34,20 @@ def parse_last_evaluated_key_binary(parsed, **kwargs):
         return
     for key, val in last_evaluated_key.items():
         if 'B' in val:
-            try:
-                val['B'] = base64.b64decode(val['B'])
-            except (binascii.Error, TypeError):
-                logger.debug('Failed to base64 decode potential binary')
+            val['B'] = base64.b64decode(val['B'])
+
+
+class DynamoDBPaginatorFix(object):
+    def __init__(self, event_emitter):
+        self._event_emitter = event_emitter
+
+    def register_events(self):
+        self._event_emitter.register(
+            'calling-command.dynamodb.*', self._maybe_register_pagination_fix
+        )
+
+    def _maybe_register_pagination_fix(self, parsed_globals, **kwargs):
+        if parsed_globals.paginate:
+            self._event_emitter.register(
+                'after-call.dynamodb.*', parse_last_evaluated_key_binary
+            )
diff --git a/tests/unit/customizations/test_dynamodb.py b/tests/unit/customizations/test_dynamodb.py
@@ -32,14 +32,8 @@ def test_ignores_response_without_last_evaluated_key(self):
     def test_parses_binary_key(self):
         self.assert_parsed_output({'B': 'Zm9v'}, {'B': b'foo'})
 
-    def test_ignores_already_parsed_binary(self):
-        self.assert_parsed_output({'B': b'foo'}, {'B': b'foo'})
-
     def test_ignores_strings(self):
         self.assert_parsed_output({'S': 'Zm9v'}, {'S': 'Zm9v'})
 
     def test_ignores_numbers(self):
         self.assert_parsed_output({'N': 2}, {'N': 2})
-
-    def test_gracefully_handles_invalid_base64(self):
-        self.assert_parsed_output({'B': 'notb64'}, {'B': 'notb64'})