From 6e5c0851e8413f1be5e92ad022097027e7185b65 Mon Sep 17 00:00:00 2001 From: Stephen Firincieli Date: Fri, 28 Oct 2022 15:23:31 -0400 Subject: [PATCH 1/2] Don't error when headers are None --- datadog_lambda/tracing.py | 2 +- tests/test_tracing.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 4db7db98..1e63a78b 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -174,7 +174,7 @@ def extract_context_from_http_event_or_context(event, lambda_context): Falls back to lambda context if no trace data is found in the `headers` """ - headers = event.get("headers", {}) + headers = event.get("headers", {}) or {} lowercase_headers = {k.lower(): v for k, v in headers.items()} trace_id = lowercase_headers.get(TraceHeader.TRACE_ID) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index c578de07..1e623347 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -182,6 +182,7 @@ def test_with_complete_datadog_trace_headers(self): {"trace-id": "123", "parent-id": "321", "sampling-priority": "1"}, ) + def test_with_extractor_function(self): def extractor_foo(event, context): foo = event.get("foo", {}) @@ -1237,3 +1238,14 @@ def test_mark_trace_as_error_for_5xx_responses_sends_error_metric_and_set_error_ ) mock_submit_errors_metric.assert_called_once() self.assertEqual(1, mock_span.error) + + + def test_no_error_with_nonetype_headers(self): + lambda_ctx = get_mock_context() + ctx, source = extract_dd_trace_context( + { + "headers": None + }, + lambda_ctx, + ) + self.assertEqual(ctx, None) \ No newline at end of file From 586f965898cb3026b5e5dc3f1c7b008dc8657a85 Mon Sep 17 00:00:00 2001 From: Stephen Firincieli Date: Fri, 28 Oct 2022 15:45:35 -0400 Subject: [PATCH 2/2] Fix formatting of file --- tests/test_tracing.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 1e623347..7f5ba4d3 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -182,7 +182,6 @@ def test_with_complete_datadog_trace_headers(self): {"trace-id": "123", "parent-id": "321", "sampling-priority": "1"}, ) - def test_with_extractor_function(self): def extractor_foo(event, context): foo = event.get("foo", {}) @@ -1239,13 +1238,10 @@ def test_mark_trace_as_error_for_5xx_responses_sends_error_metric_and_set_error_ mock_submit_errors_metric.assert_called_once() self.assertEqual(1, mock_span.error) - def test_no_error_with_nonetype_headers(self): lambda_ctx = get_mock_context() ctx, source = extract_dd_trace_context( - { - "headers": None - }, + {"headers": None}, lambda_ctx, ) - self.assertEqual(ctx, None) \ No newline at end of file + self.assertEqual(ctx, None)