fix a few things

DataDog · Oct 15, 2024 · 3941664 · 3941664
1 parent eeda013
commit 3941664
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 21 deletions.
diff --git a/utils/_context/_scenarios/integrations.py b/utils/_context/_scenarios/integrations.py
@@ -53,15 +53,25 @@ def configure(self, config):
 
 
 class AWSIntegrationsScenario(EndToEndScenario):
-    def __init__(self) -> None:
+    def __init__(
+        self,
+        name="INTEGRATIONS_AWS",
+        doc="Spawns tracer, and agent. Test AWS integrations.",
+        include_kafka=False,
+        include_rabbitmq=False,
+        include_buddies=False,
+    ) -> None:
         super().__init__(
-            "INTEGRATIONS_AWS",
+            name,
             weblog_env={
                 "DD_TRACE_SPAN_ATTRIBUTE_SCHEMA": "v1",
                 "AWS_ACCESS_KEY_ID": "my-access-key",
                 "AWS_SECRET_ACCESS_KEY": "my-access-key",
             },
-            doc="Spawns tracer, and agent. Test AWS integrations.",
+            doc=doc,
+            include_kafka=include_kafka,
+            include_rabbitmq=include_rabbitmq,
+            include_buddies=include_buddies,
             scenario_groups=[ScenarioGroup.INTEGRATIONS, ScenarioGroup.ESSENTIALS],
         )
         # Since we are using real AWS queues / topics, we need a unique message to ensure we aren't consuming messages
@@ -72,27 +82,16 @@ def __init__(self) -> None:
     def configure(self, config):
         super().configure(config)
         if not self.replay:
-            self._check_aws_variables(self)
+            self._check_aws_variables()
         self.unique_id = _get_unique_id(self.replay, self.host_log_folder)
 
 
 class CrossedTracingLibraryScenario(AWSIntegrationsScenario):
     def __init__(self) -> None:
         super().__init__(
             "CROSSED_TRACING_LIBRARIES",
-            weblog_env={
-                "DD_TRACE_API_VERSION": "v0.4",
-                "AWS_ACCESS_KEY_ID": "my-access-key",
-                "AWS_SECRET_ACCESS_KEY": "my-access-key",
-            },
             include_kafka=True,
             include_buddies=True,
             include_rabbitmq=True,
-            doc="Spawns a buddy for each supported language of APM",
-            scenario_groups=[ScenarioGroup.INTEGRATIONS, ScenarioGroup.ESSENTIALS],
+            doc="Spawns a buddy for each supported language of APM, requires AWS authentication.",
         )
-
-        # Since we are using real AWS queues / topics, we need a unique message to ensure we aren't consuming messages
-        # from other tests. This time hash is added to the message, test consumers only stops once finding the specific
-        # message.
-        self.unique_id = None
diff --git a/utils/_context/containers.py b/utils/_context/containers.py
@@ -341,7 +341,7 @@ def collect_logs(self):
         if os.environ.get("AWS_SESSION_TOKEN"):
             keys.append(bytearray(os.environ["AWS_SESSION_TOKEN"], "utf-8"))
         if os.environ.get("AWS_SECURITY_TOKEN"):
-            keys.append(bytearray(os.environ["AWS_SESSION_TOKEN"], "utf-8"))
+            keys.append(bytearray(os.environ["AWS_SECURITY_TOKEN"], "utf-8"))
 
         # set by CI runner
         if os.environ.get("SYSTEM_TESTS_AWS_ACCESS_KEY_ID"):
@@ -702,10 +702,7 @@ def configure(self, replay):
         appsec_rules_version = self.image.env.get("SYSTEM_TESTS_APPSEC_EVENT_RULES_VERSION", "0.0.0")
         self.appsec_rules_version = LibraryVersion("appsec_rules", appsec_rules_version).version
 
-        self.environment["AWS_ACCESS_KEY_ID"] = os.environ.get("AWS_ACCESS_KEY_ID", "")
-        self.environment["AWS_SECRET_ACCESS_KEY"] = os.environ.get("AWS_SECRET_ACCESS_KEY", "")
-        self.environment["AWS_DEFAULT_REGION"] = os.environ.get("AWS_DEFAULT_REGION", "")
-        self.environment["AWS_REGION"] = os.environ.get("AWS_REGION", "")
+        _set_aws_auth_environment(self)
 
         self._library = LibraryVersion(
             self.image.env.get("SYSTEM_TESTS_LIBRARY", None), self.image.env.get("SYSTEM_TESTS_LIBRARY_VERSION", None),