Prevent function crash when setting unknown log level.

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,22 @@
+import logging
+import os
+
+_level_mappping = logging.getLevelNamesMapping()
+_level_mappping.update({
+    # https://docs.datadoghq.com/agent/troubleshooting/debug_mode/?tab=agentv6v7#agent-log-level
+    "TRACE": 5,
+    "WARN": logging.WARNING,
+    "OFF": 100,
+})
+
+def initialize_logging(name):
+    logger = logging.getLogger(name)
+    str_level = os.environ.get("DD_LOG_LEVEL", "INFO").upper() or "INFO"
+    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,39 @@
+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)),
+        ("INFO", (20, 30, 40, 50)),
+        ("WARNING", (30, 40, 50)),
+        ("WARN", (30, 40, 50)),
+        ("ERROR", (40, 50)),
+        ("CRITICAL", (50,)),
+        ("OFF", ()),
+        ("", (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):
+    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