Skip to content

Commit b1797fe

Browse files
committed
Test flushing per invocation.
1 parent 2713470 commit b1797fe

File tree

2 files changed

+44
-3
lines changed

2 files changed

+44
-3
lines changed

datadog_lambda/wrapper.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -184,9 +184,9 @@ def __init__(self, func):
184184
self.min_cold_start_trace_duration = get_env_as_int(
185185
DD_MIN_COLD_START_DURATION, 3
186186
)
187-
self.local_testing_mode = (
188-
os.environ.get(DD_LOCAL_TEST, "false").lower() in ("true", "1")
189-
)
187+
self.local_testing_mode = os.environ.get(
188+
DD_LOCAL_TEST, "false"
189+
).lower() in ("true", "1")
190190
self.cold_start_trace_skip_lib = [
191191
"ddtrace.internal.compat",
192192
"ddtrace.filters",

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+
class TestLambdaWrapperFlushExtension(unittest.TestCase):
652+
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)