Add more unit test

DataDog · Nov 4, 2024 · f951348 · f951348
1 parent a3647c5
commit f951348
Showing 1 changed file with 70 additions and 1 deletion.
diff --git a/aws/logs_monitoring/tests/test_awslogs_handler.py b/aws/logs_monitoring/tests/test_awslogs_handler.py
@@ -24,7 +24,7 @@
     },
 )
 env_patch.start()
-from aws.logs_monitoring.settings import DD_SOURCE
+from aws.logs_monitoring.settings import DD_HOST, DD_SOURCE
 from steps.handlers.awslogs_handler import AwsLogsHandler
 from steps.handlers.aws_attributes import AwsAttributes
 from caching.cache_layer import CacheLayer
@@ -124,6 +124,75 @@ def test_awslogs_handler_step_functions_tags_added_properly(
             awslogs_handler.metadata[DD_SOURCE], AwsEventSource.STEPFUNCTION.value
         )
 
+        # SF customized log group
+        eventFromCustomizedLogGroup = {
+            "awslogs": {
+                "data": base64.b64encode(
+                    gzip.compress(
+                        bytes(
+                            json.dumps(
+                                {
+                                    "messageType": "DATA_MESSAGE",
+                                    "owner": "425362996713",
+                                    "logGroup": "test/logs",
+                                    "logStream": "states/logs-to-traces-sequential/2022-11-10-15-50/7851b2d9",
+                                    "subscriptionFilters": ["testFilter"],
+                                    "logEvents": [
+                                        {
+                                            "id": "37199773595581154154810589279545129148442535997644275712",
+                                            "timestamp": 1668095539607,
+                                            "message": '{"execution_arn": "arn:aws:states:us-east-1:12345678910:execution:StepFunction1:ccccccc-d1da-4c38-b32c-2b6b07d713fa","redrive_count": "0"}',
+                                        }
+                                    ],
+                                }
+                            ),
+                            "utf-8",
+                        )
+                    )
+                )
+            }
+        }
+        awslogs_handler.handle(event)
+        self.assertEqual(
+            awslogs_handler.metadata[DD_SOURCE], AwsEventSource.STEPFUNCTION.value
+        )
+        self.assertEqual(
+            awslogs_handler.metadata[DD_HOST],
+            "arn:aws:states:us-east-1:12345678910:statemachine:stepfunction1",
+        )
+
+        eventFromCustomizedLogGroup = {
+            "awslogs": {
+                "data": base64.b64encode(
+                    gzip.compress(
+                        bytes(
+                            json.dumps(
+                                {
+                                    "messageType": "DATA_MESSAGE",
+                                    "owner": "425362996713",
+                                    "logGroup": "test/logs",
+                                    "logStream": "states/logs-to-traces-sequential/2022-11-10-15-50/7851b2d9",
+                                    "subscriptionFilters": ["testFilter"],
+                                    "logEvents": [
+                                        {
+                                            "id": "37199773595581154154810589279545129148442535997644275712",
+                                            "timestamp": 1668095539607,
+                                            "message": '{"execution_arn": "arn:aws:states:us-east-1:12345678910:execution:StepFunction1:ccccccc-d1da-4c38-b32c-2b6b07d713fa","redrive_count": "0"}',
+                                        }
+                                    ],
+                                }
+                            ),
+                            "utf-8",
+                        )
+                    )
+                )
+            }
+        }
+        awslogs_handler.handle(eventFromCustomizedLogGroup)
+        self.assertNotEqual(
+            awslogs_handler.metadata[DD_SOURCE], AwsEventSource.STEPFUNCTION.value
+        )
+
     def test_process_lambda_logs(self):
         # Non Lambda log
         stepfunction_loggroup = {