feat: allow custom kwargs #895

heitorlessa · Dec 20, 2021 · 7326ba6 · 7326ba6
1 parent 39b3136
commit 7326ba6
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/aws_lambda_powertools/logging/logger.py b/aws_lambda_powertools/logging/logger.py
@@ -328,7 +328,7 @@ def handler(event, context):
         )
 
         @functools.wraps(lambda_handler)
-        def decorate(event, context):
+        def decorate(event, context, **kwargs):
             lambda_context = build_lambda_context_model(context)
             cold_start = _is_cold_start()
 

diff --git a/tests/functional/test_logger.py b/tests/functional/test_logger.py
@@ -597,3 +597,16 @@ def handler(event, context):
     first_log, second_log = capture_multiple_logging_statements_output(stdout)
     assert "my_key" in first_log
     assert "my_key" not in second_log
+
+
+def test_inject_lambda_context_allows_handler_with_kwargs(lambda_context, stdout, service_name):
+    # GIVEN
+    logger = Logger(service=service_name, stream=stdout)
+
+    # WHEN
+    @logger.inject_lambda_context(clear_state=True)
+    def handler(event, context, my_custom_option=None):
+        pass
+
+    # THEN
+    handler({}, lambda_context, my_custom_option="blah")