Prevent function crash when setting unknown log level. (#418)

Author: purple4reina

datadog_lambda/__init__.py

@@ -1,6 +1,5 @@
-import os
-import logging
 from datadog_lambda.cold_start import initialize_cold_start_tracing
+from datadog_lambda.logger import initialize_logging
 
 initialize_cold_start_tracing()
 
@@ -13,5 +12,4 @@
 
 __version__ = importlib_metadata.version(__name__)
 
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.getLevelName(os.environ.get("DD_LOG_LEVEL", "INFO").upper()))
+initialize_logging(__name__)

datadog_lambda/logger.py

@@ -0,0 +1,27 @@
+import logging
+import os
+
+try:
+    _level_mappping = logging.getLevelNamesMapping()
+except AttributeError:
+    # python 3.8
+    _level_mappping = {name: num for num, name in logging._levelToName.items()}
+# https://docs.datadoghq.com/agent/troubleshooting/debug_mode/?tab=agentv6v7#agent-log-level
+_level_mappping.update(
+    {
+        "TRACE": 5,
+        "WARN": logging.WARNING,
+        "OFF": 100,
+    }
+)
+
+
+def initialize_logging(name):
+    logger = logging.getLogger(name)
+    str_level = (os.environ.get("DD_LOG_LEVEL") or "INFO").upper()
+    level = _level_mappping.get(str_level)
+    if level is None:
+        logger.setLevel(logging.INFO)
+        logger.warning("Invalid log level: %s Defaulting to INFO", str_level)
+    else:
+        logger.setLevel(level)

tests/test_logger.py

@@ -0,0 +1,43 @@
+import io
+import logging
+import pytest
+
+from datadog_lambda.logger import initialize_logging
+
+_test_initialize_logging = (
+    ("TRACE", (10, 20, 30, 40, 50)),
+    ("DEBUG", (10, 20, 30, 40, 50)),
+    ("debug", (10, 20, 30, 40, 50)),
+    ("INFO", (20, 30, 40, 50)),
+    ("WARNING", (30, 40, 50)),
+    ("WARN", (30, 40, 50)),
+    ("ERROR", (40, 50)),
+    ("CRITICAL", (50,)),
+    ("OFF", ()),
+    ("", (20, 30, 40, 50)),
+    (None, (20, 30, 40, 50)),
+    ("PURPLE", (30, 20, 30, 40, 50)),  # log warning then default to INFO
+)
+
+
+@pytest.mark.parametrize("level,logged_levels", _test_initialize_logging)
+def test_initialize_logging(level, logged_levels, monkeypatch):
+    if level is not None:
+        monkeypatch.setenv("DD_LOG_LEVEL", level)
+
+    stream = io.StringIO()
+    handler = logging.StreamHandler(stream)
+    handler.setFormatter(logging.Formatter("%(levelno)s"))
+    logger = logging.getLogger(__name__)
+    logger.addHandler(handler)
+
+    initialize_logging(__name__)
+
+    logger.debug("debug")
+    logger.info("info")
+    logger.warning("warning")
+    logger.error("error")
+    logger.critical("critical")
+
+    logged = tuple(map(int, stream.getvalue().strip().split()))
+    assert logged == logged_levels