"""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