Merge pull request #80 from enter-at/fix/methods-decorator

fix(lambda_handler): ensure decorator forwards handler self

Author: alexsavio

lambda_handlers/handlers/lambda_handler.py

@@ -19,16 +19,20 @@ def __init__(self, handler: Optional[Callable] = None):
 
     def __call__(self, handler: Callable):  # noqa: D102
         @wraps(handler)
-        def wrapper(event, context):
-            return self._call_handler(handler, event, context)
+        def wrapper(*args):
+            try:
+                handler_self = args[0] if len(args) > 2 else None
+                return self.after(self._call_handler(handler_self, handler, *self.before(*args[-2:])))
+            except Exception as exception:
+                return self.on_exception(exception)
 
         return wrapper
 
-    def _call_handler(self, handler: Callable, event: Event, context: LambdaContext) -> Any:
-        try:
-            return self.after(handler(*self.before(event, context)))
-        except Exception as exception:
-            return self.on_exception(exception)
+    def _call_handler(self, handler_self, handler, event, context):
+        if handler_self:
+            return handler(handler_self, event, context)
+
+        return handler(event, context)
 
     def before(self, event: Event, context: LambdaContext) -> Tuple[Event, LambdaContext]:
         """Event method to be called just before the handler is executed."""

lambda_handlers/handlers/tests/test_lambda_handler.py

@@ -31,7 +31,7 @@ def on_exception(self, exception: Exception) -> Any:
 
 
 @pytest.fixture
-def handler():
+def function_handler():
     @lambda_handler
     def handler(event, context):
         if context is None:
@@ -41,20 +41,51 @@ def handler(event, context):
     return handler
 
 
-class TestLambdaHandler:
+@pytest.fixture
+def method_handler():
+    class Adapter:
+        @lambda_handler
+        def __call__(self, event, context):
+            if context is None:
+                raise EventAwareException(message='no such context', event=event)
+            return event
+
+    return Adapter()
+
+
+class TestLambdaHandlerDecorateFunction:
+
+    @pytest.fixture
+    def event(self):
+        return {'route': []}
+
+    def test_call_order(self, function_handler, event):
+        result = function_handler(event, {})
+
+        assert result == event
+        assert event['route'] == ['before', 'after']
+
+    def test_call_exception(self, function_handler, event):
+        with pytest.raises(EventAwareException, match='no such context'):
+            function_handler(event, None)
+
+        assert event['route'] == ['before', 'on_exception']
+
+
+class TestLambdaHandlerDecorateMethod:
 
     @pytest.fixture
     def event(self):
         return {'route': []}
 
-    def test_call_order(self, handler, event):
-        result = handler(event, {})
+    def test_call_order(self, method_handler, event):
+        result = method_handler(event, {})
 
         assert result == event
         assert event['route'] == ['before', 'after']
 
-    def test_call_exception(self, handler, event):
+    def test_call_exception(self, method_handler, event):
         with pytest.raises(EventAwareException, match='no such context'):
-            handler(event, None)
+            method_handler(event, None)
 
         assert event['route'] == ['before', 'on_exception']