From 20e281ec12dabe9c2d921f61b5ab1d3a2cbc9679 Mon Sep 17 00:00:00 2001 From: Rey Abolofia Date: Thu, 21 Nov 2024 09:04:31 -0800 Subject: [PATCH] Prevent error when attempting to use None context. --- datadog_lambda/tracing.py | 2 ++ tests/test_tracing.py | 7 +++++++ 2 files changed, 9 insertions(+) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 5ed306a3..347e2fb7 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -699,6 +699,8 @@ def inject_correlation_ids(): def set_dd_trace_py_root(trace_context_source, merge_xray_traces): + if not _is_context_complete(dd_trace_context): + return if trace_context_source == TraceContextSource.EVENT or merge_xray_traces: context = Context( trace_id=dd_trace_context.trace_id, diff --git a/tests/test_tracing.py b/tests/test_tracing.py index d0db05cd..4c530ad0 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -934,6 +934,9 @@ def setUp(self): self.mock_activate = patcher.start() self.mock_activate.return_value = True self.addCleanup(patcher.stop) + patcher = patch("datadog_lambda.tracing.dd_trace_context", None) + self.mock_dd_trace_context = patcher.start() + self.addCleanup(patcher.stop) def tearDown(self): global dd_tracing_enabled @@ -991,6 +994,10 @@ def test_set_dd_trace_py_root_no_span_id(self): self.mock_activate.assert_called() self.mock_activate.assert_has_calls([call(expected_context)]) + def test_set_dd_trace_py_root_none_context(self): + set_dd_trace_py_root(TraceContextSource.EVENT, True) + self.mock_activate.assert_not_called() + class TestServiceMapping(unittest.TestCase): def setUp(self):