fix(logger): clear_state should keep custom key formats (#1095)

Author: eliasbrange

aws_lambda_powertools/logging/formatter.py

@@ -127,8 +127,8 @@ def __init__(
 
         super(LambdaPowertoolsFormatter, self).__init__(datefmt=self.datefmt)
 
-        keys_combined = {**self._build_default_keys(), **kwargs}
-        self.log_format.update(**keys_combined)
+        self.keys_combined = {**self._build_default_keys(), **kwargs}
+        self.log_format.update(**self.keys_combined)
 
     def serialize(self, log: Dict) -> str:
         """Serialize structured log dict to JSON str"""
@@ -187,7 +187,7 @@ def remove_keys(self, keys: Iterable[str]):
 
     def clear_state(self):
         self.log_format = dict.fromkeys(self.log_record_order)
-        self.log_format.update(**self._build_default_keys())
+        self.log_format.update(**self.keys_combined)
 
     @staticmethod
     def _build_default_keys():

tests/functional/test_logger.py

@@ -683,6 +683,26 @@ def handler(event, context):
         assert key in second_log
 
 
+def test_clear_state_keeps_custom_keys(lambda_context, stdout, service_name):
+    # GIVEN
+    location_format = "%(module)s.%(funcName)s:clear_state"
+    logger = Logger(service=service_name, stream=stdout, location=location_format, custom_key="foo")
+
+    # WHEN clear_state is set
+    @logger.inject_lambda_context(clear_state=True)
+    def handler(event, context):
+        logger.info("Foo")
+
+    # THEN all standard keys should be available as usual
+    handler({}, lambda_context)
+    handler({}, lambda_context)
+
+    first_log, second_log = capture_multiple_logging_statements_output(stdout)
+    for log in (first_log, second_log):
+        assert "foo" == log["custom_key"]
+        assert "test_logger.handler:clear_state" == log["location"]
+
+
 def test_clear_state_keeps_exception_keys(lambda_context, stdout, service_name):
     # GIVEN
     logger = Logger(service=service_name, stream=stdout)