Skip to content

Commit 07bad6e

Browse files
authored
Do not call flush extension after each invocation. (#406)
* Do not call flush extension after each invocation. * Test flushing per invocation.
1 parent 3db46b2 commit 07bad6e

File tree

2 files changed

+49
-1
lines changed

2 files changed

+49
-1
lines changed

datadog_lambda/wrapper.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@
6060
DD_LOGS_INJECTION = "DD_LOGS_INJECTION"
6161
DD_MERGE_XRAY_TRACES = "DD_MERGE_XRAY_TRACES"
6262
AWS_LAMBDA_FUNCTION_NAME = "AWS_LAMBDA_FUNCTION_NAME"
63+
DD_LOCAL_TEST = "DD_LOCAL_TEST"
6364
DD_TRACE_EXTRACTOR = "DD_TRACE_EXTRACTOR"
6465
DD_TRACE_MANAGED_SERVICES = "DD_TRACE_MANAGED_SERVICES"
6566
DD_ENCODE_AUTHORIZER_CONTEXT = "DD_ENCODE_AUTHORIZER_CONTEXT"
@@ -183,6 +184,9 @@ def __init__(self, func):
183184
self.min_cold_start_trace_duration = get_env_as_int(
184185
DD_MIN_COLD_START_DURATION, 3
185186
)
187+
self.local_testing_mode = os.environ.get(
188+
DD_LOCAL_TEST, "false"
189+
).lower() in ("true", "1")
186190
self.cold_start_trace_skip_lib = [
187191
"ddtrace.internal.compat",
188192
"ddtrace.filters",
@@ -367,7 +371,10 @@ def _after(self, event, context):
367371

368372
if not self.flush_to_log or should_use_extension:
369373
flush_stats()
370-
if should_use_extension:
374+
if should_use_extension and self.local_testing_mode:
375+
# when testing locally, the extension does not know when an
376+
# invocation completes because it does not have access to the
377+
# logs api
371378
flush_extension()
372379

373380
if self.encode_authorizer_context and is_authorizer_response(self.response):

tests/test_wrapper.py

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -647,3 +647,44 @@ def handler(event, context):
647647
},
648648
payload["metadata"],
649649
)
650+
651+
652+
class TestLambdaWrapperFlushExtension(unittest.TestCase):
653+
def setUp(self):
654+
self.orig_environ = os.environ
655+
656+
def tearDown(self):
657+
os.environ = self.orig_environ
658+
659+
@patch("datadog_lambda.wrapper.should_use_extension", True)
660+
def test_local_test_envvar_flushing(self):
661+
662+
flushes = []
663+
lambda_event = {}
664+
lambda_context = get_mock_context()
665+
666+
def flush():
667+
flushes.append(1)
668+
669+
for environ, flush_called in (
670+
({"DD_LOCAL_TEST": "True"}, True),
671+
({"DD_LOCAL_TEST": "true"}, True),
672+
({"DD_LOCAL_TEST": "1"}, True),
673+
({"DD_LOCAL_TEST": "False"}, False),
674+
({"DD_LOCAL_TEST": "false"}, False),
675+
({"DD_LOCAL_TEST": "0"}, False),
676+
({"DD_LOCAL_TEST": ""}, False),
677+
({}, False),
678+
):
679+
680+
os.environ = environ
681+
flushes.clear()
682+
683+
@patch("datadog_lambda.wrapper.flush_extension", flush)
684+
@wrapper.datadog_lambda_wrapper
685+
def lambda_handler(event, context):
686+
pass
687+
688+
lambda_handler(lambda_event, lambda_context)
689+
690+
self.assertEqual(flush_called, len(flushes) == 1)

0 commit comments

Comments
 (0)