To pass the unittests

DataDog · Oct 30, 2023 · fbac793 · fbac793
1 parent fc8f36d
commit fbac793
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 47 deletions.
diff --git a/aws/logs_monitoring/parsing.py b/aws/logs_monitoring/parsing.py
@@ -538,7 +538,7 @@ def awslogs_handler(event, context, metadata):
     ):
         state_machine_arn = ""
         try:
-            state_machine_arn = get_state_machine_arn(logs)
+            state_machine_arn = get_state_machine_arn(json.loads(logs["logEvents"][0]["message"]))
             if state_machine_arn:  # not empty
                 metadata[DD_HOST] = state_machine_arn
         except Exception as e:
@@ -856,8 +856,7 @@ def normalize_events(events, metadata):
     return normalized
 
 
-def get_state_machine_arn(logs):
-    message = json.loads(logs["logEvents"][0]["message"])
+def get_state_machine_arn(message):
     if message.get("execution_arn") is not None:
         execution_arn = message["execution_arn"]
         arn_tokens = re.split(r"[:/\\]", execution_arn)

diff --git a/aws/logs_monitoring/tests/test_parsing.py b/aws/logs_monitoring/tests/test_parsing.py
@@ -969,58 +969,31 @@ def test_get_service_from_tags_default_to_source(self):
 
 class TestParsingStepFunctionLogs(unittest.TestCase):
     def test_get_state_machine_arn(self):
-        invalid_log_event = {
-            "logEvents": [
-                {
-                    "id": "31953106606966983378809025079804211143289615424298221568",
-                    "timestamp": 1609556645000,
-                    "message": {
+        invalid_sf_log_message =  {
                         "no_execution_arn": "xxxx/yyy"
-                    }
-                }
-            ]
+
         }
-        self.assertEqual(get_state_machine_arn(invalid_log_event), "")
-        normal_log_event = {
-            "logEvents": [
-                {
-                    "id": "31953106606966983378809025079804211143289615424298221568",
-                    "timestamp": 1609556645000,
-                    "message": {
+        self.assertEqual(get_state_machine_arn(invalid_sf_log_message), "")
+
+        normal_sf_log_message = {
                         "execution_arn": "arn:aws:states:sa-east-1:425362996713:express:my-Various-States:7f653fda-c79a-430b-91e2-3f97eb87cabb:862e5d40-a457-4ca2-a3c1-78485bd94d3f"
-                    }
-                }
-            ]
+
         }
-        self.assertEqual(get_state_machine_arn(normal_log_event),
-                         "arn:aws:states:sa-east-1:425362996713:express:my-Various-States")
-        forward_slash_log_event = {
-            "logEvents": [
-                {
-                    "id": "31953106606966983378809025079804211143289615424298221568",
-                    "timestamp": 1609556645000,
-                    "message": {
+        self.assertEqual(get_state_machine_arn(normal_sf_log_message),
+                         "arn:aws:states:sa-east-1:425362996713:stateMachine:my-Various-States")
+
+        forward_slash_sf_log_message = {
                         "execution_arn": "arn:aws:states:sa-east-1:425362996713:express:my-Various-States/7f653fda-c79a-430b-91e2-3f97eb87cabb:862e5d40-a457-4ca2-a3c1-78485bd94d3f"
-                    }
-                }
-            ]
         }
-        self.assertEqual(get_state_machine_arn(forward_slash_log_event),
-                         "arn:aws:states:sa-east-1:425362996713:express:my-Various-States")
-        back_slash_log_event = {
-            "logEvents": [
-                {
-                    "id": "31953106606966983378809025079804211143289615424298221568",
-                    "timestamp": 1609556645000,
-                    "message": {
+        self.assertEqual(get_state_machine_arn(forward_slash_sf_log_message),
+                          "arn:aws:states:sa-east-1:425362996713:stateMachine:my-Various-States")
+
+        back_slash_sf_log_message = {
                         "execution_arn": "arn:aws:states:sa-east-1:425362996713:express:my-Various-States\\7f653fda-c79a-430b-91e2-3f97eb87cabb:862e5d40-a457-4ca2-a3c1-78485bd94d3f"
                     }
-                }
-            ]
-        }
-        self.assertEqual(get_state_machine_arn(back_slash_log_event),
-                         "arn:aws:states:sa-east-1:425362996713:express:my-Various-States")
-
+        self.assertEqual(get_state_machine_arn(back_slash_sf_log_message),
+                          "arn:aws:states:sa-east-1:425362996713:stateMachine:my-Various-States")
+
 
 if __name__ == "__main__":
     unittest.main()