deepset-ai--haystack/test/tracing/test_logging_tracer.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

import logging
import threading

import pytest

from haystack import Pipeline, component, tracing
from haystack.tracing.logging_tracer import LoggingTracer


@component
class Hello:
    @component.output_types(output=str)
    def run(self, word: str | None) -> dict[str, str]:
        return {"output": f"Hello, {word}!"}


@component
class FailingComponent:
    @component.output_types(output=str)
    def run(self, word: str | None) -> dict[str, str]:
        raise Exception("Failing component")


class TestLoggingTracer:
    def test_init(self) -> None:
        tracer = LoggingTracer()
        assert tracer.tags_color_strings == {}

        tracer = LoggingTracer(tags_color_strings={"tag_name": "color_string"})
        assert tracer.tags_color_strings == {"tag_name": "color_string"}

    def test_logging_tracer(self, caplog: pytest.LogCaptureFixture) -> None:
        tracer = LoggingTracer()

        caplog.set_level(logging.DEBUG)
        with tracer.trace("test") as span:
            span.set_tag("key", "value")

        assert "Operation: test" in caplog.text
        assert "key=value" in caplog.text
        assert len(caplog.records) == 2

        # structured logging - LoggingTracer dynamically adds these attributes to LogRecord
        assert caplog.records[0].operation_name == "test"  # type: ignore[attr-defined]
        assert caplog.records[1].tag_name == "key"  # type: ignore[attr-defined]
        assert caplog.records[1].tag_value == "value"  # type: ignore[attr-defined]

    def test_concurrent_spans_do_not_interleave(self, caplog: pytest.LogCaptureFixture) -> None:
        """Spans closing from parallel threads keep each span's operation-name and tag records contiguous."""
        tracer = LoggingTracer()
        caplog.set_level(logging.DEBUG)

        num_threads = 4
        num_tags = 3
        # Line every thread up so their span-exit emissions race, maximizing the chance of interleaving without a lock.
        barrier = threading.Barrier(num_threads)

        def emit_span(idx: int) -> None:
            with tracer.trace(f"op-{idx}") as span:
                barrier.wait()
                for tag_i in range(num_tags):
                    span.set_tag(f"tag-{tag_i}", str(idx))

        threads = [threading.Thread(target=emit_span, args=(i,)) for i in range(num_threads)]
        for thread in threads:
            thread.start()
        for thread in threads:
            thread.join()

        # Bucket each tag record under the operation record that precedes it.
        groups: list[tuple[str, list[str]]] = []
        for record in caplog.records:
            if hasattr(record, "operation_name"):
                groups.append((record.operation_name, []))
            elif hasattr(record, "tag_name") and groups:
                groups[-1][1].append(record.tag_value)  # type: ignore[attr-defined]

        # One contiguous group per span, and every tag inside a group belongs to that span (no interleaving).
        assert len(groups) == num_threads
        for operation_name, tag_values in groups:
            idx = operation_name.removeprefix("op-")
            assert tag_values == [idx] * num_tags

    def test_tracing_complex_values(self, caplog: pytest.LogCaptureFixture) -> None:
        tracer = LoggingTracer()

        caplog.set_level(logging.DEBUG)

        with tracer.trace("test") as span:
            span.set_tag("key", {"a": 1, "b": [2, 3, 4]})

        assert "Operation: test" in caplog.text
        assert 'key={"a": 1, "b": [2, 3, 4]}' in caplog.text
        assert len(caplog.records) == 2

        # structured logging - LoggingTracer dynamically adds these attributes to LogRecord
        assert caplog.records[0].operation_name == "test"  # type: ignore[attr-defined]
        assert caplog.records[1].tag_name == "key"  # type: ignore[attr-defined]
        assert caplog.records[1].tag_value == '{"a": 1, "b": [2, 3, 4]}'  # type: ignore[attr-defined]

    def test_apply_color_strings(self, caplog: pytest.LogCaptureFixture) -> None:
        tracer = LoggingTracer(tags_color_strings={"key": "color_string"})

        caplog.set_level(logging.DEBUG)

        with tracer.trace("test") as span:
            span.set_tag("key", "value")

        assert "color_string" in caplog.text

    def test_logging_pipeline(self, caplog: pytest.LogCaptureFixture) -> None:
        pipeline = Pipeline()
        pipeline.add_component("hello", Hello())
        pipeline.add_component("hello2", Hello())
        pipeline.connect("hello.output", "hello2.word")

        tracing.enable_tracing(LoggingTracer())
        caplog.set_level(logging.DEBUG)

        pipeline.run(data={"word": "world"})

        records = caplog.records

        assert any(
            record.operation_name == "haystack.component.run" for record in records if hasattr(record, "operation_name")
        )
        assert any(
            record.operation_name == "haystack.pipeline.run" for record in records if hasattr(record, "operation_name")
        )

        tags_records = [record for record in records if hasattr(record, "tag_name")]
        assert any(record.tag_name == "haystack.component.name" for record in tags_records)
        assert any(record.tag_value == "hello" for record in tags_records)  # type: ignore[attr-defined]

        tracing.disable_tracing()

    def test_logging_pipeline_with_content_tracing(self, caplog: pytest.LogCaptureFixture) -> None:
        pipeline = Pipeline()
        pipeline.add_component("hello", Hello())

        tracing.tracer.is_content_tracing_enabled = True
        tracing.enable_tracing(LoggingTracer())

        caplog.set_level(logging.DEBUG)

        pipeline.run(data={"word": "world"})
        records = caplog.records

        tags_records = [record for record in records if hasattr(record, "tag_name")]

        input_tag_value = [
            record.tag_value  # type: ignore[attr-defined]
            for record in tags_records
            if record.tag_name == "haystack.component.input"
        ][0]
        assert input_tag_value == '{"word": "world"}'

        output_tag_value = [
            record.tag_value  # type: ignore[attr-defined]
            for record in tags_records
            if record.tag_name == "haystack.component.output"
        ][0]
        assert output_tag_value == '{"output": "Hello, world!"}'

        tracing.tracer.is_content_tracing_enabled = False
        tracing.disable_tracing()

    def test_logging_pipeline_on_failure(self, caplog: pytest.LogCaptureFixture) -> None:
        """
        Test that the LoggingTracer also logs events when a component fails.
        """
        pipeline = Pipeline()
        pipeline.add_component("failing_component", FailingComponent())

        tracing.enable_tracing(LoggingTracer())
        caplog.set_level(logging.DEBUG)

        try:
            pipeline.run(data={"word": "world"})
        except Exception:
            pass

        records = caplog.records

        assert any(
            record.operation_name == "haystack.component.run" for record in records if hasattr(record, "operation_name")
        )
        assert any(
            record.operation_name == "haystack.pipeline.run" for record in records if hasattr(record, "operation_name")
        )

        tags_records = [record for record in records if hasattr(record, "tag_name")]
        assert any(record.tag_name == "haystack.component.name" for record in tags_records)
        assert any(record.tag_value == "failing_component" for record in tags_records)  # type: ignore[attr-defined]

        tracing.disable_tracing()