haystack/test/tracing/test_tracer.py

import builtins
import sys
from unittest.mock import Mock

import opentelemetry.trace
import pytest
from _pytest.monkeypatch import MonkeyPatch
from opentelemetry.sdk.resources import SERVICE_NAME, Resource
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter

from haystack.tracing.opentelemetry import OpenTelemetryTracer
from haystack.tracing.tracer import (
    NullTracer,
    NullSpan,
    enable_tracing,
    Tracer,
    disable_tracing,
    is_tracing_enabled,
    ProxyTracer,
    auto_enable_tracing,
    tracer,
)
from test.tracing.utils import SpyingTracer


class TestNullTracer:
    def test_tracing(self) -> None:
        assert isinstance(tracer.actual_tracer, NullTracer)

        # None of this raises
        with tracer.trace("operation", {"key": "value"}) as span:
            span.set_tag("key", "value")
            span.set_tags({"key": "value"})

        assert isinstance(tracer.current_span(), NullSpan)
        assert isinstance(tracer.current_span().raw_span(), NullSpan)


class TestProxyTracer:
    def test_tracing(self) -> None:
        spying_tracer = SpyingTracer()
        my_tracer = ProxyTracer(provided_tracer=spying_tracer)

        enable_tracing(spying_tracer)

        with my_tracer.trace("operation", {"key": "value"}) as span:
            span.set_tag("key", "value")
            span.set_tags({"key2": "value2"})

        assert len(spying_tracer.spans) == 1
        assert spying_tracer.spans[0].operation_name == "operation"
        assert spying_tracer.spans[0].tags == {"key": "value", "key2": "value2"}


class TestConfigureTracer:
    def test_enable_tracer(self) -> None:
        my_tracer = Mock(spec=Tracer)  # anything else than `NullTracer` works for this test

        enable_tracing(my_tracer)

        assert isinstance(tracer, ProxyTracer)
        assert tracer.actual_tracer is my_tracer
        assert is_tracing_enabled()

    def test_disable_tracing(self) -> None:
        my_tracker = Mock(spec=Tracer)  # anything else than `NullTracer` works for this test

        enable_tracing(my_tracker)
        assert tracer.actual_tracer is my_tracker

        disable_tracing()
        assert isinstance(tracer.actual_tracer, NullTracer)
        assert is_tracing_enabled() is False


class TestAutoEnableTracer:
    @pytest.fixture()
    def configured_opentelemetry_tracing(self) -> None:
        resource = Resource(attributes={SERVICE_NAME: "haystack-testing"})

        traceProvider = TracerProvider(resource=resource)
        processor = SimpleSpanProcessor(InMemorySpanExporter())
        traceProvider.add_span_processor(processor)

        # We can't uset `set_tracer_provider` here, because opentelemetry has a lock to only set it once
        opentelemetry.trace._TRACER_PROVIDER = traceProvider

        yield

        # unfortunately, there's no cleaner way to reset the global tracer provider
        opentelemetry.trace._TRACER_PROVIDER = None
        disable_tracing()

    def test_skip_auto_enable_tracer_if_already_configured(self) -> None:
        my_tracker = Mock(spec=Tracer)  # anything else than `NullTracer` works for this test
        enable_tracing(my_tracker)

        auto_enable_tracing()

        assert tracer.actual_tracer is my_tracker

    def test_skip_auto_enable_if_tracing_disabled_via_env(
        self, monkeypatch: MonkeyPatch, configured_opentelemetry_tracing: None
    ) -> None:
        monkeypatch.setenv("HAYSTACK_AUTO_TRACE_ENABLED", "false")

        old_tracer = tracer.actual_tracer

        auto_enable_tracing()

        assert tracer.actual_tracer is old_tracer

    def test_enable_opentelemetry_tracer(self, configured_opentelemetry_tracing: None) -> None:
        auto_enable_tracing()

        activated_tracer = tracer.actual_tracer
        assert isinstance(activated_tracer, OpenTelemetryTracer)
        assert is_tracing_enabled()

    def test_skip_enable_opentelemetry_tracer_if_import_error(self, monkeypatch: MonkeyPatch) -> None:
        monkeypatch.delitem(sys.modules, "opentelemetry", raising=False)
        monkeypatch.setattr(builtins, "__import__", Mock(side_effect=ImportError))
        auto_enable_tracing()

        activated_tracer = tracer.actual_tracer
        assert isinstance(activated_tracer, NullTracer)
        assert not is_tracing_enabled()