Updated tests to check for new tag

DataDog · Aug 1, 2024 · dbd763b · dbd763b
1 parent 2e7ff49
commit dbd763b
Show file tree

Hide file tree

Showing 4 changed files with 3 additions and 10 deletions.
diff --git a/aws/logs_monitoring/settings.py b/aws/logs_monitoring/settings.py
@@ -58,13 +58,6 @@ def get_env_var(envvar, default, boolean=False):
 #
 DD_FORWARD_LOG = get_env_var("DD_FORWARD_LOG", "true", boolean=True)
 
-## @param DD_STEP_FUNCTION_TRACE_ENABLED - boolean - optional - default: true
-## Set this variable to `True` to enable log to trace conversion for Step Functions.
-#
-DD_STEP_FUNCTION_TRACE_ENABLED = get_env_var(
-    "DD_STEP_FUNCTION_TRACE_ENABLED", default="false", boolean=True
-)
-
 ## @param DD_USE_TCP - boolean - optional -default: false
 ## Change this value to `true` to send your logs and metrics using the TCP network client
 ## By default, it uses the HTTP client.

diff --git a/aws/logs_monitoring/steps/handlers/awslogs_handler.py b/aws/logs_monitoring/steps/handlers/awslogs_handler.py
@@ -21,7 +21,6 @@
     DD_SOURCE,
     DD_HOST,
     DD_CUSTOM_TAGS,
-    DD_STEP_FUNCTION_TRACE_ENABLED,
 )
 
 RDS_REGEX = re.compile("/aws/rds/(instance|cluster)/(?P<host>[^/]+)/(?P<name>[^/]+)")
@@ -185,7 +184,7 @@ def handle_step_function_source(self):
                 + ",".join(formatted_stepfunctions_tags)
             )
 
-        if DD_STEP_FUNCTION_TRACE_ENABLED:
+        if os.environ.get("DD_STEP_FUNCTION_TRACE_ENABLED", "false").lower() == "true":
             self.metadata[DD_CUSTOM_TAGS] = ",".join(
                 [self.metadata.get(DD_CUSTOM_TAGS, [])]
                 + ["dd_step_function_trace_enabled:true"]

diff --git a/...ogsHandler.test_awslogs_handler_step_functions_tags_added_properly.metadata.approved.json b/...ogsHandler.test_awslogs_handler_step_functions_tags_added_properly.metadata.approved.json
@@ -1,6 +1,6 @@
 {
     "ddsource": "stepfunction",
-    "ddtags": "env:dev,test_tag_key:test_tag_value",
+    "ddtags": "env:dev,test_tag_key:test_tag_value,dd_step_function_trace_enabled:true",
     "host": "/aws/vendedlogs/states/logs-to-traces-sequential-Logs",
     "service": "stepfunction"
 }
diff --git a/aws/logs_monitoring/tests/test_awslogs_handler.py b/aws/logs_monitoring/tests/test_awslogs_handler.py
@@ -71,6 +71,7 @@ def test_awslogs_handler_rds_postgresql(self, mock_cache_init):
 
     @patch("caching.cloudwatch_log_group_cache.CloudwatchLogGroupTagsCache.__init__")
     @patch("caching.cloudwatch_log_group_cache.send_forwarder_internal_metrics")
+    @patch.dict("os.environ", {"DD_STEP_FUNCTION_TRACE_ENABLED": "true"})
     def test_awslogs_handler_step_functions_tags_added_properly(
         self,
         mock_forward_metrics,