SupplyAgent/tests/supply_agent/test_publish_hook.py

"""Publish hook stays free of supply_infra at import time."""

from __future__ import annotations

import importlib
from pathlib import Path

from supply_agent.logging.publish import (
    get_run_artifact_publisher,
    publish_run_artifacts,
    set_run_artifact_publisher,
)


def test_publish_module_does_not_import_supply_infra() -> None:
    source = importlib.util.find_spec("supply_agent.logging.publish")
    assert source is not None and source.origin is not None
    text = Path(source.origin).read_text(encoding="utf-8")
    assert "import supply_infra" not in text
    assert "from supply_infra" not in text


def test_publish_is_noop_without_publisher() -> None:
    previous = get_run_artifact_publisher()
    try:
        set_run_artifact_publisher(None)

        class _Logger:
            pass

        assert publish_run_artifacts(_Logger()) is None  # type: ignore[arg-type]
    finally:
        set_run_artifact_publisher(previous)


def test_set_run_artifact_publisher_is_invoked() -> None:
    previous = get_run_artifact_publisher()
    seen: list[object] = []

    def _publisher(logger: object) -> str:
        seen.append(logger)
        return "https://example.com/log.html"

    try:
        set_run_artifact_publisher(_publisher)
        marker = object()
        assert publish_run_artifacts(marker) == "https://example.com/log.html"  # type: ignore[arg-type]
        assert seen == [marker]
    finally:
        set_run_artifact_publisher(previous)


def test_infra_registers_publisher_hook() -> None:
    previous = get_run_artifact_publisher()
    from supply_agent.logging.usage import get_llm_usage_recorder, set_llm_usage_recorder

    previous_usage = get_llm_usage_recorder()
    try:
        set_run_artifact_publisher(None)
        set_llm_usage_recorder(None)
        import supply_infra  # noqa: F401
        from supply_infra.agent_logging.register import register_agent_logging_hooks
        from supply_infra.agent_logging.publish import publish_run_artifacts_to_oss
        from supply_infra.llm_billing.recorder import record_llm_usage_to_db

        register_agent_logging_hooks()
        assert get_run_artifact_publisher() is publish_run_artifacts_to_oss
        assert get_llm_usage_recorder() is record_llm_usage_to_db
    finally:
        set_run_artifact_publisher(previous)
        set_llm_usage_recorder(previous_usage)