Add tests for structured logging

Author: Artem Krivonos

awslambdaric/bootstrap.py

@@ -13,8 +13,8 @@
 from .lambda_context import LambdaContext
 from .lambda_runtime_client import LambdaRuntimeClient
 from .lambda_runtime_exception import FaultException
+from .lambda_runtime_log_utils import INVOCATION_LOGGING_CONTEXT, JsonFormatter
 from .lambda_runtime_marshaller import to_json
-from .lambda_runtime_log_utils import JsonFormatter, INSTANCE_LOGGING_CONTEXT
 
 ERROR_LOG_LINE_TERMINATE = "\r"
 ERROR_LOG_IDENT = "\u00a0"  # NO-BREAK SPACE U+00A0
@@ -372,6 +372,25 @@ def create_log_sink():
 _GLOBAL_AWS_REQUEST_ID = None
 
 
+def setup_logging(log_format, log_sink):
+    logging.Formatter.converter = time.gmtime
+    logger = logging.getLogger()
+    logger_handler = LambdaLoggerHandler(log_sink)
+    if log_format == "JSON":
+        logger_handler.setFormatter(JsonFormatter())
+        logging._levelToName[logging.CRITICAL] = "FATAL"
+    else:
+        logger_handler.setFormatter(
+            logging.Formatter(
+                "[%(levelname)s]\t%(asctime)s.%(msecs)03dZ\t%(aws_request_id)s\t%(message)s\n",
+                "%Y-%m-%dT%H:%M:%S",
+            )
+        )
+
+    logger_handler.addFilter(LambdaLoggerFilter())
+    logger.addHandler(logger_handler)
+
+
 def run(app_root, handler, lambda_runtime_api_addr):
     sys.stdout = Unbuffered(sys.stdout)
     sys.stderr = Unbuffered(sys.stderr)
@@ -380,23 +399,7 @@ def run(app_root, handler, lambda_runtime_api_addr):
         lambda_runtime_client = LambdaRuntimeClient(lambda_runtime_api_addr)
 
         try:
-            logging.Formatter.converter = time.gmtime
-            logger = logging.getLogger()
-            logger_handler = LambdaLoggerHandler(log_sink)
-            if AWS_LAMBDA_LOG_FORMAT == "JSON":
-                logger_handler.setFormatter(JsonFormatter())
-                logging._levelToName[logging.CRITICAL] = "FATAL"
-            else:
-                logger_handler.setFormatter(
-                    logging.Formatter(
-                        "[%(levelname)s]\t%(asctime)s.%(msecs)03dZ\t%(aws_request_id)s\t%(message)s\n",
-                        "%Y-%m-%dT%H:%M:%S",
-                    )
-                )
-
-            logger_handler.addFilter(LambdaLoggerFilter())
-            logger.addHandler(logger_handler)
-
+            setup_logging(AWS_LAMBDA_LOG_FORMAT, log_sink)
             global _GLOBAL_AWS_REQUEST_ID
 
             request_handler = _get_handler(handler)
@@ -410,7 +413,7 @@ def run(app_root, handler, lambda_runtime_api_addr):
 
         while True:
             event_request = lambda_runtime_client.wait_next_invocation()
-            INSTANCE_LOGGING_CONTEXT.clear()
+            INVOCATION_LOGGING_CONTEXT.clear()
 
             _GLOBAL_AWS_REQUEST_ID = event_request.invoke_id
 

awslambdaric/lambda_runtime_log_utils.py

@@ -6,7 +6,6 @@
 
 from .lambda_runtime_marshaller import to_json
 
-
 RESERVED_FIELDS = {
     "name",
     "msg",
@@ -57,6 +56,9 @@ def format_exception(exc_info):
 
     @staticmethod
     def format_location(record: logging.LogRecord):
+        if not record.exc_info:
+            return None
+
         return f"{record.pathname}:{record.funcName}:{record.lineno}"
 
     def format(self, record: logging.LogRecord) -> str:
@@ -81,7 +83,7 @@ def format(self, record: logging.LogRecord) -> str:
 
         result = {k: v for k, v in result.items() if v is not None}
 
-        return to_json(result)
+        return to_json(result) + "\n"
 
 
 INSTANCE_LOGGING_CONTEXT = dict()

tests/test_bootstrap.py

@@ -4,6 +4,7 @@
 
 import importlib
 import json
+import logging
 import os
 import re
 import tempfile
@@ -17,6 +18,12 @@
 import awslambdaric.bootstrap as bootstrap
 from awslambdaric.lambda_runtime_exception import FaultException
 from awslambdaric.lambda_runtime_marshaller import LambdaMarshaller
+from awslambdaric.lambda_runtime_log_utils import (
+    lambda_instance_logs_update,
+    lambda_invocation_logs_update,
+    INVOCATION_LOGGING_CONTEXT,
+    INSTANCE_LOGGING_CONTEXT,
+)
 
 
 class TestUpdateXrayEnv(unittest.TestCase):
@@ -1174,6 +1181,118 @@ def test_multiple_frame(self):
                 self.assertEqual(content[pos:], b"")
 
 
+class TestLogging(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        bootstrap.setup_logging("JSON", bootstrap.StandardLogSink())
+
+    def tearDown(self) -> None:
+        INVOCATION_LOGGING_CONTEXT.clear()
+        INSTANCE_LOGGING_CONTEXT.clear()
+
+    @patch("sys.stderr", new_callable=StringIO)
+    def test_json_formatter(self, mock_stderr):
+        logger = logging.getLogger("a.b")
+
+        test_cases = [
+            (
+                logging.ERROR,
+                "TEST 1",
+                {
+                    "level": "ERROR",
+                    "logger": "a.b",
+                    "message": "TEST 1",
+                    "requestId": "",
+                },
+            ),
+            (
+                logging.ERROR,
+                "test \nwith \nnew \nlines",
+                {
+                    "level": "ERROR",
+                    "logger": "a.b",
+                    "message": "test \nwith \nnew \nlines",
+                    "requestId": "",
+                },
+            ),
+            (
+                logging.CRITICAL,
+                "TEST FATAL",
+                {
+                    "level": "FATAL",
+                    "logger": "a.b",
+                    "message": "TEST FATAL",
+                    "requestId": "",
+                },
+            ),
+        ]
+        for level, msg, expected in test_cases:
+            with self.subTest(msg):
+                with patch("sys.stdout", new_callable=StringIO) as mock_stdout:
+                    logger.log(level, msg)
+
+                    data = json.loads(mock_stdout.getvalue())
+                    data.pop("timestamp")
+                    self.assertEqual(
+                        data,
+                        expected,
+                    )
+        self.assertEqual(mock_stderr.getvalue(), "")
+
+    @patch("sys.stdout", new_callable=StringIO)
+    @patch("sys.stderr", new_callable=StringIO)
+    def test_exception(self, mock_stderr, mock_stdout):
+        try:
+            raise ValueError("error message")
+        except ValueError:
+            logging.getLogger("test.logger").exception("test exception")
+
+        exception_log = json.loads(mock_stdout.getvalue())
+        self.assertIn("location", exception_log)
+        self.assertIn("stackTrace", exception_log)
+        exception_log.pop("timestamp")
+        exception_log.pop("location")
+        exception_log.pop("stackTrace")
+
+        self.assertEqual(
+            exception_log,
+            {
+                "errorMessage": "error message",
+                "errorType": "ValueError",
+                "level": "ERROR",
+                "logger": "test.logger",
+                "message": "test exception",
+                "requestId": "",
+            },
+        )
+
+        self.assertEqual(mock_stderr.getvalue(), "")
+
+    @patch("sys.stdout", new_callable=StringIO)
+    @patch("sys.stderr", new_callable=StringIO)
+    def test_log_with_extra_params(self, mock_stderr, mock_stdout):
+        lambda_instance_logs_update("instance_key", "instance_value")
+        lambda_invocation_logs_update("invocation_key", "invocation_value")
+
+        logging.getLogger("test.logger").error("test extra params")
+
+        data = json.loads(mock_stdout.getvalue())
+        data.pop("timestamp")
+
+        self.assertEqual(
+            data,
+            {
+                "instance_key": "instance_value",
+                "invocation_key": "invocation_value",
+                "level": "ERROR",
+                "logger": "test.logger",
+                "message": "test extra params",
+                "requestId": "",
+            },
+        )
+        self.assertEqual(mock_stderr.getvalue(), "")
+
+
 class TestBootstrapModule(unittest.TestCase):
     @patch("awslambdaric.bootstrap.handle_event_request")
     @patch("awslambdaric.bootstrap.LambdaRuntimeClient")