add tests an d comments

Author: michael-zhao459

datadog_lambda/tracing.py

@@ -220,6 +220,7 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
     try:
         context = _extract_context_from_eventbridge_sqs_event(event)
         if _is_context_complete(context):
+            # Unsupported DSM context propagation method
             return context, None
     except Exception:
         logger.debug("Failed extracting context as EventBridge to SQS.")
@@ -299,6 +300,7 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
                                 span_id=int(x_ray_context["parent_id"], 16),
                                 sampling_priority=float(x_ray_context["sampled"]),
                             ),
+                            # Unsupported DSM context propagation method
                             None,
                         )
         # return a empty dict so dsm checkpoint can be set even if no context is found

datadog_lambda/wrapper.py

@@ -392,10 +392,9 @@ def extract_source_arn(event):
 
 
 def _create_dsm_carrier_func(dd_data):
-    # {}.get(key) returns None, allowing DSM checkpoint to be set even if context is missing
-    # None.get(key) errors, preventing checkpoint for unsupported dsm context propagation methods
-
     def carrier_get(key):
+        # {}.get(key) returns None, DSM checkpoint can be set even if DSM context extraction fails
+        # None.get(key) errors, preventing checkpoint for unsupported dsm context propagation method
         return dd_data.get(key)
 
     return carrier_get

tests/test_wrapper.py

@@ -7,13 +7,11 @@
 from datadog_lambda.constants import TraceHeader
 
 import datadog_lambda.wrapper as wrapper
-import datadog_lambda.xray as xray
 
-from datadog_lambda.config import config
 from datadog_lambda.metric import lambda_metric
 from datadog_lambda.thread_stats_writer import ThreadStatsWriter
 from datadog_lambda.trigger import EventTypes
-from ddtrace.trace import Span, tracer
+from ddtrace.trace import tracer
 from ddtrace.internal.constants import MAX_UINT_64BITS
 
 from tests.utils import get_mock_context, reset_xray_connection
@@ -285,7 +283,9 @@ def lambda_handler(event, context):
     def test_5xx_sends_errors_metric_and_set_tags(self, mock_extract_trigger_tags):
         mock_extract_trigger_tags.return_value = {
             "function_trigger.event_source": "api-gateway",
-            "function_trigger.event_source_arn": "arn:aws:apigateway:us-west-1::/restapis/1234567890/stages/prod",
+            "function_trigger.event_source_arn": (
+                "arn:aws:apigateway:us-west-1::/restapis/1234567890/stages/prod"
+            ),
             "http.url": "https://70ixmpl4fl.execute-api.us-east-2.amazonaws.com",
             "http.url_details.path": "/prod/path/to/resource",
             "http.method": "GET",
@@ -1015,3 +1015,28 @@ def test_extract_source_arn_sns_to_sqs_event(self):
         self.assertEqual(
             result, "arn:aws:sqs:us-east-1:123456789012:test-queue-from-sns"
         )
+
+
+class TestCreateDsmCarrierFunc(unittest.TestCase):
+    def test_carrier_with_data(self):
+        """Test that the carrier function works with a dict-like object."""
+        mock_dd_data = Mock()
+        mock_dd_data.get.return_value = "mock_value"
+
+        carrier = wrapper._create_dsm_carrier_func(mock_dd_data)
+        self.assertTrue(callable(carrier))
+        self.assertEqual(carrier("some_key"), "mock_value")
+        mock_dd_data.get.assert_called_once_with("some_key")
+
+    def test_carrier_with_empty_dict(self):
+        """Test that the carrier function works with an empty dictionary."""
+        carrier = wrapper._create_dsm_carrier_func({})
+        self.assertTrue(callable(carrier))
+        self.assertIsNone(carrier("some_key"))
+
+    def test_carrier_with_none(self):
+        """Test that the carrier function fails correctly when given None."""
+        carrier = wrapper._create_dsm_carrier_func(None)
+        self.assertTrue(callable(carrier))
+        with self.assertRaises(AttributeError):
+            carrier("some_key")