diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index 7f5efcdaec3..170eb450339 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -421,7 +421,10 @@ class SpendLogCleanup: general_settings=None, redis_cache: RedisCache | None = None, partition_manager: SpendLogsPartitionManager | None = None, + *, + clock: Callable[[], float] = time.monotonic, ): + self._clock = clock self.retention_seconds: int | None = None self.partition_manager = partition_manager or SpendLogsPartitionManager() from litellm.proxy.proxy_server import general_settings as default_settings @@ -586,7 +589,7 @@ class SpendLogCleanup: async def _probe_storage(self, prisma_client: PrismaClient, deadline: float) -> StorageProbe | None: """Read what the spend-log tables weigh from the catalog, or None when the probe fails or has no budget.""" - if time.monotonic() >= deadline: + if self._clock() >= deadline: return None timeout_ms: Final = self._timeout_ms(deadline) try: @@ -624,7 +627,7 @@ class SpendLogCleanup: stale.allocated_bytes, _ANALYZE_SQL, ) - if time.monotonic() >= deadline: + if self._clock() >= deadline: return None timeout_ms: Final = self._timeout_ms(deadline) try: @@ -653,7 +656,7 @@ class SpendLogCleanup: async def _load_size_state(self, prisma_client: PrismaClient, deadline: float) -> SizeCleanupState | None: """The reclaimable count saved by the last check, or None on the first check or when it cannot be read.""" - if time.monotonic() >= deadline: + if self._clock() >= deadline: return None timeout_ms: Final = self._timeout_ms(deadline) try: @@ -686,7 +689,7 @@ class SpendLogCleanup: the tool-index rows that expire with their spend logs. When fewer rows than that remain, the step is everything up to now. """ - if time.monotonic() >= deadline: + if self._clock() >= deadline: return None timeout_ms: Final = self._timeout_ms(deadline) async with db_span("size_step_cutoff", "LiteLLM_SpendLogs"), bounded_tx(prisma_client, timeout_ms) as tx: @@ -718,11 +721,11 @@ class SpendLogCleanup: Interpolating this into SQL is safe by construction: an int cannot carry SQL, and SET does not accept a bind parameter. """ - remaining_ms: Final = int((deadline - time.monotonic()) * 1000) + remaining_ms: Final = int((deadline - self._clock()) * 1000) return max(1, min(int(self.batch_timeout_seconds * 1000), remaining_ms)) @staticmethod - def _group_deadline(overall_deadline: float, groups_remaining: int) -> float: + def _group_deadline(overall_deadline: float, groups_remaining: int, now: float) -> float: """ Give each pending cleanup group an equal share of the time left. @@ -731,10 +734,9 @@ class SpendLogCleanup: """ if groups_remaining == 1: return overall_deadline - current_time: Final = time.monotonic() - if current_time >= overall_deadline: + if now >= overall_deadline: return overall_deadline - return current_time + (overall_deadline - current_time) / groups_remaining + return now + (overall_deadline - now) / groups_remaining def _remaining_timeout_ms(self, deadline: float) -> RemainingTimeoutMs: """ @@ -748,7 +750,7 @@ class SpendLogCleanup: """ def remaining() -> int | None: - return None if time.monotonic() >= deadline else self._timeout_ms(deadline) + return None if self._clock() >= deadline else self._timeout_ms(deadline) return remaining @@ -841,7 +843,7 @@ class SpendLogCleanup: run_count = 0 consecutive_failures = 0 while True: - if time.monotonic() >= deadline: + if self._clock() >= deadline: verbose_proxy_logger.info( "Run budget exhausted during %s cleanup after %d rows; the next run resumes from here", table_name, @@ -858,13 +860,13 @@ class SpendLogCleanup: prisma_client, cutoff_date, table_name, time_column, total_deleted, "batch_cap_reached", deadline ) # Find rows and delete them in one go without fetching to application - batch_started_at = time.monotonic() + batch_started_at = self._clock() try: batch_result = await self._execute_delete_batch( prisma_client, delete_sql, cutoff_date, table_name, deadline ) except Exception as batch_exc: - if time.monotonic() >= deadline: + if self._clock() >= deadline: # The statement timeout was clamped to the budget that was # left, so this batch was cancelled by the deadline itself. # That is the bound working, not a database fault, and @@ -920,7 +922,7 @@ class SpendLogCleanup: consecutive_failures = 0 deleted_count = batch_result - SpendLogCleanupMetrics.record_batch(table_name, deleted_count, time.monotonic() - batch_started_at) + SpendLogCleanupMetrics.record_batch(table_name, deleted_count, self._clock() - batch_started_at) verbose_proxy_logger.info("Deleted %s %s rows in this batch", deleted_count, table_name) if deleted_count == 0: @@ -957,7 +959,7 @@ class SpendLogCleanup: operator the backlog was not drained; the gauge simply keeps its value from the last run that finished inside its budget. """ - if time.monotonic() >= deadline: + if self._clock() >= deadline: return TableCleanupResult(table_name=table_name, rows_deleted=rows_deleted, stop_reason=stop_reason) remaining: Final = await self._count_remaining(prisma_client, cutoff_date, table_name, time_column, deadline) if remaining is not None: @@ -1070,7 +1072,7 @@ class SpendLogCleanup: # behind any long-running reader for as long as that reader lives, which # is the one way this job could still outlast its budget without bound. remaining_timeout_ms: Final = self._remaining_timeout_ms(deadline) - if time.monotonic() >= deadline: + if self._clock() >= deadline: verbose_proxy_logger.info("Run budget already spent, skipping partition maintenance this run") elif self.general_settings.get( "use_spend_logs_partitioning", False @@ -1082,7 +1084,7 @@ class SpendLogCleanup: verbose_proxy_logger.info("Dropped %d expired spend-log partitions: %s", len(dropped), dropped) logs_result: Final = await self._delete_old_logs( - prisma_client, cutoff_date, self._group_deadline(deadline, groups_remaining=2) + prisma_client, cutoff_date, self._group_deadline(deadline, groups_remaining=2, now=self._clock()) ) verbose_proxy_logger.info("Deleted %s logs", logs_result.rows_deleted) @@ -1210,7 +1212,7 @@ class SpendLogCleanup: self, prisma_client: PrismaClient, deadline: float, still_over_target: Callable[[], Awaitable[bool]] ) -> tuple[str, ...]: remaining_timeout_ms: Final = self._remaining_timeout_ms(deadline) - if time.monotonic() >= deadline or not self._setting("use_spend_logs_partitioning"): + if self._clock() >= deadline or not self._setting("use_spend_logs_partitioning"): return () if not await self.partition_manager.is_partitioned(prisma_client, remaining_timeout_ms): return () @@ -1257,7 +1259,7 @@ class SpendLogCleanup: needed: Final = rows_still_needed(state.logs_deleted + state.index_deleted) if needed <= 0: return state.stopped("target_reached") - if time.monotonic() >= deadline: + if self._clock() >= deadline: return state.stopped("budget_exhausted") if state.batches_left <= 0: return state.stopped("batch_cap_reached") @@ -1318,11 +1320,11 @@ class SpendLogCleanup: except Exception: # noqa: BLE001 # retained observations are retried by the next cleanup job verbose_proxy_logger.warning("Auto-router baseline retention remains pending") sessions_result: Final = await self._delete_old_autorouter_session_rows( - prisma_client, session_cutoff, self._group_deadline(deadline, 3) + prisma_client, session_cutoff, self._group_deadline(deadline, 3, self._clock()) ) verbose_proxy_logger.info("Deleted %s expired auto-router session rollup rows", sessions_result.rows_deleted) user_sessions_result: Final = await self._delete_old_autorouter_user_session_rows( - prisma_client, session_cutoff, self._group_deadline(deadline, 2) + prisma_client, session_cutoff, self._group_deadline(deadline, 2, self._clock()) ) verbose_proxy_logger.info( "Deleted %s expired auto-router user session rollup rows", user_sessions_result.rows_deleted @@ -1443,7 +1445,7 @@ class SpendLogCleanup: async def _run_passes(self, prisma_client: PrismaClient, age_rules: bool, size_rule: bool) -> None: lock_acquired = False - run_started_at: Final = time.monotonic() + run_started_at: Final = self._clock() progress: Final = _RunProgress() progress_token: Final = _run_progress.set(progress) try: @@ -1493,7 +1495,7 @@ class SpendLogCleanup: SpendLogCleanupMetrics.record_run("skipped_locked") return - deadline: Final = time.monotonic() + self.run_budget_seconds + deadline: Final = self._clock() + self.run_budget_seconds size_group: Final = int(active_size_rule is not None) configured_group_count: Final = ( int(delete_spend_logs and self.retention_seconds is not None) @@ -1506,7 +1508,7 @@ class SpendLogCleanup: spend_log_results: Final = ( await self._clean_spend_log_tables( prisma_client, - self._group_deadline(deadline, configured_group_count), + self._group_deadline(deadline, configured_group_count, self._clock()), ) if delete_spend_logs and self.retention_seconds is not None else () @@ -1526,7 +1528,7 @@ class SpendLogCleanup: await self._clean_session_rollup( prisma_client, autorouter_retention_seconds, - self._group_deadline(deadline, remaining_groups_after_spend_logs), + self._group_deadline(deadline, remaining_groups_after_spend_logs, self._clock()), ) if autorouter_retention_seconds is not None else () @@ -1535,7 +1537,7 @@ class SpendLogCleanup: await self._clean_health_checks( prisma_client, health_check_retention_seconds, - self._group_deadline(deadline, remaining_groups_after_sessions), + self._group_deadline(deadline, remaining_groups_after_sessions, self._clock()), ) if health_check_retention_seconds is not None else () @@ -1544,7 +1546,7 @@ class SpendLogCleanup: await self._clean_daily_tag_spend( prisma_client, daily_tag_spend_retention_seconds, - self._group_deadline(deadline, 1 + size_group), + self._group_deadline(deadline, 1 + size_group, self._clock()), ) if daily_tag_spend_retention_seconds is not None else () @@ -1561,12 +1563,12 @@ class SpendLogCleanup: ) outcome: Final = self._run_outcome(results) SpendLogCleanupMetrics.record_run(outcome) - self._log_run_summary(outcome, results, time.monotonic() - run_started_at, size_report) + self._log_run_summary(outcome, results, self._clock() - run_started_at, size_report) except asyncio.CancelledError: verbose_proxy_logger.error( "Spend log cleanup cancelled after %.2fs (rows_deleted=%d, batches=%d); the next run resumes from here", - time.monotonic() - run_started_at, + self._clock() - run_started_at, progress.rows_deleted, progress.batches, ) diff --git a/tests/integration/coordination_redis_proxy_config.yaml b/tests/integration/coordination_redis_proxy_config.yaml index e28e0f5f2ad..49a7a24e0ac 100644 --- a/tests/integration/coordination_redis_proxy_config.yaml +++ b/tests/integration/coordination_redis_proxy_config.yaml @@ -9,5 +9,6 @@ general_settings: coordination_redis: host: os.environ/REDIS_HOST port: os.environ/REDIS_PORT + socket_timeout: os.environ/COORDINATION_REDIS_SOCKET_TIMEOUT router_settings: disable_cooldowns: true diff --git a/tests/integration/management/test_user_updates_wedged_coordination_redis.py b/tests/integration/management/test_user_updates_wedged_coordination_redis.py index 7934e927d93..2b05d35d553 100644 --- a/tests/integration/management/test_user_updates_wedged_coordination_redis.py +++ b/tests/integration/management/test_user_updates_wedged_coordination_redis.py @@ -16,6 +16,7 @@ from pydantic import JsonValue from redis import Redis from redis.client import PubSub +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import _BULK_EVICTION_WAIT_SECONDS from tests.integration._support.budget import case_timeout_seconds from tests.integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value from tests.integration._support.database import read_rows @@ -27,6 +28,8 @@ _BURST: Final = 30 _WEDGED_PUBLISHES_PER_BURST_USER: Final = 2 _HANDLER_BUDGET_SECONDS: Final = 0.75 _BULK_BUDGET_SECONDS: Final = 2.0 +_COORDINATION_SOCKET_TIMEOUT_SECONDS: Final = 5.0 +_SHORTEST_REDIS_WAIT_SECONDS: Final = min(_BULK_EVICTION_WAIT_SECONDS, _COORDINATION_SOCKET_TIMEOUT_SECONDS) _CHANNEL: Final = "litellm_proxy.auth_cache_invalidation" @@ -102,6 +105,7 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( "DATABASE_URL": database_url, "REDIS_HOST": coordination.host, "REDIS_PORT": str(coordination.port), + "COORDINATION_REDIS_SOCKET_TIMEOUT": str(_COORDINATION_SOCKET_TIMEOUT_SECONDS), }, config=Path("tests/integration/coordination_redis_proxy_config.yaml"), remove_environment=("DATABASE_URL_READ_REPLICA",), @@ -235,8 +239,9 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( elapsed for i, (status, elapsed) in enumerate(results) if i % 5 in (0, 1, 4) and status == 200 ) timings["c1_burst_p95"] = elapsed_sorted[int(len(elapsed_sorted) * 0.95) - 1] - assert timings["c1_burst_p95"] < _HANDLER_BUDGET_SECONDS, ( - f"burst p95 {timings['c1_burst_p95']:.3f}s" + assert timings["c1_burst_p95"] < _SHORTEST_REDIS_WAIT_SECONDS, ( + f"burst p95 {timings['c1_burst_p95']:.3f}s reached the {_SHORTEST_REDIS_WAIT_SECONDS}s " + "shortest wait the proxy puts on Redis, so burst updates waited on the wedged Redis" ) eventually( lambda: candidate.request("GET", "/health/liveliness").status_code, diff --git a/tests/unit/proxy/auth/test_login_utils.py b/tests/unit/proxy/auth/test_login_utils.py index 37ea05f8ba9..5148e38aa61 100644 --- a/tests/unit/proxy/auth/test_login_utils.py +++ b/tests/unit/proxy/auth/test_login_utils.py @@ -10,6 +10,7 @@ import os from collections.abc import Mapping from contextlib import ExitStack from datetime import datetime, timedelta, timezone +from types import SimpleNamespace from typing import TYPE_CHECKING, Final from unittest.mock import AsyncMock, MagicMock, patch @@ -678,6 +679,12 @@ def _throttle( ) +def _set_throttle_clock(monkeypatch: pytest.MonkeyPatch, now: float) -> None: + from litellm.proxy.auth import login_throttle + + monkeypatch.setattr(login_throttle, "time", SimpleNamespace(time=lambda: now)) + + def _stores(): from litellm.caching.in_memory_cache import InMemoryCache @@ -756,12 +763,14 @@ async def test_too_many_failures_for_one_username_block_that_pair_and_carry_retr monkeypatch.setenv("UI_USERNAME", "admin") monkeypatch.setenv("UI_PASSWORD", "right") + _set_throttle_clock(monkeypatch, 1_000.0) throttle = _throttle(user_limit=2, block_seconds=77) keys = throttle._keys("admin") assert [await _fail(throttle) for _ in range(3)] == ["401", "401", "401"], "the limit itself is a plain 401" assert throttle._local_block_ttl(keys.pair_block) == 77 + _set_throttle_clock(monkeypatch, 1_000.6) with pytest.raises(ProxyException) as blocked: await _guess(throttle) assert blocked.value.code == "429" @@ -814,10 +823,12 @@ async def test_a_correct_password_is_refused_while_its_pair_is_blocked(monkeypat monkeypatch.setenv("UI_USERNAME", "admin") monkeypatch.setenv("UI_PASSWORD", "right") monkeypatch.setenv("DATABASE_URL", "postgresql://stub") + _set_throttle_clock(monkeypatch, 1_000.0) throttle = _throttle(user_limit=1, block_seconds=90) assert [await _fail(throttle, username="user@corp.com") for _ in range(3)] == ["401", "401", "429"] + _set_throttle_clock(monkeypatch, 1_000.6) with pytest.raises(ProxyException) as refused: await _db_login(throttle, "user@corp.com", "right", correct=True) assert refused.value.code == "429" @@ -1328,12 +1339,14 @@ async def test_a_source_block_outranks_a_pair_block_in_the_retry_after(monkeypat monkeypatch.setenv("UI_USERNAME", "admin") monkeypatch.setenv("UI_PASSWORD", "right") + _set_throttle_clock(monkeypatch, 1_000.0) throttle = _throttle(user_limit=1, source_limit=3, block_seconds=120, client_ip="203.0.113.45") assert [await _fail(throttle) for _ in range(2)] == ["401", "401"], "the admin pair is now blocked" throttle.blocks.set_cache(throttle._keys("admin").pair_block, 1, ttl=30) assert [await _fail(throttle, username=f"spray-{i}@corp.com") for i in range(3)] == ["401"] * 3 assert throttle._local_block_ttl(throttle._keys("admin").source_block) == 120, "the source is now blocked too" + _set_throttle_clock(monkeypatch, 1_000.6) for name in ("admin", "spray-0@corp.com", "never-seen@corp.com"): with pytest.raises(ProxyException) as refused: await _guess(throttle, username=name) @@ -1457,6 +1470,7 @@ async def test_a_redis_outage_falls_back_to_this_workers_own_counter(monkeypatch monkeypatch.setenv("UI_USERNAME", "admin") monkeypatch.setenv("UI_PASSWORD", "right") + _set_throttle_clock(monkeypatch, 1_000.0) throttle = _throttle(user_limit=2, block_seconds=300, redis_cache=_DownRedis()) records: list[logging.LogRecord] = [] @@ -1465,6 +1479,7 @@ async def test_a_redis_outage_falls_back_to_this_workers_own_counter(monkeypatch verbose_proxy_logger.addHandler(handler) try: assert [await _fail(throttle) for _ in range(3)] == ["401"] * 3 + _set_throttle_clock(monkeypatch, 1_000.6) with pytest.raises(ProxyException) as blocked: await _guess(throttle) finally: diff --git a/tests/unit/proxy/proxy_server/test_routes_login_sso.py b/tests/unit/proxy/proxy_server/test_routes_login_sso.py index 43629b926ad..a6de5c7b6a0 100644 --- a/tests/unit/proxy/proxy_server/test_routes_login_sso.py +++ b/tests/unit/proxy/proxy_server/test_routes_login_sso.py @@ -11,6 +11,7 @@ Routes covered: from __future__ import annotations from datetime import datetime, timezone +from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -504,6 +505,12 @@ def test_login_form_ignores_open_redirect_return_to(client, monkeypatch): # --------------------------------------------------------------------------- +def _set_throttle_clock(monkeypatch: pytest.MonkeyPatch, now: float) -> None: + from litellm.proxy.auth import login_throttle + + monkeypatch.setattr(login_throttle, "time", SimpleNamespace(time=lambda: now)) + + def _install_real_auth(monkeypatch, **settings): """Run the real authenticate_user so the throttle inside it is exercised. @@ -583,11 +590,12 @@ def test_budget_is_shared_across_username_casing(client, monkeypatch, reset_logi def test_a_refused_attempt_carries_retry_after(client, monkeypatch, reset_login_throttle): - """The 429 tells the caller how long the block has left.""" + """The 429 tells the caller how long the block has left, rounded up to whole seconds.""" + _set_throttle_clock(monkeypatch, 1_000.0) _install_real_auth(monkeypatch, max_failed_login_attempts_per_source=2, failed_login_block_seconds=77) - assert [_json_login(client, "/v2/login") for _ in range(2)] == [401, 401] + _set_throttle_clock(monkeypatch, 1_000.6) refused = client.post("/v2/login", json={"username": "admin", "password": "wrong"}) assert refused.status_code == 429 assert refused.headers.get("retry-after") == "77" @@ -595,10 +603,11 @@ def test_a_refused_attempt_carries_retry_after(client, monkeypatch, reset_login_ def test_the_form_returns_a_human_readable_lockout_page(client, monkeypatch, reset_login_throttle): """The no-JavaScript form must render a wait page when its POST is throttled.""" + _set_throttle_clock(monkeypatch, 1_000.0) _install_real_auth(monkeypatch, max_failed_login_attempts_per_source=2, failed_login_block_seconds=77) - assert [_form_login(client) for _ in range(2)] == [401, 401] + _set_throttle_clock(monkeypatch, 1_000.6) refused = client.post("/login", data={"username": "admin", "password": "wrong"}) assert refused.status_code == 429 assert refused.headers.get("content-type", "").startswith("text/html") @@ -697,12 +706,7 @@ def test_the_master_key_as_a_bearer_token_still_works_while_the_ui_password_is_b def test_a_database_users_correct_password_is_refused_while_blocked(client, monkeypatch, reset_login_throttle): """The block is hard: while it lasts, nothing from that source signs in as that user, right password or not, and the block is not extended by the refused attempts.""" - from types import SimpleNamespace - - from litellm.proxy.auth import login_throttle - - frozen_now = login_throttle.time.time() - monkeypatch.setattr(login_throttle, "time", SimpleNamespace(time=lambda: frozen_now)) + _set_throttle_clock(monkeypatch, 1_000.0) _install_real_auth(monkeypatch, max_failed_login_attempts_per_source=2, failed_login_block_seconds=64) _db_user(monkeypatch, "user@corp.com") diff --git a/tests/unit/proxy/test_spend_log_cleanup.py b/tests/unit/proxy/test_spend_log_cleanup.py index 876ffe96584..114c970e163 100644 --- a/tests/unit/proxy/test_spend_log_cleanup.py +++ b/tests/unit/proxy/test_spend_log_cleanup.py @@ -8,6 +8,7 @@ import math import time from collections.abc import Awaitable, Callable from contextlib import asynccontextmanager +from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -34,6 +35,17 @@ from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import ( from tests.unit.proxy.db.fake_prisma_engine import engine_call +@dataclass(slots=True) +class _ManualClock: + now: float = 0.0 + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + def _far_deadline() -> float: """A run deadline far enough out that only the other bounds can stop a batch loop.""" return time.monotonic() + 3600 @@ -908,26 +920,31 @@ async def test_run_budget_stops_the_loop_and_leaves_the_backlog_for_the_next_run _wire_tx(mock_prisma_client.db) mock_db = MagicMock() _wire_tx(mock_db) - mock_db.execute_raw = AsyncMock(return_value=1000) + clock: Final = _ManualClock() + + async def _delete_batch(_query: str, *_args: object) -> int: + clock.advance(0.1) + return 1000 + + mock_db.execute_raw = AsyncMock(side_effect=_delete_batch) mock_prisma_client.db = mock_db cleaner = SpendLogCleanup( general_settings={ "maximum_spend_logs_retention_period": "7d", # Comfortably more batches than a sub-second budget can reach (each - # batch sleeps 0.1s), but small enough that a broken deadline fails - # this test in seconds instead of hanging it + # batch advances the clock 0.1s), but small enough that a broken + # deadline still ends the loop "maximum_spend_logs_cleanup_max_batches": 50, - } + }, + clock=clock, ) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) - started_at = time.monotonic() - result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, time.monotonic() + 1.0) - elapsed = time.monotonic() - started_at + result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, clock() + 1.0) assert result.stop_reason == "budget_exhausted" - assert elapsed < 3, f"budgeted run overran its deadline: {elapsed}s" + assert clock.now < 1.15, f"budgeted run overran its deadline: {clock.now}s" assert mock_db.execute_raw.call_count < 50 assert result.rows_deleted > 0 @@ -943,7 +960,13 @@ async def test_run_budget_is_shared_across_tables_not_granted_per_table(): _wire_tx(mock_prisma_client.db) mock_db = MagicMock() _wire_tx(mock_db) - mock_db.execute_raw = AsyncMock(return_value=1000) + clock: Final = _ManualClock() + + async def _delete_batch(_query: str, *_args: object) -> int: + clock.advance(0.1) + return 1000 + + mock_db.execute_raw = AsyncMock(side_effect=_delete_batch) mock_prisma_client.db = mock_db cleaner = SpendLogCleanup( @@ -951,20 +974,19 @@ async def test_run_budget_is_shared_across_tables_not_granted_per_table(): "maximum_spend_logs_retention_period": "7d", "maximum_autorouter_session_retention_period": "365d", # Comfortably more batches than a sub-second budget can reach (each - # batch sleeps 0.1s), but small enough that a broken deadline fails - # this test in seconds instead of hanging it + # batch advances the clock 0.1s), but small enough that a broken + # deadline still ends the loop "maximum_spend_logs_cleanup_max_batches": 50, "maximum_spend_logs_cleanup_run_budget": "1s", - } + }, + clock=clock, ) cleaner.pod_lock_manager = None - started_at = time.monotonic() await cleaner.cleanup_old_spend_logs(mock_prisma_client) - elapsed = time.monotonic() - started_at # three tables are eligible; a per-table budget would push this past 3s - assert elapsed < 2.5, f"budget was granted per table, not per run: {elapsed}s" + assert clock.now < 1.15, f"budget was granted per table, not per run: {clock.now}s" tables_touched = {call[0][0].split('"')[1] for call in mock_db.execute_raw.call_args_list} assert "LiteLLM_SpendLogs" in tables_touched @@ -972,10 +994,15 @@ async def test_run_budget_is_shared_across_tables_not_granted_per_table(): @pytest.mark.asyncio async def test_cleanup_groups_share_budget_so_health_checks_still_get_a_delete(): mock_prisma_client = MagicMock() - _wire_tx(mock_prisma_client.db) mock_db = MagicMock() _wire_tx(mock_db) - mock_db.execute_raw = AsyncMock(return_value=1000) + clock: Final = _ManualClock() + + async def _delete_batch(_query: str, *_args: object) -> int: + clock.advance(0.1) + return 1000 + + mock_db.execute_raw = AsyncMock(side_effect=_delete_batch) mock_prisma_client.db = mock_db cleaner = SpendLogCleanup( @@ -984,7 +1011,8 @@ async def test_cleanup_groups_share_budget_so_health_checks_still_get_a_delete() "maximum_health_check_retention_period": "30d", "maximum_spend_logs_cleanup_max_batches": 50, "maximum_spend_logs_cleanup_run_budget": "1s", - } + }, + clock=clock, ) cleaner.pod_lock_manager = None @@ -1598,7 +1626,13 @@ async def test_spend_logs_backlog_cannot_starve_tool_index_cleanup(): mock_prisma_client = MagicMock() mock_db = MagicMock() _wire_tx(mock_db) - mock_db.execute_raw = AsyncMock(return_value=1000) + clock: Final = _ManualClock() + + async def _delete_batch(_query: str, *_args: object) -> int: + clock.advance(0.1) + return 1000 + + mock_db.execute_raw = AsyncMock(side_effect=_delete_batch) mock_prisma_client.db = mock_db cleaner = SpendLogCleanup( @@ -1606,18 +1640,17 @@ async def test_spend_logs_backlog_cannot_starve_tool_index_cleanup(): "maximum_spend_logs_retention_period": "7d", "maximum_spend_logs_cleanup_max_batches": 500, "maximum_spend_logs_cleanup_run_budget": "1s", - } + }, + clock=clock, ) cleaner.pod_lock_manager = None - started_at = time.monotonic() await cleaner.cleanup_old_spend_logs(mock_prisma_client) - elapsed = time.monotonic() - started_at tables = [call[0][0].split('"')[1] for call in mock_db.execute_raw.call_args_list] assert tables.count("LiteLLM_SpendLogs") > 0 assert tables.count("LiteLLM_SpendLogToolIndex") > 0, "tool index cleanup was starved by the spend-log backlog" - assert elapsed < 2.5, f"splitting the budget must not extend the run: {elapsed}s" + assert clock.now < 1.15, f"splitting the budget must not extend the run: {clock.now}s" @pytest.mark.asyncio