AIRegulation-DocAnalysis/backend/tests/observability/test_model_usage_bootstrap.py

"""Unit tests for the model-usage persistence wiring in app.shared.bootstrap.

get_model_usage_store()'s settings-gating is tested the same way
tests/test_reranker_bootstrap.py tests get_reranker() — by patching
"app.shared.bootstrap.settings" wholesale, matching this codebase's
established convention for testing @lru_cache settings-gated factories.
The remaining tests isolate _start_model_usage_persistence() /
_stop_model_usage_persistence() from get_model_usage_store() entirely (via
monkeypatch on the module-level function), so no real database or event loop
is needed anywhere in this file — asyncio.create_task itself is also mocked.
"""

from __future__ import annotations

import sys
from unittest.mock import MagicMock, patch

# Patch psycopg2 before importing anything that transitively imports it, in
# case this file is collected before test_model_usage_persistence.py.
mock_psycopg2 = MagicMock()
mock_psycopg2.extras = MagicMock()
sys.modules.setdefault("psycopg2", mock_psycopg2)
sys.modules.setdefault("psycopg2.extras", mock_psycopg2.extras)
sys.modules.setdefault("psycopg2.pool", MagicMock())

from app.shared import bootstrap
from app.shared.model_usage_tracker import ModelUsageEntry, ModelUsageTracker


def test_get_model_usage_store_returns_none_when_not_postgres_backend():
    """get_model_usage_store() must be None unless document_repository_backend == 'postgres'."""
    bootstrap.get_model_usage_store.cache_clear()

    with patch("app.shared.bootstrap.settings") as mock_settings:
        mock_settings.document_repository_backend = "json"
        result = bootstrap.get_model_usage_store()

    bootstrap.get_model_usage_store.cache_clear()
    assert result is None


def test_get_model_usage_store_returns_instance_when_postgres_backend():
    """get_model_usage_store() must return a PostgresModelUsageStore when enabled.

    ThreadedConnectionPool is mocked so no real connection is attempted; the
    postgres_host/port/user/password/db values PostgresModelUsageStore reads
    come from app.config.settings.settings directly (not from the
    app.shared.bootstrap.settings reference mocked below), so they don't need
    to be set here — only document_repository_backend gates this factory.
    """
    bootstrap.get_model_usage_store.cache_clear()

    with patch("psycopg2.pool.ThreadedConnectionPool"), \
         patch(
             "app.infrastructure.storage.postgres_model_usage_store.PostgresModelUsageStore._ensure_schema"
         ), \
         patch("app.shared.bootstrap.settings") as mock_settings:
        mock_settings.document_repository_backend = "postgres"
        result = bootstrap.get_model_usage_store()

    bootstrap.get_model_usage_store.cache_clear()
    from app.infrastructure.storage.postgres_model_usage_store import PostgresModelUsageStore
    assert isinstance(result, PostgresModelUsageStore)


def test_start_model_usage_persistence_seeds_tracker_and_starts_flush_loop(monkeypatch):
    """When a store is available, startup must seed the tracker and schedule the flush task."""
    fake_store = MagicMock()
    fake_store.load_all.return_value = {
        "deepseek:deepseek-v4-flash": ModelUsageEntry(
            provider="deepseek", model="deepseek-v4-flash", total_tokens=99,
        ),
    }
    tracker = ModelUsageTracker()
    monkeypatch.setattr(bootstrap, "get_model_usage_store", lambda: fake_store)
    monkeypatch.setattr(bootstrap, "get_model_usage_tracker", lambda: tracker)

    with patch("asyncio.create_task") as mock_create_task:
        bootstrap._start_model_usage_persistence()
        # Close the coroutine object passed to the mock so pytest doesn't warn
        # about "coroutine was never awaited" — it was never meant to run here.
        mock_create_task.call_args[0][0].close()

    assert tracker.get("deepseek", "deepseek-v4-flash").total_tokens == 99
    mock_create_task.assert_called_once()

    bootstrap._stop_model_usage_persistence()  # reset the module-level task handle


def test_start_model_usage_persistence_is_a_no_op_without_a_store(monkeypatch):
    """No store configured (json backend) — startup must not touch asyncio or the tracker."""
    monkeypatch.setattr(bootstrap, "get_model_usage_store", lambda: None)

    with patch("asyncio.create_task") as mock_create_task:
        bootstrap._start_model_usage_persistence()

    mock_create_task.assert_not_called()


def test_stop_model_usage_persistence_cancels_task_and_flushes(monkeypatch):
    """Shutdown must cancel the running flush task and perform one final flush."""
    fake_store = MagicMock()
    monkeypatch.setattr(bootstrap, "get_model_usage_store", lambda: fake_store)
    fake_task = MagicMock()
    bootstrap._model_usage_flush_task = fake_task

    bootstrap._stop_model_usage_persistence()

    fake_task.cancel.assert_called_once()
    fake_store.flush.assert_called_once()
    assert bootstrap._model_usage_flush_task is None