From 27134701c7c32788e2ad3748c29ff1f461b0867c Mon Sep 17 00:00:00 2001
From: Rey Abolofia <rey.abolofia@datadoghq.com>
Date: Thu, 14 Dec 2023 13:43:28 -0800
Subject: [PATCH 1/2] Do not call flush extension after each invocation.

---
 datadog_lambda/wrapper.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py
index 15dec7f1..19ed3ebb 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):

From 516ded2258b848b24cf940e071b175e3f485fd22 Mon Sep 17 00:00:00 2001
From: Rey Abolofia <rey.abolofia@datadoghq.com>
Date: Fri, 15 Dec 2023 10:31:03 -0800
Subject: [PATCH 2/2] Test flushing per invocation.

---
 datadog_lambda/wrapper.py |  6 +++---
 tests/test_wrapper.py     | 41 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py
index 19ed3ebb..73d17887 100644
--- a/datadog_lambda/wrapper.py
+++ b/datadog_lambda/wrapper.py
@@ -184,9 +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.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",
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)