diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 15dec7f1..73d17887 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -60,6 +60,7 @@ DD_LOGS_INJECTION = "DD_LOGS_INJECTION" DD_MERGE_XRAY_TRACES = "DD_MERGE_XRAY_TRACES" AWS_LAMBDA_FUNCTION_NAME = "AWS_LAMBDA_FUNCTION_NAME" +DD_LOCAL_TEST = "DD_LOCAL_TEST" DD_TRACE_EXTRACTOR = "DD_TRACE_EXTRACTOR" DD_TRACE_MANAGED_SERVICES = "DD_TRACE_MANAGED_SERVICES" DD_ENCODE_AUTHORIZER_CONTEXT = "DD_ENCODE_AUTHORIZER_CONTEXT" @@ -183,6 +184,9 @@ def __init__(self, func): self.min_cold_start_trace_duration = get_env_as_int( DD_MIN_COLD_START_DURATION, 3 ) + self.local_testing_mode = os.environ.get( + DD_LOCAL_TEST, "false" + ).lower() in ("true", "1") self.cold_start_trace_skip_lib = [ "ddtrace.internal.compat", "ddtrace.filters", @@ -367,7 +371,10 @@ def _after(self, event, context): if not self.flush_to_log or should_use_extension: flush_stats() - if should_use_extension: + if should_use_extension and self.local_testing_mode: + # when testing locally, the extension does not know when an + # invocation completes because it does not have access to the + # logs api flush_extension() if self.encode_authorizer_context and is_authorizer_response(self.response): diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index 29c842d2..39998921 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -647,3 +647,44 @@ def handler(event, context): }, payload["metadata"], ) + + +class TestLambdaWrapperFlushExtension(unittest.TestCase): + def setUp(self): + self.orig_environ = os.environ + + def tearDown(self): + os.environ = self.orig_environ + + @patch("datadog_lambda.wrapper.should_use_extension", True) + def test_local_test_envvar_flushing(self): + + flushes = [] + lambda_event = {} + lambda_context = get_mock_context() + + def flush(): + flushes.append(1) + + for environ, flush_called in ( + ({"DD_LOCAL_TEST": "True"}, True), + ({"DD_LOCAL_TEST": "true"}, True), + ({"DD_LOCAL_TEST": "1"}, True), + ({"DD_LOCAL_TEST": "False"}, False), + ({"DD_LOCAL_TEST": "false"}, False), + ({"DD_LOCAL_TEST": "0"}, False), + ({"DD_LOCAL_TEST": ""}, False), + ({}, False), + ): + + os.environ = environ + flushes.clear() + + @patch("datadog_lambda.wrapper.flush_extension", flush) + @wrapper.datadog_lambda_wrapper + def lambda_handler(event, context): + pass + + lambda_handler(lambda_event, lambda_context) + + self.assertEqual(flush_called, len(flushes) == 1)