scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-31 06:56:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 244

test: fix flaky spend-log cleanup budget, login lockout, and wedged Redis tests (#47108)

Browse source 
* test: stabilize cleanup, lockout, and post-kill timing tests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): bound the wedged Redis burst p95 by the socket timeout the test sets

The burst p95 under a fixed 0.75s cap tracked how fast the box ran the 30
concurrent writes, so a loaded runner failed it at unchanged code. A handler
that waits on the wedged publish takes at least the coordination Redis socket
timeout, so the test now sets that timeout and asserts the p95 stays under it.

* test: pin the lockout clock and bound the wedged Redis burst by the proxy's 2s wait cap

The lockout tests now set the login throttle's clock instead of reading the real one, so Retry-After is pinned at 77 and a round-down fails on every run

The wedged Redis burst p95 is now bounded by the 2s cap the proxy puts on its own Redis waits, so a burst wait capped under the 5s socket timeout still fails the test

The fresh post-kill clients keep the harness's 2s keep-alive expiry, here and in the two tests #47080 changed

* test: tell ceil from round in the lockout tests and bound the burst by the imported eviction cap

The refused attempt now runs with 76.4s left on the 77s block, where ceil gives 77 but round and int give 76, so
both lockout tests fail if Retry-After stops rounding up.

The burst bound imports _BULK_EVICTION_WAIT_SECONDS instead of copying its 2s value, so lowering the cap moves
the bound with it.

* test: bound the wedged Redis burst by the shorter of the eviction cap and the socket timeout

Any wait the proxy puts on the stopped Redis lasts at least the shorter of the two, so the bound now follows
whichever one a later change lowers.

* test: drive the tool index cleanup test with the manual clock and keep the post-kill block as on main

The post-kill reset is already fixed on main by making kill() wait for the
worker's last thread, so the fresh-client and keep-alive edits go back to
main's version. The tool index cleanup test now advances the injected clock
instead of timing the real one, and the lockout clock helper gets a type hint.

* test: advance the spend log cleanup tests' manual clock through its own method

* test: freeze the throttle clock in the login_utils lockout tests and drive the run budget tests with the manual clock

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-29 03:46:32 -07:00 • committed by GitHub
parent b93d26bc85
commit e85cb4052e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 123 additions and 63 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

60
litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py
View file

@ -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,
)

1
tests/integration/coordination_redis_proxy_config.yaml
View file

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

9
tests/integration/management/test_user_updates_wedged_coordination_redis.py
View file

@ -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,

15
tests/unit/proxy/auth/test_login_utils.py
View file

@ -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:

22
tests/unit/proxy/proxy_server/test_routes_login_sso.py
View file

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

79
tests/unit/proxy/test_spend_log_cleanup.py
View file

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