diff --git a/litellm/proxy/spend_tracking/spend_capture_rate.py b/litellm/proxy/spend_tracking/spend_capture_rate.py index 4536ea0ee42..f78adeee538 100644 --- a/litellm/proxy/spend_tracking/spend_capture_rate.py +++ b/litellm/proxy/spend_tracking/spend_capture_rate.py @@ -86,7 +86,7 @@ async def captured_spend_by_day( end_date: date, ) -> Mapping[str, float]: """LiteLLM's tracked spend per UTC day (ISO date) for the given ``custom_llm_provider`` values.""" - rows: Final = await prisma_client.db.query_raw( + rows: Final = await prisma_client.replica_db.query_raw( _CAPTURED_SPEND_BY_DAY_SQL, start_date.isoformat(), end_date.isoformat(), tuple(litellm_providers) ) return MappingProxyType({row.date: row.spend for row in _CAPTURED_SPEND_ROWS.validate_python(rows)}) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py index 8fb8b41a0a8..eab5b263983 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py @@ -733,6 +733,7 @@ async def test_update_agent_in_db_preserves_secret_when_echoed_back_redacted(): with the redaction marker.""" registry: Final = AgentRegistry() mock_prisma: Final = MagicMock() + mock_prisma.replica_db = mock_prisma.db mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( return_value=SimpleNamespace( @@ -830,6 +831,7 @@ async def test_update_agent_in_db_preserves_secret_nested_under_a_non_sensitive_ redaction marker, not just top-level secret keys.""" registry: Final = AgentRegistry() mock_prisma: Final = MagicMock() + mock_prisma.replica_db = mock_prisma.db mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( return_value=SimpleNamespace( @@ -967,6 +969,7 @@ async def test_patch_agent_in_db_preserves_secret_when_echoed_back_redacted(): the stored credential.""" registry: Final = AgentRegistry() mock_prisma: Final = MagicMock() + mock_prisma.replica_db = mock_prisma.db mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( return_value={ @@ -1165,6 +1168,7 @@ _KILL_SWITCH: Final = { async def test_add_agent_to_db_stores_kill_switch_json_and_a_json_null_when_unset(): registry: Final = AgentRegistry() mock_prisma: Final = MagicMock() + mock_prisma.replica_db = mock_prisma.db mock_create = AsyncMock(return_value=_agent_row_mock([])) mock_prisma.db.litellm_agentstable.create = mock_create @@ -1191,6 +1195,7 @@ async def test_add_agent_to_db_stores_kill_switch_json_and_a_json_null_when_unse async def test_add_agent_to_db_rejects_a_kill_switch_with_a_non_http_url(): registry: Final = AgentRegistry() mock_prisma: Final = MagicMock() + mock_prisma.replica_db = mock_prisma.db mock_prisma.db.litellm_agentstable.create = AsyncMock(return_value=_agent_row_mock([])) with pytest.raises(Exception, match="absolute http"): @@ -1210,6 +1215,7 @@ async def test_add_agent_to_db_rejects_a_kill_switch_with_a_non_http_url(): async def test_patch_agent_in_db_keeps_kill_switch_when_omitted_and_clears_it_on_null(): registry: Final = AgentRegistry() mock_prisma: Final = MagicMock() + mock_prisma.replica_db = mock_prisma.db mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( return_value={ "agent_id": "agent-123", @@ -1239,6 +1245,7 @@ async def test_patch_agent_in_db_keeps_kill_switch_when_omitted_and_clears_it_on async def test_patch_agent_in_db_restores_the_stored_kill_switch_secret_behind_the_marker(): registry: Final = AgentRegistry() mock_prisma: Final = MagicMock() + mock_prisma.replica_db = mock_prisma.db mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( return_value={ "agent_id": "agent-123", @@ -1274,6 +1281,7 @@ async def test_patch_agent_in_db_restores_the_stored_kill_switch_secret_behind_t async def test_update_agent_in_db_clears_kill_switch_when_omitted_and_restores_secret_when_echoed(): registry: Final = AgentRegistry() mock_prisma: Final = MagicMock() + mock_prisma.replica_db = mock_prisma.db mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( return_value=SimpleNamespace(litellm_params={}, object_permission_id=None, kill_switch=json.dumps(_KILL_SWITCH)) ) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index 0cb422bddd8..74b83406bce 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -1330,6 +1330,7 @@ def test_kill_switch_trigger_fires_the_db_row_config_over_a_stale_in_memory_copy kill_switch={"url": "https://ops.example.com/kill-v2", "method": "DELETE", "auth": None}, ) prisma: Final = MagicMock() + prisma.replica_db = prisma.db prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=db_row) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) fake: Final = _FakeKillSwitchClient(httpx.Response(204)) @@ -1353,6 +1354,7 @@ def test_get_agent_redacts_kill_switch_secret_for_admins_and_hides_it_from_other def _get_as(role: LitellmUserRoles): with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_prisma.replica_db = mock_prisma.db mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) return _make_app_with_role(role).get("/v1/agents/agent-123", headers={"Authorization": "Bearer k"}) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_capture_rate.py b/tests/test_litellm/proxy/spend_tracking/test_spend_capture_rate.py index ebcdc95b8a2..d129633d0dd 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_capture_rate.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_capture_rate.py @@ -68,6 +68,7 @@ def _page(*buckets: dict[str, object], next_page: str | None = None) -> dict[str def _fake_prisma(rows: list[dict[str, object]]) -> MagicMock: prisma = MagicMock() + prisma.replica_db = prisma.db prisma.db.query_raw = AsyncMock(return_value=rows) return prisma @@ -306,6 +307,7 @@ async def test_a_healthy_scheduled_check_publishes_and_never_touches_the_alert_l async def test_a_scheduled_check_that_fails_never_claims_the_alert_window(monkeypatch): lock = _pod_lock(acquired=True) prisma = MagicMock() + prisma.replica_db = prisma.db prisma.db.query_raw = AsyncMock(side_effect=RuntimeError("database gone")) with pytest.raises(RuntimeError, match="database gone"): @@ -395,6 +397,7 @@ class _PsycopgPrisma: def __init__(self, conn: psycopg.Connection) -> None: self.db = self + self.replica_db = self self._conn = conn async def query_raw(self, sql: str, *params: object) -> list[dict[str, object]]: