diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index b5324b1e..3cadc538 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -332,7 +332,7 @@ def extract_dd_trace_context(event, lambda_context, extractor=None): parent_id, sampling_priority, ) = extract_context_custom_extractor(extractor, event, lambda_context) - elif "headers" in event: + elif isinstance(event, (list, dict)) and "headers" in event: ( trace_id, parent_id, diff --git a/tests/test_tracing.py b/tests/test_tracing.py index be52697e..e44e9b0c 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -101,6 +101,28 @@ def test_without_datadog_trace_headers(self): {}, ) + def test_with_non_object_event(self): + lambda_ctx = get_mock_context() + ctx, source = extract_dd_trace_context(b"", lambda_ctx) + self.assertEqual(source, "xray") + self.assertDictEqual( + ctx, + { + "trace-id": fake_xray_header_value_root_decimal, + "parent-id": fake_xray_header_value_parent_decimal, + "sampling-priority": "2", + }, + ) + self.assertDictEqual( + get_dd_trace_context(), + { + TraceHeader.TRACE_ID: fake_xray_header_value_root_decimal, + TraceHeader.PARENT_ID: fake_xray_header_value_parent_decimal, + TraceHeader.SAMPLING_PRIORITY: "2", + }, + {}, + ) + def test_with_incomplete_datadog_trace_headers(self): lambda_ctx = get_mock_context() ctx, source = extract_dd_trace_context(