mirror of
https://github.com/deepset-ai/haystack.git
synced 2025-07-08 09:31:31 +00:00
94 lines
3.5 KiB
Python
94 lines
3.5 KiB
Python
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
|
#
|
|
# 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,
|
|
}
|