changed to explicit check of dsm context

Author: michael-zhao459

datadog_lambda/tracing.py

@@ -78,7 +78,7 @@ def _extract_context_with_data_streams(context_json, event_type, arn):
     if not config.data_streams_enabled:
         return context
 
-    if not _is_context_complete(context):
+    if "dd-pathway-ctx-base64" not in context_json:
         return context
     try:
         carrier_get = _create_carrier_get(context_json)

tests/test_dsm.py

@@ -64,13 +64,13 @@ def test_extract_context_data_streams_enabled_complete_context(self):
             self.assertEqual(kwargs["manual_checkpoint"], False)
             self.assertEqual(result, mock_context)
 
-    def test_extract_context_data_streams_enabled_incomplete_context(self):
+    def test_extract_context_data_streams_enabled_invalid_context(self):
         with patch.dict(os.environ, {"DD_DATA_STREAMS_ENABLED": "true"}):
-            context_json = {"dd-pathway-ctx-base64": "12345"}
+            context_json = {"something-malformed": "12345"}
             event_type = "sqs"
             arn = "arn:aws:sqs:us-east-1:123456789012:test-queue"
 
-            mock_context = Context(trace_id=12345, span_id=None, sampling_priority=1)
+            mock_context = Context(trace_id=12345, span_id=12345, sampling_priority=1)
             self.mock_extract.return_value = mock_context
 
             result = _extract_context_with_data_streams(context_json, event_type, arn)