# SPDX-FileCopyrightText: 2022-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 import opentelemetry.trace import pytest from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.sdk.resources import SERVICE_NAME, Resource from opentelemetry.sdk.trace import TracerProvider from haystack.tracing.opentelemetry import OpenTelemetryTracer @pytest.fixture() def span_exporter() -> InMemorySpanExporter: return InMemorySpanExporter() @pytest.fixture def opentelemetry_tracer(span_exporter: InMemorySpanExporter) -> opentelemetry.trace.Tracer: # Service name is required for most backends resource = Resource(attributes={SERVICE_NAME: "haystack-testing"}) traceProvider = TracerProvider(resource=resource) processor = SimpleSpanProcessor(span_exporter) traceProvider.add_span_processor(processor) return traceProvider.get_tracer("my_test") class TestOpenTelemetryTracer: def test_opentelemetry_tracer( self, opentelemetry_tracer: opentelemetry.trace.Tracer, span_exporter: InMemorySpanExporter ) -> None: tracer = OpenTelemetryTracer(opentelemetry_tracer) with tracer.trace("test") as span: span.set_tag("key", "value") def test_tagging( self, opentelemetry_tracer: opentelemetry.trace.Tracer, span_exporter: InMemorySpanExporter ) -> None: tracer = OpenTelemetryTracer(opentelemetry_tracer) with tracer.trace("test", tags={"key1": "value1"}) as span: span.set_tag("key2", "value2") spans = list(span_exporter.get_finished_spans()) assert len(spans) == 1 assert spans[0].attributes == {"key1": "value1", "key2": "value2"} def test_current_span( self, opentelemetry_tracer: opentelemetry.trace.Tracer, span_exporter: InMemorySpanExporter ) -> None: tracer = OpenTelemetryTracer(opentelemetry_tracer) with tracer.trace("test"): current_span = tracer.current_span() assert tracer.current_span() is not None current_span.set_tag("key1", "value1") raw_span = current_span.raw_span() assert raw_span is not None assert isinstance(raw_span, opentelemetry.trace.Span) raw_span.set_attribute("key2", "value2") spans = list(span_exporter.get_finished_spans()) assert len(spans) == 1 assert spans[0].attributes == {"key1": "value1", "key2": "value2"} def test_tracing_complex_values( self, opentelemetry_tracer: opentelemetry.trace.Tracer, span_exporter: InMemorySpanExporter ) -> None: tracer = OpenTelemetryTracer(opentelemetry_tracer) with tracer.trace("test") as span: span.set_tag("key", {"a": 1, "b": [2, 3, 4]}) spans = list(span_exporter.get_finished_spans()) assert len(spans) == 1 assert spans[0].attributes == {"key": '{"a": 1, "b": [2, 3, 4]}'} def test_log_correlation_info(self, opentelemetry_tracer: opentelemetry.trace.Tracer) -> None: tracer = OpenTelemetryTracer(opentelemetry_tracer) with tracer.trace("test") as span: span.set_tag("key", "value") correlation_data = span.get_correlation_data_for_logs() assert correlation_data == { "trace_id": span.raw_span().get_span_context().trace_id, "span_id": span.raw_span().get_span_context().span_id, }