Use a better API to handle agent trace polymorphism

tests/ai_guard/test_ai_guard_sdk.py

@@ -2,7 +2,7 @@
 
 from utils import context, interfaces, scenarios, weblog, features
 from utils.dd_constants import SamplingPriority
-from utils.dd_types import DataDogSpan
+from utils.dd_types import DataDogLibrarySpan
 
 BLOCKING_HEADER: str = "X-AI-Guard-Block"
 MESSAGES: dict = {
@@ -60,7 +60,7 @@ def _assert_key(values: dict, key: str, value: object | None = None):
 @scenarios.ai_guard
 class Test_Evaluation:
     def _assert_span(self, action: str, messages: list, *, blocking: str):
-        def validate(span: DataDogSpan):
+        def validate(span: DataDogLibrarySpan):
             if span["resource"] != "ai_guard":
                 return False
 
@@ -201,7 +201,7 @@ def test_root_span_user_keep(self):
 @scenarios.ai_guard
 class Test_Full_Response_And_Tags:
     def _assert_span(self, response: dict, action: str):
-        def validate(span: DataDogSpan):
+        def validate(span: DataDogLibrarySpan):
             if span["resource"] != "ai_guard":
                 return False
 
@@ -242,7 +242,7 @@ def test_evaluation(self):
 @features.ai_guard
 @scenarios.default
 class Test_SDK_Disabled:
-    def _validate_no_ai_guard_span(self, span: DataDogSpan):
+    def _validate_no_ai_guard_span(self, span: DataDogLibrarySpan):
         assert span["resource"] != "ai_guard"
         return True
 
@@ -268,7 +268,7 @@ class Test_ContentParts:
     """Test AI Guard with multi-modal content parts (text + image_url)."""
 
     def _assert_span_with_content_parts(self, messages: list):
-        def validate(span: DataDogSpan):
+        def validate(span: DataDogLibrarySpan):
             if span["resource"] != "ai_guard":
                 return False
 

tests/apm_tracing_e2e/test_otel.py

@@ -1,5 +1,5 @@
 from utils import weblog, scenarios, interfaces, features
-from utils.dd_constants import TraceAgentPayloadFormat
+from utils.dd_types import DataDogAgentSpan
 
 
 @features.otel_api
@@ -26,20 +26,20 @@ def test_datadog_otel_span(self):
         assert len(spans) >= 2, "Agent did not submit the spans we want!"
 
         # Assert the parent span sent by the agent.
-        parent, parent_span_format = _get_span_by_resource(spans, "root-otel-name.dd-resource")
+        parent = _get_span_by_resource(spans, "root-otel-name.dd-resource")
         assert parent.get("parentID") is None
-        parent_meta = interfaces.agent.get_span_meta(parent, parent_span_format)
+        parent_meta = parent.meta
         if parent_meta["language"] != "jvm":  # Java OpenTelemetry API does not provide Span ID API
             assert parent.get("spanID") == "10000"
         assert parent_meta.get("attributes") == "values"
         assert parent_meta.get("error.message") == "testing_end_span_options"
-        parent_metrics = interfaces.agent.get_span_metrics(parent, parent_span_format)
+        parent_metrics = interfaces.agent.get_span_metrics(parent)
         assert parent_metrics["_dd.top_level"] == 1.0
         # Assert the child sent by the agent.
         # childName is no longer the operation name, rather the resource name
         # after remapping the OTel attributes to Datadog semantics
-        child, child_span_format = _get_span_by_resource(spans, "otel-name.dd-resource")
-        child_meta = interfaces.agent.get_span_meta(child, child_span_format)
+        child = _get_span_by_resource(spans, "otel-name.dd-resource")
+        child_meta = child.meta
         assert child.get("parentID") == parent.get("spanID")
         assert child.get("spanID") != "10000"
         assert child.get("duration") == "1000000000"
@@ -59,41 +59,37 @@ def test_distributed_otel_trace(self):
         assert len(spans) >= 3, "Agent did not submit the spans we want!"
 
         # Assert the parent span sent by the agent.
-        parent, parent_span_format = _get_span_by_resource(spans, "root-otel-name.dd-resource")
-        assert interfaces.agent.get_span_name(parent, parent_span_format) == "internal"
+        parent = _get_span_by_resource(spans, "root-otel-name.dd-resource")
+        assert parent.get_span_name() == "internal"
         assert parent.get("parentID") is None
-        parent_metrics = interfaces.agent.get_span_metrics(parent, parent_span_format)
+        parent_metrics = interfaces.agent.get_span_metrics(parent)
         assert parent_metrics["_dd.top_level"] == 1.0
 
         # Assert the Roundtrip child span sent by the agent, this span is created by an external OTel contrib package
-        roundtrip_span, roundtrip_span_format = _get_span_by_name(spans, "client.request")
-        assert interfaces.agent.get_span_name(roundtrip_span, roundtrip_span_format) == "client.request"
-        assert interfaces.agent.get_span_resource(roundtrip_span, roundtrip_span_format) == "HTTP GET"
+        roundtrip_span = _get_span_by_name(spans, "client.request")
+        assert roundtrip_span.get_span_name() == "client.request"
+        assert roundtrip_span.get_span_resource() == "HTTP GET"
         assert roundtrip_span.get("parentID") == parent.get("spanID")
 
         # Assert the Handler function child span sent by the agent.
-        handler_span, handler_span_format = _get_span_by_name(spans, "server.request")
-        assert interfaces.agent.get_span_resource(handler_span, handler_span_format) == "testOperation"
+        handler_span = _get_span_by_name(spans, "server.request")
+        assert handler_span.get_span_resource() == "testOperation"
         assert handler_span.get("parentID") == roundtrip_span.get("spanID")
 
         # Assert the spans received from the backend!
         spans = interfaces.backend.assert_request_spans_exist(self.req, query_filter="", retries=10)
         assert len(spans) == 3
 
 
-def _get_span_by_name(
-    spans: list[tuple[dict, TraceAgentPayloadFormat]], span_name: str
-) -> tuple[dict, TraceAgentPayloadFormat]:
-    for s, span_format in spans:
-        if interfaces.agent.get_span_name(s, span_format) == span_name:
-            return s, span_format
-    return {}, TraceAgentPayloadFormat.legacy
+def _get_span_by_name(spans: list[DataDogAgentSpan], span_name: str) -> DataDogAgentSpan:
+    for s in spans:
+        if s.get_span_name() == span_name:
+            return s
+    raise ValueError("Span not found")
 
 
-def _get_span_by_resource(
-    spans: list[tuple[dict, TraceAgentPayloadFormat]], resource_name: str
-) -> tuple[dict, TraceAgentPayloadFormat]:
-    for s, span_format in spans:
-        if interfaces.agent.get_span_resource(s, span_format) == resource_name:
-            return s, span_format
-    return {}, TraceAgentPayloadFormat.legacy
+def _get_span_by_resource(spans: list[DataDogAgentSpan], resource_name: str) -> DataDogAgentSpan:
+    for s in spans:
+        if s.get_span_resource() == resource_name:
+            return s
+    raise ValueError("Span not found")

tests/apm_tracing_e2e/test_process_tags.py

@@ -17,7 +17,7 @@ def setup_tracing_process_tags_svc(self):
     def check_tracing_process_tags(self, validate_process_tags_func: Callable):
         # Get all the spans from the agent
         found = False
-        for data, _, _ in interfaces.agent.get_traces(self.req):
+        for data, _ in interfaces.agent.get_traces(self.req):
             # Check that the agent managed to extract the process tags from the first chunk
             if "idxTracerPayloads" in data["request"]["content"]:
                 for payload in data["request"]["content"]["idxTracerPayloads"]:

tests/apm_tracing_e2e/test_single_span.py

@@ -5,8 +5,8 @@
     SINGLE_SPAN_SAMPLING_MECHANISM_VALUE,
     SINGLE_SPAN_SAMPLING_RATE,
     SINGLE_SPAN_SAMPLING_MAX_PER_SEC,
-    TraceAgentPayloadFormat,
 )
+from utils.dd_types import DataDogAgentSpan
 
 
 @rfc("ATI-2419")
@@ -32,12 +32,12 @@ def test_parent_span_is_single_span(self):
         assert len(spans) == 1, "Agent did not submit the spans we want!"
 
         # Assert the spans sent by the agent.
-        span, span_format = spans[0]
-        assert interfaces.agent.get_span_name(span, span_format) == "parent.span.single_span_submitted"
+        span = spans[0]
+        assert span.get_span_name() == "parent.span.single_span_submitted"
         assert span.get("parentID") is None
-        metrics = interfaces.agent.get_span_metrics(span, span_format)
+        metrics = interfaces.agent.get_span_metrics(span)
         assert metrics["_dd.top_level"] == 1.0
-        _assert_single_span_metrics(span, span_format)
+        _assert_single_span_metrics(span)
 
         # Assert the spans received from the backend!
         backend_spans = interfaces.backend.assert_single_spans_exist(self.req)
@@ -56,10 +56,10 @@ def test_child_span_is_single_span(self):
         assert len(spans) == 1, "Agent did not submit the spans we want!"
 
         # Assert the spans sent by the agent.
-        span, span_format = spans[0]
-        assert interfaces.agent.get_span_name(span, span_format) == "child.span.single_span_submitted"
+        span = spans[0]
+        assert span.get_span_name() == "child.span.single_span_submitted"
         assert span.get("parentID") is not None
-        _assert_single_span_metrics(span, span_format)
+        _assert_single_span_metrics(span)
 
         # Assert the spans received from the backend!
         backend_spans = interfaces.backend.assert_single_spans_exist(self.req)
@@ -79,8 +79,8 @@ def _assert_single_span_event(event: dict, name: str, *, is_root: bool):
         assert len(parent_id) > 0, f"In a child span the parent_id should be specified. Actual: {parent_id}"
 
 
-def _assert_single_span_metrics(span: dict, span_format: TraceAgentPayloadFormat):
-    metrics = interfaces.agent.get_span_metrics(span, span_format)
+def _assert_single_span_metrics(span: DataDogAgentSpan):
+    metrics = interfaces.agent.get_span_metrics(span)
     assert metrics[SAMPLING_PRIORITY_KEY] == -1  # due to the global sampling rate = 0
     assert metrics[SINGLE_SPAN_SAMPLING_RATE] == 1.0
     assert metrics[SINGLE_SPAN_SAMPLING_MECHANISM] == SINGLE_SPAN_SAMPLING_MECHANISM_VALUE

tests/appsec/rasp/test_api10.py

@@ -6,7 +6,7 @@
 import urllib.parse
 
 from utils import features, weblog, interfaces, scenarios, rfc, context
-from utils.dd_types import DataDogSpan
+from utils.dd_types import DataDogLibrarySpan
 
 from tests.appsec.rasp.utils import (
     find_series,
@@ -28,7 +28,7 @@ class API10:
     TAGS_EXPECTED: list[tuple[str, str]] = []
     TAGS_EXPECTED_METRIC: list[tuple[str, str]] = []
 
-    def validate(self, span: DataDogSpan):
+    def validate(self, span: DataDogLibrarySpan):
         if span.get("parent_id") not in (0, None):
             return None
 
@@ -44,7 +44,7 @@ def validate(self, span: DataDogSpan):
 
         return True
 
-    def validate_metric(self, span: DataDogSpan):
+    def validate_metric(self, span: DataDogLibrarySpan):
         for tag, expected in self.TAGS_EXPECTED_METRIC:
             # check also in meta to be safe
             assert tag in span["metrics"] or tag in span["meta"], f"Missing {tag} from span's meta/metrics"
@@ -291,7 +291,7 @@ def setup_api10_res_body(self):
             "/external_request", data=json.dumps(self.BODY), headers={"Content-Type": "application/json"}
         )
 
-    def validate_absence(self, span: DataDogSpan):
+    def validate_absence(self, span: DataDogLibrarySpan):
         if span.get("parent_id") not in (0, None):
             return None
 
@@ -322,7 +322,7 @@ def setup_api10_res_body(self):
             "/external_request", data=json.dumps(self.BODY), headers={"Content-Type": "application/json"}
         )
 
-    def validate_absence(self, span: DataDogSpan):
+    def validate_absence(self, span: DataDogLibrarySpan):
         if span.get("parent_id") not in (0, None):
             return None
 

tests/appsec/test_asm_standalone.py

@@ -9,7 +9,7 @@
 from utils.telemetry_utils import TelemetryUtils
 from utils._weblog import HttpResponse, _Weblog
 from utils import context, weblog, interfaces, scenarios, features, rfc, missing_feature, logger
-from utils.dd_types import DataDogSpan, DataDogTrace, TraceLibraryPayloadFormat
+from utils.dd_types import DataDogLibrarySpan, DataDogLibraryTrace, LibraryTraceFormat
 
 USER = "test"
 NEW_USER = "testnew"
@@ -34,9 +34,9 @@
 #   - The value can be a string to assert the value of the tag
 #   - The value can be a lambda function that will be used to assert the value of the tag (special case for _sampling_priority_v1)
 def assert_tags(
-    first_span: DataDogSpan, span: DataDogSpan, obj: str, expected_tags: dict[str, str | None | Callable]
+    first_span: DataDogLibrarySpan, span: DataDogLibrarySpan, obj: str, expected_tags: dict[str, str | None | Callable]
 ) -> bool:
-    def _assert_tags_value(span: DataDogSpan, obj: str, expected_tags: dict[str, str | None | Callable]):
+    def _assert_tags_value(span: DataDogLibrarySpan, obj: str, expected_tags: dict[str, str | None | Callable]):
         struct = span if obj is None else span[obj]
         for tag, value in expected_tags.items():
             if value is None:
@@ -61,8 +61,8 @@ def _assert_tags_value(span: DataDogSpan, obj: str, expected_tags: dict[str, str
         return False
 
 
-def _assert_trace_id(trace: DataDogTrace, span: DataDogSpan, trace_id: int) -> None:
-    if trace.format == TraceLibraryPayloadFormat.v10:
+def _assert_trace_id(trace: DataDogLibraryTrace, span: DataDogLibrarySpan, trace_id: int) -> None:
+    if trace.format == LibraryTraceFormat.v10:
         assert trace.trace_id_equals(trace_id)
     else:
         assert span.raw_span["trace_id"] == trace_id
@@ -137,7 +137,7 @@ def setup_no_appsec_upstream__no_asm_event__is_kept_with_priority_1__from_minus_
             )
 
     def fix_priority_lambda(
-        self, span: DataDogSpan, default_checks: dict[str, str | Callable | None]
+        self, span: DataDogLibrarySpan, default_checks: dict[str, str | Callable | None]
     ) -> dict[str, str | Callable | None]:
         if "_dd.appsec.s.req.headers" in span["meta"]:
             return {

tests/appsec/test_automated_login_events.py

@@ -12,7 +12,7 @@
 from utils import scenarios
 from utils import weblog
 from utils.dd_constants import Capabilities, SamplingPriority
-from utils.dd_types import DataDogSpan
+from utils.dd_types import DataDogLibrarySpan
 
 
 def login_data(username: str, password: str):
@@ -529,7 +529,7 @@ def setup_login_success_headers(self):
     def test_login_success_headers(self):
         # Validate that all relevant headers are included on user login success on extended mode
 
-        def validate_login_success_headers(span: DataDogSpan):
+        def validate_login_success_headers(span: DataDogLibrarySpan):
             if span.get("parent_id") not in (0, None):
                 return None
 
@@ -549,7 +549,7 @@ def setup_login_failure_headers(self):
     def test_login_failure_headers(self):
         # Validate that all relevant headers are included on user login failure on extended mode
 
-        def validate_login_failure_headers(span: DataDogSpan):
+        def validate_login_failure_headers(span: DataDogLibrarySpan):
             if span.get("parent_id") not in (0, None):
                 return None
 
@@ -1023,7 +1023,7 @@ def setup_login_success_headers(self):
     def test_login_success_headers(self):
         # Validate that all relevant headers are included on user login success on extended mode
 
-        def validate_login_success_headers(span: DataDogSpan):
+        def validate_login_success_headers(span: DataDogLibrarySpan):
             if span.get("parent_id") not in (0, None):
                 return None
 
@@ -1043,7 +1043,7 @@ def setup_login_failure_headers(self):
     def test_login_failure_headers(self):
         # Validate that all relevant headers are included on user login failure on extended mode
 
-        def validate_login_failure_headers(span: DataDogSpan):
+        def validate_login_failure_headers(span: DataDogLibrarySpan):
             if span.get("parent_id") not in (0, None):
                 return None
 
@@ -1054,7 +1054,7 @@ def validate_login_failure_headers(span: DataDogSpan):
         interfaces.library.validate_one_span(self.r_hdr_failure, validator=validate_login_failure_headers)
 
 
-def assert_priority(span: DataDogSpan, trace: list[DataDogSpan]):
+def assert_priority(span: DataDogLibrarySpan, trace: list[DataDogLibrarySpan]):
     if "_sampling_priority_v1" not in span["metrics"]:
         # some tracers like java only send the priority in the first and last span of the trace
         assert trace[0]["metrics"].get("_sampling_priority_v1") == SamplingPriority.USER_KEEP
@@ -1457,7 +1457,7 @@ def setup_login_success_headers(self):
     def test_login_success_headers(self):
         # Validate that all relevant headers are included on user login success on extended mode
 
-        def validate_login_success_headers(span: DataDogSpan):
+        def validate_login_success_headers(span: DataDogLibrarySpan):
             if span.get("parent_id") not in (0, None):
                 return None
 
@@ -1477,7 +1477,7 @@ def setup_login_failure_headers(self):
     def test_login_failure_headers(self):
         # Validate that all relevant headers are included on user login failure on extended mode
 
-        def validate_login_failure_headers(span: DataDogSpan):
+        def validate_login_failure_headers(span: DataDogLibrarySpan):
             if span.get("parent_id") not in (0, None):
                 return None