add tests

aws/logs_monitoring/steps/enrichment.py

@@ -54,11 +54,16 @@ def add_metadata_to_lambda_log(event, cache_layer):
     if not lambda_log_arn:
         return
 
+    # Function name is the seventh piece of the ARN
+    try:
+        function_name = lambda_log_arn.split(":")[6]
+    except IndexError:
+        logger.error(f"Failed to extract function name from ARN: {lambda_log_arn}")
+        return
+
     # Set Lambda ARN to "host"
     event[DD_HOST] = lambda_log_arn
 
-    # Function name is the seventh piece of the ARN
-    function_name = lambda_log_arn.split(":")[6]
     tags = [f"functionname:{function_name}"]
 
     # Get custom tags of the Lambda function
@@ -94,7 +99,10 @@ def add_metadata_to_lambda_log(event, cache_layer):
     tags = list(set(tags))
     tags.sort()  # Keep order deterministic
 
-    event[DD_CUSTOM_TAGS] = ",".join([event.get(DD_CUSTOM_TAGS)] + tags)
+    if custom_tags := event.get(DD_CUSTOM_TAGS):
+        event[DD_CUSTOM_TAGS] = custom_tags + ",".join(tags)
+    else:
+        event[DD_CUSTOM_TAGS] = ",".join(tags)
 
 
 def get_enriched_lambda_log_tags(log_event, cache_layer):
@@ -164,7 +172,7 @@ def extract_ddtags_from_message(event):
                 [
                     tag
                     for tag in event[DD_CUSTOM_TAGS].split(",")
-                    if not tag.startswith("service")
+                    if not tag.startswith("service:")
                 ]
             )
 

aws/logs_monitoring/tests/approved_files/TestLambdaMetadataEnrichment.test_lambda_event_bad_arn.approved.json

@@ -0,0 +1,5 @@
+{
+    "lambda": {
+        "arn": "bad_arn"
+    }
+}

aws/logs_monitoring/tests/approved_files/TestLambdaMetadataEnrichment.test_lambda_event_w_custom_tags_env.approved.json

@@ -0,0 +1,8 @@
+{
+    "ddtags": "account_id:123456789012,aws_account:123456789012,env:customtags_env,functionname:my-function,region:us-east-1,service:my-function",
+    "host": "arn:aws:lambda:us-east-1:123456789012:function:my-function",
+    "lambda": {
+        "arn": "arn:aws:lambda:us-east-1:123456789012:function:my-function"
+    },
+    "service": "my-function"
+}

aws/logs_monitoring/tests/approved_files/TestLambdaMetadataEnrichment.test_lambda_event_w_custom_tags_w_service.approved.json

@@ -0,0 +1,8 @@
+{
+    "ddtags": "account_id:123456789012,aws_account:123456789012,functionname:my-function,region:us-east-1",
+    "host": "arn:aws:lambda:us-east-1:123456789012:function:my-function",
+    "lambda": {
+        "arn": "arn:aws:lambda:us-east-1:123456789012:function:my-function"
+    },
+    "service": "my_service"
+}

aws/logs_monitoring/tests/approved_files/TestLambdaMetadataEnrichment.test_lambda_event_w_custom_tags_wo_service.approved.json

@@ -0,0 +1,8 @@
+{
+    "ddtags": "account_id:123456789012,aws_account:123456789012,functionname:my-function,region:us-east-1,service:customtags_service",
+    "host": "arn:aws:lambda:us-east-1:123456789012:function:my-function",
+    "lambda": {
+        "arn": "arn:aws:lambda:us-east-1:123456789012:function:my-function"
+    },
+    "service": "customtags_service"
+}

aws/logs_monitoring/tests/approved_files/TestLambdaMetadataEnrichment.test_lambda_event_w_service.approved.json

@@ -0,0 +1,8 @@
+{
+    "ddtags": "account_id:123456789012,aws_account:123456789012,functionname:my-function,region:us-east-1",
+    "host": "arn:aws:lambda:us-east-1:123456789012:function:my-function",
+    "lambda": {
+        "arn": "arn:aws:lambda:us-east-1:123456789012:function:my-function"
+    },
+    "service": "my_service"
+}

aws/logs_monitoring/tests/approved_files/TestLambdaMetadataEnrichment.test_lambda_event_w_service_and_ddtags.approved.json

@@ -0,0 +1,8 @@
+{
+    "ddtags": "service:ddtags_serviceaccount_id:123456789012,aws_account:123456789012,functionname:my-function,region:us-east-1",
+    "host": "arn:aws:lambda:us-east-1:123456789012:function:my-function",
+    "lambda": {
+        "arn": "arn:aws:lambda:us-east-1:123456789012:function:my-function"
+    },
+    "service": "my_service"
+}

aws/logs_monitoring/tests/approved_files/TestLambdaMetadataEnrichment.test_lambda_event_wo_service.approved.json

@@ -0,0 +1,8 @@
+{
+    "ddtags": "account_id:123456789012,aws_account:123456789012,functionname:my-function,region:us-east-1,service:my-function",
+    "host": "arn:aws:lambda:us-east-1:123456789012:function:my-function",
+    "lambda": {
+        "arn": "arn:aws:lambda:us-east-1:123456789012:function:my-function"
+    },
+    "service": "my-function"
+}

aws/logs_monitoring/tests/approved_files/TestLambdaMetadataEnrichment.test_non_lambda_event.approved.json

@@ -0,0 +1,3 @@
+{
+    "lambda": {}
+}

aws/logs_monitoring/tests/test_enrichment.py

@@ -1,6 +1,12 @@
-import unittest
 import json
+
+import unittest
+from unittest.mock import MagicMock
+
+from caching.cache_layer import CacheLayer
+from approvaltests.approvals import verify_as_json
 from steps.enrichment import (
+    add_metadata_to_lambda_log,
     extract_host_from_cloudtrails,
     extract_host_from_guardduty,
     extract_host_from_route53,
@@ -153,5 +159,99 @@ def test_parse_source_route53(self):
         self.assertEqual(event["host"], "i-99999999")
 
 
+class TestLambdaMetadataEnrichment(unittest.TestCase):
+    def test_empty_event(self):
+        cache_layer = CacheLayer("")
+        event = {}
+        add_metadata_to_lambda_log(event, cache_layer)
+
+        self.assertEqual(event, {})
+
+    def test_non_lambda_event(self):
+        cache_layer = CacheLayer("")
+        event = {"lambda": {}}
+        add_metadata_to_lambda_log(event, cache_layer)
+
+        verify_as_json(event)
+
+    def test_lambda_event_bad_arn(self):
+        cache_layer = CacheLayer("")
+        event = {"lambda": {"arn": "bad_arn"}}
+        add_metadata_to_lambda_log(event, cache_layer)
+        verify_as_json(event)
+
+    def test_lambda_event_wo_service(self):
+        cache_layer = CacheLayer("")
+        event = {
+            "lambda": {
+                "arn": "arn:aws:lambda:us-east-1:123456789012:function:my-function"
+            }
+        }
+        add_metadata_to_lambda_log(event, cache_layer)
+        verify_as_json(event)
+
+    def test_lambda_event_w_custom_tags_wo_service(self):
+        cache_layer = CacheLayer("")
+        cache_layer._lambda_cache.get = MagicMock(
+            return_value=["service:customtags_service"]
+        )
+        event = {
+            "lambda": {
+                "arn": "arn:aws:lambda:us-east-1:123456789012:function:my-function"
+            }
+        }
+        add_metadata_to_lambda_log(event, cache_layer)
+        verify_as_json(event)
+
+    def test_lambda_event_w_custom_tags_w_service(self):
+        cache_layer = CacheLayer("")
+        cache_layer._lambda_cache.get = MagicMock(
+            return_value=["service:customtags_service"]
+        )
+        event = {
+            "lambda": {
+                "arn": "arn:aws:lambda:us-east-1:123456789012:function:my-function"
+            },
+            "service": "my_service",
+        }
+        add_metadata_to_lambda_log(event, cache_layer)
+        verify_as_json(event)
+
+    def test_lambda_event_w_service(self):
+        cache_layer = CacheLayer("")
+        event = {
+            "lambda": {
+                "arn": "arn:aws:lambda:us-east-1:123456789012:function:my-function"
+            },
+            "service": "my_service",
+        }
+        add_metadata_to_lambda_log(event, cache_layer)
+        verify_as_json(event)
+
+    def test_lambda_event_w_service_and_ddtags(self):
+        cache_layer = CacheLayer("")
+        event = {
+            "lambda": {
+                "arn": "arn:aws:lambda:us-east-1:123456789012:function:my-function"
+            },
+            "service": "my_service",
+            "ddtags": "service:ddtags_service",
+        }
+        add_metadata_to_lambda_log(event, cache_layer)
+        verify_as_json(event)
+
+    def test_lambda_event_w_custom_tags_env(self):
+        cache_layer = CacheLayer("")
+        cache_layer._lambda_cache.get = MagicMock(return_value=["env:customtags_env"])
+        event = {
+            "lambda": {
+                "arn": "arn:aws:lambda:us-east-1:123456789012:function:my-function"
+            },
+            "ddtags": "env:none",
+        }
+        add_metadata_to_lambda_log(event, cache_layer)
+        verify_as_json(event)
+
+
 if __name__ == "__main__":
     unittest.main()