fix(logger): support additional args for handlers when injecting lambda context (#1276)

Author: RonRademaker

aws_lambda_powertools/logging/logger.py

@@ -335,7 +335,7 @@ def handler(event, context):
         )
 
         @functools.wraps(lambda_handler)
-        def decorate(event, context, **kwargs):
+        def decorate(event, context, *args, **kwargs):
             lambda_context = build_lambda_context_model(context)
             cold_start = _is_cold_start()
 
@@ -351,7 +351,7 @@ def decorate(event, context, **kwargs):
                 logger.debug("Event received")
                 self.info(getattr(event, "raw_event", event))
 
-            return lambda_handler(event, context)
+            return lambda_handler(event, context, *args, **kwargs)
 
         return decorate
 

tests/functional/test_logger.py

@@ -773,3 +773,20 @@ def handler(event, context):
     # THEN logger should log event received from Lambda
     logged_event, _ = capture_multiple_logging_statements_output(stdout)
     assert logged_event["message"] == lambda_event
+
+
+def test_inject_lambda_context_with_additional_args(lambda_context, stdout, service_name):
+    # GIVEN Logger is initialized
+    logger = Logger(service=service_name, stream=stdout)
+
+    # AND a handler that use additional parameters
+    @logger.inject_lambda_context
+    def handler(event, context, planet, str_end="."):
+        logger.info(f"Hello {planet}{str_end}")
+
+    handler({}, lambda_context, "World", str_end="!")
+
+    # THEN the decorator should included them
+    log = capture_logging_output(stdout)
+
+    assert log["message"] == "Hello World!"