renamed to extract_context_with_datastreams

Author: michael-zhao459

datadog_lambda/tracing.py

@@ -66,7 +66,7 @@
 LOWER_64_BITS = "LOWER_64_BITS"
 
 
-def _extract_context(context_json, event_type, arn):
+def _extract_context_with_data_streams(context_json, event_type, arn):
     from ddtrace.data_streams import set_consume_checkpoint
 
     """
@@ -307,7 +307,9 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
                         logger.debug(
                             "Failed to extract Step Functions context from SQS/SNS event."
                         )
-                return _extract_context(dd_data, "sqs" if is_sqs else "sns", arn)
+                return _extract_context_with_data_streams(
+                    dd_data, "sqs" if is_sqs else "sns", arn
+                )
         else:
             # Handle case where trace context is injected into attributes.AWSTraceHeader
             # example: Root=1-654321ab-000000001234567890abcdef;Parent=0123456789abcdef;Sampled=1
@@ -408,7 +410,7 @@ def extract_context_from_kinesis_event(event, lambda_context):
             data_obj = json.loads(data_str)
             dd_ctx = data_obj.get("_datadog")
             if dd_ctx:
-                return _extract_context(dd_ctx, "kinesis", arn)
+                return _extract_context_with_data_streams(dd_ctx, "kinesis", arn)
     except Exception as e:
         logger.debug("The trace extractor returned with error %s", e)
 

tests/test_dsm.py

@@ -7,7 +7,7 @@
 from ddtrace.trace import Context
 
 from datadog_lambda.tracing import (
-    _extract_context,
+    _extract_context_with_data_streams,
     _create_carrier_get,
     extract_context_from_sqs_or_sns_event_or_context,
     extract_context_from_kinesis_event,
@@ -38,7 +38,7 @@ def test_extract_context_data_streams_disabled(self):
             mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1)
             self.mock_extract.return_value = mock_context
 
-            result = _extract_context(context_json, event_type, arn)
+            result = _extract_context_with_data_streams(context_json, event_type, arn)
 
             self.mock_extract.assert_called_once_with(context_json)
             self.mock_checkpoint.assert_not_called()
@@ -53,7 +53,7 @@ def test_extract_context_data_streams_enabled_complete_context(self):
             mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1)
             self.mock_extract.return_value = mock_context
 
-            result = _extract_context(context_json, event_type, arn)
+            result = _extract_context_with_data_streams(context_json, event_type, arn)
 
             self.mock_extract.assert_called_once_with(context_json)
             self.mock_checkpoint.assert_called_once()
@@ -73,7 +73,7 @@ def test_extract_context_data_streams_enabled_incomplete_context(self):
             mock_context = Context(trace_id=12345, span_id=None, sampling_priority=1)
             self.mock_extract.return_value = mock_context
 
-            result = _extract_context(context_json, event_type, arn)
+            result = _extract_context_with_data_streams(context_json, event_type, arn)
 
             self.mock_extract.assert_called_once_with(context_json)
             self.mock_checkpoint.assert_not_called()
@@ -91,7 +91,7 @@ def test_extract_context_exception_path(self):
             test_exception = Exception("Test exception")
             self.mock_checkpoint.side_effect = test_exception
 
-            result = _extract_context(context_json, event_type, arn)
+            result = _extract_context_with_data_streams(context_json, event_type, arn)
 
             self.mock_extract.assert_called_once_with(context_json)
             self.mock_checkpoint.assert_called_once()
@@ -136,8 +136,10 @@ class TestExtractContextFromSqsOrSnsEvent(unittest.TestCase):
     def setUp(self):
         self.lambda_context = get_mock_context()
 
-    @patch("datadog_lambda.tracing._extract_context")
-    def test_sqs_event_with_datadog_message_attributes(self, mock_extract_context):
+    @patch("datadog_lambda.tracing._extract_context_with_data_streams")
+    def test_sqs_event_with_datadog_message_attributes(
+        self, mock_extract_context_with_data_streams
+    ):
         dd_data = {"dd-pathway-ctx-base64": "12345"}
         dd_json_data = json.dumps(dd_data)
 
@@ -153,20 +155,20 @@ def test_sqs_event_with_datadog_message_attributes(self, mock_extract_context):
         }
 
         mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1)
-        mock_extract_context.return_value = mock_context
+        mock_extract_context_with_data_streams.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
             event, self.lambda_context
         )
 
-        mock_extract_context.assert_called_once_with(
+        mock_extract_context_with_data_streams.assert_called_once_with(
             dd_data, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue"
         )
         self.assertEqual(result, mock_context)
 
-    @patch("datadog_lambda.tracing._extract_context")
+    @patch("datadog_lambda.tracing._extract_context_with_data_streams")
     def test_sqs_event_with_binary_datadog_message_attributes(
-        self, mock_extract_context
+        self, mock_extract_context_with_data_streams
     ):
         dd_data = {"dd-pathway-ctx-base64": "12345"}
         dd_json_data = json.dumps(dd_data)
@@ -184,19 +186,21 @@ def test_sqs_event_with_binary_datadog_message_attributes(
         }
 
         mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1)
-        mock_extract_context.return_value = mock_context
+        mock_extract_context_with_data_streams.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
             event, self.lambda_context
         )
 
-        mock_extract_context.assert_called_once_with(
+        mock_extract_context_with_data_streams.assert_called_once_with(
             dd_data, "sqs", "arn:aws:sqs:us-east-1:123456789012:test-queue"
         )
         self.assertEqual(result, mock_context)
 
-    @patch("datadog_lambda.tracing._extract_context")
-    def test_sns_event_with_datadog_message_attributes(self, mock_extract_context):
+    @patch("datadog_lambda.tracing._extract_context_with_data_streams")
+    def test_sns_event_with_datadog_message_attributes(
+        self, mock_extract_context_with_data_streams
+    ):
         dd_data = {"dd-pathway-ctx-base64": "12345"}
         dd_json_data = json.dumps(dd_data)
 
@@ -215,13 +219,13 @@ def test_sns_event_with_datadog_message_attributes(self, mock_extract_context):
         }
 
         mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1)
-        mock_extract_context.return_value = mock_context
+        mock_extract_context_with_data_streams.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
             event, self.lambda_context
         )
 
-        mock_extract_context.assert_called_once_with(
+        mock_extract_context_with_data_streams.assert_called_once_with(
             dd_data, "sns", "arn:aws:sns:us-east-1:123456789012:test-topic"
         )
         self.assertEqual(result, mock_context)
@@ -231,8 +235,10 @@ class TestExtractContextFromKinesisEvent(unittest.TestCase):
     def setUp(self):
         self.lambda_context = get_mock_context()
 
-    @patch("datadog_lambda.tracing._extract_context")
-    def test_kinesis_event_with_datadog_data(self, mock_extract_context):
+    @patch("datadog_lambda.tracing._extract_context_with_data_streams")
+    def test_kinesis_event_with_datadog_data(
+        self, mock_extract_context_with_data_streams
+    ):
         dd_data = {"dd-pathway-ctx-base64": "12345"}
         kinesis_data = {"_datadog": dd_data, "message": "test"}
         kinesis_data_str = json.dumps(kinesis_data)
@@ -248,11 +254,11 @@ def test_kinesis_event_with_datadog_data(self, mock_extract_context):
         }
 
         mock_context = Context(trace_id=12345, span_id=67890, sampling_priority=1)
-        mock_extract_context.return_value = mock_context
+        mock_extract_context_with_data_streams.return_value = mock_context
 
         result = extract_context_from_kinesis_event(event, self.lambda_context)
 
-        mock_extract_context.assert_called_once_with(
+        mock_extract_context_with_data_streams.assert_called_once_with(
             dd_data,
             "kinesis",
             "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream",