Add local_test.

Author: purple4reina

datadog_lambda/config.py

@@ -68,6 +68,8 @@ class Config:
     llmobs_enabled = _get_env("DD_LLMOBS_ENABLED", "false", as_bool)
     exception_replay_enabled = _get_env("DD_EXCEPTION_REPLAY_ENABLED", "false", as_bool)
 
+    local_test = _get_env("DD_LOCAL_TEST", "false", as_bool)
+
     @property
     def fips_mode_enabled(self):
         if not hasattr(self, "_config_fips_mode_enabled"):

datadog_lambda/wrapper.py

@@ -58,7 +58,6 @@
 
 logger = logging.getLogger(__name__)
 
-DD_LOCAL_TEST = "DD_LOCAL_TEST"
 DD_TRACE_MANAGED_SERVICES = "DD_TRACE_MANAGED_SERVICES"
 DD_ENCODE_AUTHORIZER_CONTEXT = "DD_ENCODE_AUTHORIZER_CONTEXT"
 DD_DECODE_AUTHORIZER_CONTEXT = "DD_DECODE_AUTHORIZER_CONTEXT"
@@ -70,16 +69,6 @@
 DD_DATA_STREAMS_ENABLED = "DD_DATA_STREAMS_ENABLED"
 
 
-def get_env_as_int(env_key, default_value: int) -> int:
-    try:
-        return int(os.environ.get(env_key, default_value))
-    except Exception as e:
-        logger.warn(
-            f"Failed to parse {env_key} as int. Using default value: {default_value}. Error: {e}"
-        )
-        return default_value
-
-
 init_timestamp_ns = time_ns()
 
 """
@@ -153,9 +142,6 @@ def __init__(self, func):
             self.cold_start_tracing = depends_on_dd_tracing_enabled(
                 os.environ.get(DD_COLD_START_TRACING, "true").lower() == "true"
             )
-            self.min_cold_start_trace_duration = get_env_as_int(
-                DD_MIN_COLD_START_DURATION, 3
-            )
             self.data_streams_enabled = (
                 os.environ.get(DD_DATA_STREAMS_ENABLED, "false").lower() == "true"
             )
@@ -368,7 +354,7 @@ def _after(self, event, context):
                 from datadog_lambda.metric import flush_stats
 
                 flush_stats(context)
-            if should_use_extension and self.local_testing_mode:
+            if should_use_extension and config.local_test:
                 # when testing locally, the extension does not know when an
                 # invocation completes because it does not have access to the
                 # logs api

tests/test_config.py

@@ -183,6 +183,15 @@ def set_env(key, value):
     ("DD_MIN_COLD_START_DURATION", "min_cold_start_trace_duration", "2.5", 3),
     ("DD_MIN_COLD_START_DURATION", "min_cold_start_trace_duration", "-1", -1),
     ("DD_MIN_COLD_START_DURATION", "min_cold_start_trace_duration", "purple", 3),
+    ("DD_LOCAL_TEST", "local_test", None, False),
+    ("DD_LOCAL_TEST", "local_test", "", False),
+    ("DD_LOCAL_TEST", "local_test", "true", True),
+    ("DD_LOCAL_TEST", "local_test", "TRUE", True),
+    ("DD_LOCAL_TEST", "local_test", "false", False),
+    ("DD_LOCAL_TEST", "local_test", "FALSE", False),
+    ("DD_LOCAL_TEST", "local_test", "1", True),  # CHANGED
+    ("DD_LOCAL_TEST", "local_test", "0", False),
+    ("DD_LOCAL_TEST", "local_test", "purple", False),
 )