video-image-production-build/tests/observability/test_collector.py

from __future__ import annotations

from production_build_agents.observability import collector as collector_module
from production_build_agents.observability.collector import MessageToolCollector
from production_build_agents.observability.readable_events import (
    BusinessReadableFlowCollector,
)


def test_message_tool_collector_disables_automatic_modules(
    monkeypatch,
) -> None:
    calls: list[tuple[str, object]] = []
    handle = object()
    monkeypatch.setattr(
        collector_module,
        "_start_business_collector",
        lambda run_handle: (
            calls.append(("start", run_handle)) or "token"
        ),
    )
    monkeypatch.setattr(
        collector_module,
        "_stop_business_collector",
        lambda token: calls.append(("stop", token)),
    )
    collector = MessageToolCollector(
        warn_once=lambda *_args: None,
        warning_key="test",
    )

    collector.start(handle, enabled=True)
    collector.stop()

    assert calls == [
        ("start", handle),
        ("stop", "token"),
    ]


def test_message_tool_collector_failure_never_escapes(
    monkeypatch,
) -> None:
    warnings: list[tuple[str, str]] = []

    def fail_start(*_args, **_kwargs):
        raise RuntimeError("collector unavailable")

    monkeypatch.setattr(
        collector_module,
        "_start_business_collector",
        fail_start,
    )
    collector = MessageToolCollector(
        warn_once=lambda key, exc: warnings.append(
            (key, type(exc).__name__)
        ),
        warning_key="test",
    )

    collector.start(object(), enabled=True)
    collector.stop()

    assert warnings == [("test_start", "RuntimeError")]


def test_real_business_collector_keeps_messages_but_disables_modules() -> None:
    from obagent_sdk.observe.current import _COLLECTOR

    token = collector_module._start_business_collector(object())
    try:
        active = _COLLECTOR.get()
        assert isinstance(active, BusinessReadableFlowCollector)
        assert active.auto_module is False
    finally:
        collector_module._stop_business_collector(token)