dd_service: Support DD_INTEGRATION_SERVICE

This change adds support in the dd_service checks for libraries which allow users to set per-integration service values with DD_INTEGRATION_SERVICE.
DataDog · Jul 3, 2023 · 4fab7a1 · 4fab7a1
1 parent 86f9a26
commit 4fab7a1
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/ddapm_test_agent/agent.py b/ddapm_test_agent/agent.py
@@ -2,6 +2,7 @@
 import atexit
 import base64
 from collections import OrderedDict
+from collections import defaultdict
 import json
 import logging
 import os
@@ -162,7 +163,7 @@ def __init__(self):
         # Token to be used if running test cases synchronously
         self._requests: List[Request] = []
         self._rc_server = RemoteConfigServer()
-        self._trace_failures: Dict[str, List[str]] = {"default": []}
+        self._trace_failures: Dict[str, List[str]] = defaultdict(list)
         self._trace_check_results_by_check: Dict[str, Dict[str, int]] = {}
         self._forward_endpoints: List[str] = [
             "/v0.4/traces",
@@ -758,7 +759,7 @@ async def check_failure_middleware(self, request: Request, handler: _Handler) ->
             msg = str(trace) + str(e)
             if request.app["pool_trace_check_failures"]:
                 log.info(f"Storing Trace Check Failure for Session Token: {token}.")
-                self._trace_failures[token] = [msg] if token not in self._trace_failures else self._trace_failures[token].append(msg)
+                self._trace_failures[token].append(msg)
             log.error(msg)
             return web.HTTPBadRequest(body=msg)
         else:
@@ -771,7 +772,7 @@ async def check_failure_middleware(self, request: Request, handler: _Handler) ->
                 msg = str(trace)
                 if request.app["pool_trace_check_failures"]:
                     log.info(f"Storing Trace Check Failure for Session Token: {token}.")
-                    self._trace_failures[token] = [msg] if token not in self._trace_failures else self._trace_failures[token].append(msg)
+                    self._trace_failures[token].append(msg)
                 log.error(msg)
                 if request.app["disable_error_responses"]:
                     return response

diff --git a/ddapm_test_agent/trace_checks.py b/ddapm_test_agent/trace_checks.py
@@ -182,8 +182,12 @@ def check(self, trace: List[Span], dd_config_env: dict) -> None:
                     )
                     return
 
-            if meta.get("component", "") != "":
-                dd_service = dd_config_env.get("DD_SERVICE", None)
+            component = meta.get("component", "")
+            if component != "":
+                dd_service = dd_config_env.get(f"DD_{component.upper()}_SERVICE", None)
+                if not dd_service:
+                    log.error("DD_SERVICE not set for component: %s. Args %s", component, dd_config_env)
+                dd_service = dd_service or dd_config_env.get("DD_SERVICE", None)
                 if dd_service is None:
                     self.fail(
                         json.dumps(dd_config_env, indent=4)