diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index ece85b09432..e277b45e19a 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -327,6 +327,7 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset( "LiteLLM_PromptTable", "LiteLLM_HealthCheckTable", "LiteLLM_SearchToolsTable", + "LiteLLM_ScheduledJobRun", "LiteLLM_SSOConfig", "LiteLLM_ManagedVectorStoreIndexTable", "LiteLLM_CacheConfig", diff --git a/litellm/proxy/db/scheduled_job_runs.py b/litellm/proxy/db/scheduled_job_runs.py index 710f230d7e5..41915b9f3a8 100644 --- a/litellm/proxy/db/scheduled_job_runs.py +++ b/litellm/proxy/db/scheduled_job_runs.py @@ -22,8 +22,6 @@ from litellm.proxy.db.db_span import db_span from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import bounded_tx, trusted_sql if TYPE_CHECKING: - from prisma import Prisma - from litellm.proxy.utils import PrismaClient _TABLE: Final = "LiteLLM_ScheduledJobRun" @@ -246,7 +244,7 @@ _STATE_SQL: Final = f""" (SELECT max("finished_at") FROM "{_TABLE}" WHERE "job_name" = $1 AND "status" IN ('SUCCEEDED', 'FAILED')) AS "last_finished_at", (SELECT max("scheduled_for") FROM "{_TABLE}" - WHERE "job_name" = $1 AND "status" <> 'CANCELLED') AS "last_scheduled_for" + WHERE "job_name" = $1 AND "status" NOT IN ('CANCELLED', 'PENDING')) AS "last_scheduled_for" """ _PENDING_SQL: Final = f""" SELECT "id", "schedule", "scheduled_for" FROM "{_TABLE}" @@ -334,7 +332,10 @@ class ScheduledJobRunner: async with db_span("scheduled_job_tick", _TABLE), bounded_tx(prisma_client, timeout_ms) as tx: await tx.execute_raw(trusted_sql(f"SET LOCAL statement_timeout = {timeout_ms}")) await tx.execute_raw(trusted_sql(f"SET LOCAL lock_timeout = {timeout_ms}")) - if not await _try_group_lock(tx, self.job.lock_group): + locked: Final = _LOCKED.validate_python( + await tx.query_raw(trusted_sql(_GROUP_LOCK_SQL), f"litellm_scheduled_job:{self.job.lock_group}") + ) + if not locked or not locked[0].locked: return None for reaped in _CLAIMED.validate_python(await tx.query_raw(trusted_sql(_REAP_SQL), self.job.lock_group)): verbose_proxy_logger.warning( @@ -419,16 +420,6 @@ class ScheduledJobRunner: verbose_proxy_logger.info("Scheduled job %s run %s %s", self.job.name, run_id, status) -async def _try_group_lock(tx: "Prisma", lock_group: str) -> bool: - rows: Final = _LOCKED.validate_python( - await tx.query_raw( - trusted_sql('SELECT pg_try_advisory_xact_lock(hashtext($1)) AS "locked"'), - f"litellm_scheduled_job:{lock_group}", - ) - ) - return bool(rows) and rows[0].locked - - class _Locked(BaseModel): model_config = ConfigDict(frozen=True) @@ -436,6 +427,7 @@ class _Locked(BaseModel): _LOCKED: Final = TypeAdapter(list[_Locked]) +_GROUP_LOCK_SQL: Final = 'SELECT pg_try_advisory_xact_lock(hashtext($1)) AS "locked"' @lru_cache(maxsize=64) diff --git a/tests/integration/database/test_scheduled_job_runs.py b/tests/integration/database/test_scheduled_job_runs.py index ed5b6ff4538..896969c694a 100644 --- a/tests/integration/database/test_scheduled_job_runs.py +++ b/tests/integration/database/test_scheduled_job_runs.py @@ -265,6 +265,22 @@ async def test_a_schedule_change_replaces_the_pending_run_planned_from_the_old_s assert job.started == 0 +async def test_moving_a_cron_earlier_plans_the_earlier_occurrence_instead_of_skipping_a_day() -> None: + async with _scratch() as database: + now: Final = datetime.now(timezone.utc) + job: Final = _Job(schedule=f"0 {(now + timedelta(hours=2)).hour} * * *") + runners: Final = _runners(job) + await _tick_all(runners, database) + (old,) = database.rows("\"status\" = 'PENDING'") + + job.schedule_value = f"0 {(now + timedelta(hours=1)).hour} * * *" + await _tick_all(runners, database) + + (replacement,) = database.rows("\"status\" = 'PENDING'") + assert replacement["scheduled_for"] < old["scheduled_for"] + assert replacement["scheduled_for"] - datetime.now(timezone.utc) < timedelta(hours=1) + + async def test_a_run_cancelled_by_shutdown_is_recorded_as_failed_instead_of_holding_its_lease() -> None: async with _scratch() as database: job: Final = _Job() diff --git a/tests/unit/proxy/db/test_scheduled_job_runs.py b/tests/unit/proxy/db/test_scheduled_job_runs.py index 439c01f7ad4..47306dcd3ac 100644 --- a/tests/unit/proxy/db/test_scheduled_job_runs.py +++ b/tests/unit/proxy/db/test_scheduled_job_runs.py @@ -1,21 +1,36 @@ +import asyncio +import logging +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from datetime import datetime, timedelta, timezone -from typing import Final +from typing import TYPE_CHECKING, Final, cast import pytest +from apscheduler.schedulers.asyncio import AsyncIOScheduler from litellm.proxy.db.scheduled_job_runs import ( InvalidSchedule, + JobFailed, + JobResult, + JobScheduler, JobState, + JobSucceeded, NewRun, PendingRun, Schedule, + ScheduledJob, + ScheduledJobRunner, TickPlan, cron_schedule, interval_schedule, next_run_at, plan_tick, + register_poll, ) +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + NOW: Final = datetime(2026, 11, 4, 12, 0, tzinfo=timezone.utc) @@ -167,3 +182,183 @@ def test_cron_schedule_rejects_what_is_not_a_crontab(value: object) -> None: @pytest.mark.parametrize("value", ["soon", "99999999999d", "0s", 86400, None]) def test_interval_schedule_rejects_what_is_not_a_positive_duration(value: object) -> None: assert isinstance(interval_schedule("maximum_spend_logs_retention_interval", value), InvalidSchedule) + + +class _ScriptedDatabase: + """Answers the runner's statements the way Postgres would for one due run, and records every write.""" + + def __init__(self, *, locked: bool = True, finish_updates: int = 1, broken: bool = False) -> None: + self.locked = locked + self.finish_updates = finish_updates + self.broken = broken + self.writes: tuple[tuple[str, tuple[object, ...]], ...] = () + self.db = self + + @asynccontextmanager + async def tx(self, timeout: timedelta) -> AsyncIterator["_ScriptedDatabase"]: + if self.broken: + raise ConnectionError("database unreachable") + yield self + + async def query_raw(self, sql: str, *args: object) -> list[dict[str, object]]: + if "pg_try_advisory_xact_lock" in sql: + return [{"locked": self.locked}] + if "SELECT now()" in sql: + return [ + { + "now": NOW, + "group_running": False, + "job_running": False, + "last_finished_at": None, + "last_scheduled_for": None, + } + ] + if "'RUNNING', \"claimed_by\"" in sql: + self.writes = (*self.writes, ("claim", args)) + return [{"id": args[0], "job_name": "job"}] + return [] + + async def execute_raw(self, sql: str, *args: object) -> int: + if sql.lstrip().startswith("SET LOCAL"): + return 0 + if 'SET "status" = $3' in sql: + self.writes = (*self.writes, ("finish", args)) + self.locked = False + return self.finish_updates + if sql.lstrip().startswith("INSERT"): + self.writes = (*self.writes, ("insert", args)) + return 1 + + def client(self) -> "PrismaClient": + return cast("PrismaClient", self) # cast-ok: the runner only reaches the client through .db + + +def _finishes(database: _ScriptedDatabase) -> tuple[tuple[object, ...], ...]: + return tuple(args for kind, args in database.writes if kind == "finish") + + +class _Job: + def __init__(self, result: JobResult | BaseException = JobSucceeded("cleaned")) -> None: + self.result = result + self.runs = 0 + self.release = asyncio.Event() + self.release.set() + + async def run(self, _prisma_client: object) -> JobResult: + self.runs += 1 + await self.release.wait() + if isinstance(self.result, BaseException): + raise self.result + return self.result + + def runner(self) -> ScheduledJobRunner: + definition: Final = ScheduledJob( + name="job", + lock_group="group", + schedule=lambda: _schedule("1h"), + lease=lambda: timedelta(minutes=5), + run=self.run, + ) + return ScheduledJobRunner(definition, worker_id="pod-a", new_id=lambda: "run-1") + + +async def test_a_due_run_is_claimed_by_this_pod_run_and_recorded_with_its_log() -> None: + database: Final = _ScriptedDatabase() + job: Final = _Job(JobSucceeded("cleaned")) + await job.runner().tick(database.client()) + + assert job.runs == 1 + assert ("claim", ("run-1", "pod-a", 300.0)) in database.writes + assert _finishes(database) == (("run-1", "pod-a", "SUCCEEDED", "cleaned"),) + + +@pytest.mark.parametrize( + ("result", "expected_log"), + [ + pytest.param(JobFailed("aborted after 3 batch failures"), "aborted after 3 batch failures", id="reported"), + pytest.param(RuntimeError("disk on fire"), "RuntimeError: disk on fire", id="raised"), + ], +) +async def test_a_failed_run_is_recorded_as_failed_with_why( + result: JobResult | BaseException, expected_log: str +) -> None: + database: Final = _ScriptedDatabase() + await _Job(result).runner().tick(database.client()) + assert _finishes(database) == (("run-1", "pod-a", "FAILED", expected_log),) + + +async def test_a_pod_that_does_not_get_the_group_lock_neither_plans_nor_runs() -> None: + database: Final = _ScriptedDatabase(locked=False) + job: Final = _Job() + await job.runner().tick(database.client()) + assert job.runs == 0 + assert database.writes == () + + +async def test_an_unreachable_database_skips_the_tick_instead_of_raising_into_the_scheduler() -> None: + job: Final = _Job() + await job.runner().tick(_ScriptedDatabase(broken=True).client()) + assert job.runs == 0 + + +async def test_a_poll_landing_while_this_pod_still_runs_the_job_does_not_touch_the_database() -> None: + database: Final = _ScriptedDatabase() + job: Final = _Job() + job.release.clear() + runner: Final = job.runner() + first: Final = asyncio.create_task(runner.tick(database.client())) + while job.runs == 0: + await asyncio.sleep(0) + writes_during_run: Final = database.writes + + await runner.tick(database.client()) + assert database.writes == writes_during_run + + job.release.set() + await first + assert job.runs == 1 + + +async def test_a_run_cancelled_mid_flight_records_failed_and_still_cancels() -> None: + database: Final = _ScriptedDatabase() + job: Final = _Job() + job.release.clear() + tick: Final = asyncio.create_task(job.runner().tick(database.client())) + while job.runs == 0: + await asyncio.sleep(0) + tick.cancel() + with pytest.raises(asyncio.CancelledError): + await tick + ((run_id, worker, status, log),) = _finishes(database) + assert (run_id, worker, status) == ("run-1", "pod-a", "FAILED") + assert "cancelled" in str(log) + + +async def test_a_result_that_arrives_after_the_lease_expired_is_reported_not_recorded( + caplog: pytest.LogCaptureFixture, +) -> None: + database: Final = _ScriptedDatabase(finish_updates=0) + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await _Job().runner().tick(database.client()) + assert any("after its lease expired" in record.getMessage() for record in caplog.records) + + +async def test_register_poll_adds_a_wanted_poll_once_and_removes_it_when_no_longer_wanted() -> None: + scheduler: Final = AsyncIOScheduler() + scheduler.start(paused=True) + job_scheduler: Final = cast(JobScheduler, scheduler) # cast-ok: apscheduler ships no type information + definition: Final = _Job().runner().job + + def register(wanted: bool) -> None: + register_poll(job_scheduler, None, "poll-id", definition, wanted=wanted, poll_seconds=30, misfire_grace_time=60) + + try: + register(True) + first: Final = scheduler.get_job("poll-id") + assert first is not None + register(True) + assert scheduler.get_job("poll-id") is first, "re-registering replaced the live poll and its runner" + register(False) + assert scheduler.get_job("poll-id") is None + finally: + scheduler.shutdown(wait=False) diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 9dd15649d3e..be6b3c0108e 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -6,10 +6,7 @@ import json import logging import os import re -import socket -import subprocess import sys -import time import types from collections.abc import Callable, Iterator, Mapping, Sequence from datetime import datetime, timedelta, timezone @@ -19,9 +16,6 @@ from typing import Final from unittest import mock from unittest.mock import AsyncMock, MagicMock, call, create_autospec, mock_open, patch -from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs - -import click import fastapi.routing import httpx import pytest @@ -34,7 +28,6 @@ from starlette.datastructures import URL import litellm import litellm.proxy.proxy_server as proxy_server_module -from litellm._internal_context import current_service_target from litellm._logging import turn_on_json, verbose_proxy_logger from litellm.caching.caching import RedisCache from litellm.caching.dual_cache import DualCache @@ -51,12 +44,13 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.login_throttle import LoginThrottle from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.banner import LITELLM_BANNER from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendUpdateQueue from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash from litellm.proxy.proxy_server import app, aws_error_type_exception_handler, initialize, openai_exception_handler from litellm.utils import _invalidate_model_cost_lowercase_map -from litellm.proxy.common_utils.banner import LITELLM_BANNER from tests._master_key import MASTER_KEY +from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs PUBLICLY_KNOWN_KEY: Final = "sk-" + "1234" @@ -1002,43 +996,6 @@ async def test_initialize_scheduled_jobs_registers_cleanup_when_retention_lives_ scheduler.shutdown(wait=False) -@pytest.mark.asyncio -async def test_initialize_scheduled_jobs_does_not_fall_back_to_the_interval_for_a_non_string_cron(monkeypatch): - """A truthy non-string cron is invalid, so startup must log it and register no cleanup job - rather than silently pruning on the default interval the admin never configured.""" - monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False) - monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) - from apscheduler.schedulers.asyncio import AsyncIOScheduler - - from litellm.proxy.proxy_server import ProxyStartupEvent - from litellm.proxy.utils import ProxyLogging - - mock_prisma_client = MagicMock() - mock_proxy_logging = MagicMock(spec=ProxyLogging) - mock_proxy_logging.slack_alerting_instance = MagicMock() - mock_proxy_logging.db_spend_update_writer = MagicMock() - settings = {"maximum_daily_tag_spend_retention_period": "30d", "maximum_spend_logs_cleanup_cron": 5} - scheduler = AsyncIOScheduler() - try: - with ( - patch("litellm.proxy.proxy_server.proxy_config", _mock_scheduled_proxy_config()), - patch("litellm.proxy.proxy_server.store_model_in_db", False), - patch("litellm.proxy.proxy_server.general_settings", settings), - patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler), - ): - await ProxyStartupEvent.initialize_scheduled_background_jobs( - general_settings=settings, - prisma_client=mock_prisma_client, - proxy_budget_rescheduler_min_time=1, - proxy_budget_rescheduler_max_time=2, - proxy_batch_write_at=5, - proxy_logging_obj=mock_proxy_logging, - ) - assert scheduler.get_job("spend_log_cleanup_job") is None, "invalid cron fell back to the interval" - finally: - scheduler.shutdown(wait=False) - - @pytest.mark.asyncio async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval(monkeypatch): """ @@ -1701,7 +1658,7 @@ async def test_aaaproxy_startup_master_key(mock_prisma, monkeypatch, tmp_path): with open(config_path, "w") as f: yaml.dump(test_config, f) - print(f"SET ENV VARIABLE - CONFIG_FILE_PATH, str(config_path): {str(config_path)}") + print(f"SET ENV VARIABLE - CONFIG_FILE_PATH, str(config_path): {config_path!s}") # Second setting of CONFIG_FILE_PATH to a different value monkeypatch.setenv("CONFIG_FILE_PATH", str(config_path)) print(f"config_path: {config_path}") @@ -1720,7 +1677,7 @@ async def test_aaaproxy_startup_master_key(mock_prisma, monkeypatch, tmp_path): yaml.dump(empty_config, f) monkeypatch.setenv("LITELLM_MASTER_KEY", test_env_master_key) - print("test_env_master_key: {}".format(test_env_master_key)) + print(f"test_env_master_key: {test_env_master_key}") async with proxy_startup_event(app): from litellm.proxy.proxy_server import master_key @@ -2043,7 +2000,7 @@ def test_team_info_masking(): all_teams_config=[team1_info], ) - print("Got exception: {}".format(exc_info.value)) + print(f"Got exception: {exc_info.value}") assert "secret-test-key" not in str(exc_info.value) assert "public-test-key" not in str(exc_info.value) @@ -2055,7 +2012,6 @@ async def test_get_all_team_models(): """ from unittest.mock import AsyncMock, MagicMock - from litellm.proxy._types import LiteLLM_TeamTable from litellm.proxy.proxy_server import get_all_team_models # Mock team data @@ -4037,7 +3993,7 @@ async def test_load_config_warns_per_worker_login_counters_without_general_setti Redis must hear that its counters are per worker even when the config has no general_settings.""" import logging - import litellm.proxy.proxy_server as proxy_server + from litellm.proxy import proxy_server from litellm.proxy.auth.login_throttle import warn_login_counters_are_per_worker from litellm.proxy.proxy_server import ProxyConfig @@ -4197,7 +4153,6 @@ async def test_load_environment_variables_blocks_dangerous_keys(): Test that _load_environment_variables rejects dangerous env var keys like PATH, LD_PRELOAD, PYTHONPATH, etc. """ - import logging from litellm.proxy.proxy_server import ProxyConfig @@ -4279,7 +4234,7 @@ async def test_write_config_to_file(monkeypatch): """ Do not write config to file if store_model_in_db is True """ - from unittest.mock import AsyncMock, MagicMock, patch + from unittest.mock import AsyncMock, patch from litellm.proxy.proxy_server import ProxyConfig @@ -4327,7 +4282,7 @@ async def test_write_config_to_file_when_store_model_in_db_false(monkeypatch): """ Test that config IS written to file when store_model_in_db is False """ - from unittest.mock import AsyncMock, MagicMock, patch + from unittest.mock import patch from litellm.proxy.proxy_server import ProxyConfig @@ -4424,7 +4379,7 @@ async def test_async_data_generator_midstream_error(): try: async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) - except Exception as e: + except Exception: # If there's an exception, that's also part of what we want to test pass @@ -4493,7 +4448,6 @@ async def test_chat_completion_result_no_nested_none_values(): from unittest.mock import AsyncMock, MagicMock, patch from fastapi import Request, Response - from pydantic import BaseModel import litellm from litellm.proxy._types import UserAPIKeyAuth @@ -5560,7 +5514,7 @@ async def test_add_router_settings_from_db_config_merge_logic(): This tests how router settings from config file and database are combined, including scenarios where nested dictionaries should be properly merged. """ - from unittest.mock import AsyncMock, MagicMock, patch + from unittest.mock import AsyncMock, MagicMock from litellm.proxy.proxy_server import ProxyConfig @@ -6713,7 +6667,7 @@ async def test_init_sso_settings_in_db_error_handling(): """ Test that _init_sso_settings_in_db handles database errors gracefully. """ - from unittest.mock import AsyncMock, MagicMock, patch + from unittest.mock import AsyncMock, MagicMock from litellm.proxy.proxy_server import ProxyConfig @@ -6843,7 +6797,6 @@ async def test_init_hashicorp_vault_config_override_retries_on_transport_error() invocations.append(None) if len(invocations) == 1: raise prisma.errors.ClientNotConnectedError() - return None # No config in DB → function returns early after retry. mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_configoverrides.find_unique = AsyncMock(side_effect=_flaky_find_unique) @@ -6991,9 +6944,7 @@ def test_root_redirect_when_docs_url_not_root_and_redirect_url_set(monkeypatch): routes_to_remove = [] for route in app.routes: if hasattr(route, "path") and route.path == "/": - if hasattr(route, "methods") and "GET" in route.methods: - routes_to_remove.append(route) - elif not hasattr(route, "methods"): # Catch-all routes + if hasattr(route, "methods") and "GET" in route.methods or not hasattr(route, "methods"): routes_to_remove.append(route) for route in routes_to_remove: @@ -10041,7 +9992,7 @@ async def test_reseed_locks_dict_is_bounded(): long-lived deployments with high counter-key churn. Inserting more than the cap evicts the oldest entries. """ - import litellm.constants as constants + from litellm import constants from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed orig_locks = SpendCounterReseed._locks.copy() diff --git a/tests/unit/proxy/test_spend_log_cleanup.py b/tests/unit/proxy/test_spend_log_cleanup.py index 5905594ba7a..c37d9702d56 100644 --- a/tests/unit/proxy/test_spend_log_cleanup.py +++ b/tests/unit/proxy/test_spend_log_cleanup.py @@ -22,6 +22,8 @@ from litellm.constants import ( ) from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import ( SPEND_LOG_CLEANUP_BOUND_SETTINGS, + SPEND_LOG_RETENTION_JOB, + SPEND_LOG_SIZE_CHECK_JOB, SizeCleanupState, SpendLogCleanup, SpendLogStorage, @@ -33,7 +35,7 @@ from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import ( from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import ( SpendLogCleanupMetrics, ) -from litellm.proxy.db.scheduled_job_runs import InvalidSchedule +from litellm.proxy.db.scheduled_job_runs import InvalidSchedule, JobFailed, JobSucceeded from tests.unit.proxy.db.fake_prisma_engine import engine_call @@ -2178,3 +2180,51 @@ def test_the_retention_job_follows_the_cleanup_cron_when_set_and_the_retention_i schedule = retention_schedule(settings) assert not isinstance(schedule, InvalidSchedule) assert schedule.spec == expected_spec + + +@pytest.mark.parametrize("cron", [5, ["0 4 * * *"], "not a cron"]) +def test_an_invalid_cleanup_cron_is_rejected_instead_of_falling_back_to_the_retention_interval(cron: object): + schedule = retention_schedule( + {"maximum_spend_logs_cleanup_cron": cron, "maximum_spend_logs_retention_interval": "1h"} + ) + assert isinstance(schedule, InvalidSchedule) + + +async def test_the_retention_job_records_an_aborted_cleanup_as_failed(monkeypatch): + import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module + + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 1) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"maximum_spend_logs_retention_period": "7d"}) + mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) + mock_prisma_client.db.execute_raw = AsyncMock(side_effect=TimeoutError("DB down")) + + result = await SPEND_LOG_RETENTION_JOB.run(mock_prisma_client) + + assert isinstance(result, JobFailed) + assert "stop_reason=aborted" in result.log + + +@pytest.mark.parametrize("job", [SPEND_LOG_RETENTION_JOB, SPEND_LOG_SIZE_CHECK_JOB], ids=["retention", "size"]) +async def test_a_cleanup_job_with_no_rule_configured_succeeds_without_touching_the_database(monkeypatch, job): + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + mock_prisma_client = MagicMock() + mock_prisma_client.db.execute_raw = AsyncMock() + + result = await job.run(mock_prisma_client) + + assert result == JobSucceeded("no retention or size rule is configured") + mock_prisma_client.db.execute_raw.assert_not_awaited() + + +def test_a_cleanup_lease_outlasts_the_configured_run_budget(monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"maximum_spend_logs_cleanup_run_budget": "10m"}) + for job in (SPEND_LOG_RETENTION_JOB, SPEND_LOG_SIZE_CHECK_JOB): + assert job.lease().total_seconds() > 600 + + +def test_the_size_check_job_follows_the_size_check_interval(monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"maximum_spend_logs_size_check_interval": "6h"}) + schedule = SPEND_LOG_SIZE_CHECK_JOB.schedule() + assert not isinstance(schedule, InvalidSchedule) + assert schedule.spec == "interval:21600s"