Do not call flush extension after each invocation.

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):

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)