test fixes

michael-zhao459 · michael-zhao459 · commit b19baaee4053 · 2025-06-03T15:31:18.000-04:00
diff --git a/tests/test_dsm.py b/tests/test_dsm.py
@@ -0,0 +1,113 @@
+import unittest
+from unittest.mock import patch, MagicMock
+
+from datadog_lambda.dsm import set_dsm_context, _dsm_set_sqs_context
+from datadog_lambda.trigger import EventTypes
+
+
+class TestDsmContext(unittest.TestCase):
+    def test_non_sqs_event_source_does_nothing(self):
+        """Test that non-SQS event sources don't trigger DSM context setting"""
+        event = {"Records": [{"body": "test"}]}
+
+        mock_event_source = MagicMock()
+        mock_event_source.equals.return_value = False  # Not SQS
+
+        with patch("datadog_lambda.dsm._dsm_set_sqs_context") as mock_sqs_context:
+            set_dsm_context(event, mock_event_source)
+
+            mock_event_source.equals.assert_called_once_with(EventTypes.SQS)
+            mock_sqs_context.assert_not_called()
+
+    def test_event_with_no_records_does_nothing(self):
+        """Test that events where Records is None don't trigger DSM processing"""
+        events_with_no_records = [
+            {},
+            {"Records": None},
+            {"someOtherField": "value"},
+        ]
+
+        for event in events_with_no_records:
+            with patch(
+                "ddtrace.internal.datastreams.data_streams_processor"
+            ) as mock_processor:
+                _dsm_set_sqs_context(event)
+
+                mock_processor.assert_not_called()
+
+    def test_sqs_event_triggers_dsm_sqs_context(self):
+        """Test that SQS event sources trigger the SQS-specific DSM context function"""
+        sqs_event = {
+            "Records": [
+                {
+                    "eventSource": "aws:sqs",
+                    "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:my-queue",
+                    "body": "Hello from SQS!",
+                }
+            ]
+        }
+
+        mock_event_source = MagicMock()
+        mock_event_source.equals.return_value = True
+
+        with patch("datadog_lambda.dsm._dsm_set_sqs_context") as mock_sqs_context:
+            set_dsm_context(sqs_event, mock_event_source)
+
+            mock_sqs_context.assert_called_once_with(sqs_event)
+
+    def test_multiple_records_process_each_record(self):
+        """Test that each record in an SQS event gets processed individually"""
+        multi_record_event = {
+            "Records": [
+                {
+                    "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:queue1",
+                    "body": "Message 1",
+                },
+                {
+                    "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:queue2",
+                    "body": "Message 2",
+                },
+                {
+                    "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:queue3",
+                    "body": "Message 3",
+                },
+            ]
+        }
+
+        mock_processor = MagicMock()
+        mock_context = MagicMock()
+
+        with patch(
+            "ddtrace.internal.datastreams.data_streams_processor",
+            return_value=mock_processor,
+        ):
+            with patch(
+                "ddtrace.internal.datastreams.botocore.get_datastreams_context",
+                return_value={},
+            ):
+                with patch(
+                    "ddtrace.internal.datastreams.botocore.calculate_sqs_payload_size",
+                    return_value=100,
+                ):
+                    with patch(
+                        "ddtrace.internal.datastreams.processor.DsmPathwayCodec.decode",
+                        return_value=mock_context,
+                    ):
+                        _dsm_set_sqs_context(multi_record_event)
+
+                        assert mock_context.set_checkpoint.call_count == 3
+
+                        calls = mock_context.set_checkpoint.call_args_list
+                        expected_arns = [
+                            "arn:aws:sqs:us-east-1:123456789012:queue1",
+                            "arn:aws:sqs:us-east-1:123456789012:queue2",
+                            "arn:aws:sqs:us-east-1:123456789012:queue3",
+                        ]
+
+                        for i, call in enumerate(calls):
+                            args, kwargs = call
+                            tags = args[0]
+                            self.assertIn("direction:in", tags)
+                            self.assertIn(f"topic:{expected_arns[i]}", tags)
+                            self.assertIn("type:sqs", tags)
+                            self.assertEqual(kwargs["payload_size"], 100)
diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py
@@ -563,200 +563,69 @@ def return_type_test(event, context):
             self.assertEqual(result, test_result)
             self.assertFalse(MockPrintExc.called)
 
-    @patch.dict(os.environ, {"DD_DATA_STREAMS_ENABLED": "true"})
-    def test_datadog_lambda_wrapper_dsm_sqs_context_pathway_verification(self):
-        from datadog_lambda.trigger import _EventSource, EventTypes
-
-        sqs_event_source = _EventSource(EventTypes.SQS)
-        self.mock_extract_dd_trace_context.return_value = ({}, None, sqs_event_source)
-
-        with patch(
-            "ddtrace.internal.datastreams.processor.get_connection"
-        ) as mock_get_connection:
-
-            mock_conn = unittest.mock.MagicMock()
-            mock_response = unittest.mock.MagicMock()
-            mock_response.status = 200
-            mock_conn.getresponse.return_value = mock_response
-            mock_get_connection.return_value = mock_conn
-
-            def updated_get_datastreams_context(message):
-                """
-                Updated version that handles the correct message formats
-                """
-                import base64
-                import json
-
-                context_json = None
-                message_body = message
-                try:
-                    body = message.get("Body")
-                    if body:
-                        message_body = json.loads(body)
-                except (ValueError, TypeError):
-                    pass
-
-                message_attributes = message_body.get(
-                    "MessageAttributes"
-                ) or message_body.get("messageAttributes")
-                if not message_attributes:
-                    return None
-
-                if "_datadog" not in message_attributes:
-                    return None
-
-                datadog_attr = message_attributes["_datadog"]
-
-                if message_body.get("Type") == "Notification":
-                    if datadog_attr.get("Type") == "Binary":
-                        context_json = json.loads(
-                            base64.b64decode(datadog_attr["Value"]).decode()
-                        )
-                elif "StringValue" in datadog_attr:
-                    context_json = json.loads(datadog_attr["StringValue"])
-                elif "stringValue" in datadog_attr:
-                    context_json = json.loads(datadog_attr["stringValue"])
-                elif "BinaryValue" in datadog_attr:
-                    context_json = json.loads(datadog_attr["BinaryValue"].decode())
-                else:
-                    print(f"DEBUG: Unhandled datadog_attr format: {datadog_attr}")
-
-                return context_json
-
-            # Step 1: Create a message with some context in the message attributes
-            from ddtrace.internal.datastreams.processor import DataStreamsProcessor
-
-            processor_instance = DataStreamsProcessor()
-
-            with patch(
-                "ddtrace.internal.datastreams.botocore.get_datastreams_context",
-                updated_get_datastreams_context,
-            ), patch(
-                "ddtrace.internal.datastreams.data_streams_processor",
-                return_value=processor_instance,
-            ):
-
-                parent_ctx = processor_instance.new_pathway()
-
-                parent_ctx.set_checkpoint(
-                    ["direction:out", "topic:upstream-topic", "type:sqs"],
-                    now_sec=1640995200.0,
-                    payload_size=512,
-                )
-                parent_hash = parent_ctx.hash
-                encoded_parent_context = parent_ctx.encode_b64()
-
-                sqs_event = {
-                    "Records": [
-                        {
-                            "eventSource": "aws:sqs",
-                            "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test",
-                            "Body": "test message body",
-                            "messageAttributes": {
-                                "_datadog": {
-                                    "stringValue": json.dumps(
-                                        {
-                                            "dd-pathway-ctx-base64": encoded_parent_context
-                                        }
-                                    )
-                                }
-                            },
-                        }
-                    ]
-                }
-
-                # Step 2: Call the handler
-                @wrapper.datadog_lambda_wrapper
-                def lambda_handler(event, context):
-                    return {"statusCode": 200, "body": "processed"}
-
-                result = lambda_handler(sqs_event, get_mock_context())
-                self.assertEqual(result["statusCode"], 200)
-
-                # New context set after handler call
-                current_ctx = processor_instance._current_context.value
-                self.assertIsNotNone(
-                    current_ctx,
-                    "Data streams context should be set after processing SQS message",
-                )
-
-                # Step 3: Check that hash in this context is the child of the hash you passed
-                # Step 4: Check that the right checkpoint was produced during call to handler
-                # The buckets hold the aggregated stats for all checkpoints
-                found_sqs_checkpoint = False
-                for bucket_time, bucket in processor_instance._buckets.items():
-                    for aggr_key, stats in bucket.pathway_stats.items():
-                        edge_tags_str, hash_value, parent_hash_recorded = aggr_key
-                        edge_tags = edge_tags_str.split(",")
-
-                        if (
-                            "direction:in" in edge_tags
-                            and "topic:arn:aws:sqs:us-east-1:123456789012:test"
-                            in edge_tags
-                            and "type:sqs" in edge_tags
-                        ):
-                            found_sqs_checkpoint = True
-
-                            # EXPLICIT PARENT-CHILD HASH RELATIONSHIP TEST
-                            self.assertEqual(
-                                parent_hash_recorded,
-                                parent_hash,
-                                f"Parent hash must be preserved: "
-                                f"expected {parent_hash}, got {parent_hash_recorded}",
-                            )
-                            self.assertEqual(
-                                hash_value,
-                                current_ctx.hash,
-                                f"Child hash must match current context: "
-                                f"expected {current_ctx.hash}, got {hash_value}",
-                            )
-                            self.assertNotEqual(
-                                hash_value,
-                                parent_hash_recorded,
-                                f"Child hash ({hash_value}) must be different from "
-                                f"parent hash ({parent_hash_recorded}) - proves parent-child",
-                            )
-                            self.assertGreaterEqual(
-                                stats.payload_size.count,
-                                1,
-                                "Should have one payload size measurement",
-                            )
-
-                            break
-
-                self.assertTrue(
-                    found_sqs_checkpoint,
-                    "Should have found SQS consumption checkpoint in processor stats",
-                )
-
-            processor_instance.shutdown(timeout=0.1)
-
-    @patch.dict(os.environ, {"DD_DATA_STREAMS_ENABLED": "true"})
-    @patch("datadog_lambda.wrapper.set_dsm_context")
-    def test_set_dsm_context_called_when_enabled(self, mock_set_dsm_context):
-        @wrapper.datadog_lambda_wrapper
-        def lambda_handler(event, context):
-            return {"statusCode": 200, "body": "processed"}
-
-        lambda_event = {}
-        lambda_handler(lambda_event, get_mock_context())
-
-        mock_set_dsm_context.assert_called_once()
-
-    @patch("datadog_lambda.wrapper.set_dsm_context")
-    def test_set_dsm_context_not_called_when_disabled(self, mock_set_dsm_context):
-        # Ensure DD_DATA_STREAMS_ENABLED is not in environment
-        if "DD_DATA_STREAMS_ENABLED" in os.environ:
-            del os.environ["DD_DATA_STREAMS_ENABLED"]
-
-        @wrapper.datadog_lambda_wrapper
-        def lambda_handler(event, context):
-            return {"statusCode": 200, "body": "processed"}
-
-        lambda_event = {}
-        lambda_handler(lambda_event, get_mock_context())
-
-        mock_set_dsm_context.assert_not_called()
+    def test_set_dsm_context_called_when_DSM_and_tracing_enabled(self):
+        env_vars = {"DD_DATA_STREAMS_ENABLED": "true"}
+        with patch.dict(os.environ, env_vars):
+            with patch("datadog_lambda.wrapper.dd_tracing_enabled", True):
+                with patch(
+                    "datadog_lambda.wrapper.set_dsm_context"
+                ) as set_dsm_context_patch:
+
+                    @wrapper.datadog_lambda_wrapper
+                    def lambda_handler(event, context):
+                        return "ok"
+
+                    result = lambda_handler({}, get_mock_context())
+                    assert result == "ok"
+                    assert set_dsm_context_patch.called_once()
+
+    def test_set_dsm_context_not_called_when_only_DSM_enabled(self):
+        env_vars = {"DD_DATA_STREAMS_ENABLED": "true"}
+        with patch.dict(os.environ, env_vars):
+            with patch("datadog_lambda.wrapper.dd_tracing_enabled", False):
+                with patch(
+                    "datadog_lambda.wrapper.set_dsm_context"
+                ) as set_dsm_context_patch:
+
+                    @wrapper.datadog_lambda_wrapper
+                    def lambda_handler(event, context):
+                        return "ok"
+
+                    result = lambda_handler({}, get_mock_context())
+                    assert result == "ok"
+                    assert set_dsm_context_patch.call_count == 0
+
+    def test_set_dsm_context_not_called_when_only_tracing_enabled(self):
+        env_vars = {"DD_DATA_STREAMS_ENABLED": "false"}
+        with patch.dict(os.environ, env_vars):
+            with patch("datadog_lambda.wrapper.dd_tracing_enabled", True):
+                with patch(
+                    "datadog_lambda.wrapper.set_dsm_context"
+                ) as set_dsm_context_patch:
+
+                    @wrapper.datadog_lambda_wrapper
+                    def lambda_handler(event, context):
+                        return "ok"
+
+                    result = lambda_handler({}, get_mock_context())
+                    assert result == "ok"
+                    assert set_dsm_context_patch.call_count == 0
+
+    def test_set_dsm_context_not_called_when_tracing_and_DSM_disabled(self):
+        env_vars = {"DD_DATA_STREAMS_ENABLED": "false"}
+        with patch.dict(os.environ, env_vars):
+            with patch("datadog_lambda.wrapper.dd_tracing_enabled", True):
+                with patch(
+                    "datadog_lambda.wrapper.set_dsm_context"
+                ) as set_dsm_context_patch:
+
+                    @wrapper.datadog_lambda_wrapper
+                    def lambda_handler(event, context):
+                        return "ok"
+
+                    result = lambda_handler({}, get_mock_context())
+                    assert result == "ok"
+                    assert set_dsm_context_patch.call_count == 0
 
 
 class TestLambdaDecoratorSettings(unittest.TestCase):
