chore: import upstream snapshot with attribution

test/tracing/utils.py

@@ -0,0 +1,54 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import contextlib
+import dataclasses
+import uuid
+from collections.abc import Iterator
+from typing import Any
+
+from haystack.tracing import Span, Tracer
+
+
+@dataclasses.dataclass
+class SpyingSpan(Span):
+    operation_name: str
+    parent_span: Span | None = None
+
+    tags: dict[str, Any] = dataclasses.field(default_factory=dict)
+
+    trace_id: str | None = dataclasses.field(default_factory=lambda: str(uuid.uuid4()))
+    span_id: str | None = dataclasses.field(default_factory=lambda: str(uuid.uuid4()))
+
+    def set_tag(self, key: str, value: Any) -> None:
+        self.tags[key] = value
+
+    def get_correlation_data_for_logs(self) -> dict[str, Any]:
+        return {"trace_id": self.trace_id, "span_id": self.span_id}
+
+    def set_content_tag(self, key: str, value: Any) -> None:
+        """
+        Set a content tag, but only if content tracing is enabled in the tracer.
+        """
+        self.set_tag(key, value)
+
+
+class SpyingTracer(Tracer):
+    def current_span(self) -> Span | None:
+        return self.spans[-1] if self.spans else None
+
+    def __init__(self) -> None:
+        self.spans: list[SpyingSpan] = []
+
+    @contextlib.contextmanager
+    def trace(
+        self, operation_name: str, tags: dict[str, Any] | None = None, parent_span: Span | None = None
+    ) -> Iterator[Span]:
+        new_span = SpyingSpan(operation_name, parent_span)
+        for key, value in (tags or {}).items():
+            new_span.set_tag(key, value)
+
+        self.spans.append(new_span)
+
+        yield new_span