diff --git a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py index 7defd6cef6c..87adc464608 100644 --- a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py +++ b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py @@ -7,11 +7,11 @@ from datetime import datetime, timedelta import pytest from apscheduler.schedulers.asyncio import AsyncIOScheduler -import litellm.proxy.shutdown.scheduled_jobs as scheduled_jobs +from litellm.proxy.shutdown import scheduled_jobs from litellm.proxy.shutdown.scheduled_jobs import ( AwaitableAsyncIOExecutor, - stop_in_flight_scheduler_jobs, pause_scheduled_jobs, + stop_in_flight_scheduler_jobs, ) diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index 1691b2d174a..b8f28d0c780 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -83,10 +83,10 @@ def test_spend_log_cleanup_cron_scheduling(): assert trigger_weekly is not None # Invalid cron expression should raise ValueError - with pytest.raises(ValueError, match='Wrong number of fields; got'): + with pytest.raises(ValueError, match="Wrong number of fields; got"): CronTrigger.from_crontab("invalid cron") - with pytest.raises(ValueError, match='is higher than the maximum value'): + with pytest.raises(ValueError, match="is higher than the maximum value"): CronTrigger.from_crontab("60 25 * * *") # Invalid minute and hour @@ -99,6 +99,7 @@ def test_spend_log_cleanup_cron_scheduler_integration(): a real database connection. """ from unittest.mock import MagicMock + from apscheduler.triggers.cron import CronTrigger # Mock scheduler @@ -145,15 +146,11 @@ def test_spend_log_cleanup_cron_scheduler_integration(): # No cron, so it should fall back to interval } - cleanup_cron_fallback = general_settings_interval.get( - "maximum_spend_logs_cleanup_cron" - ) + cleanup_cron_fallback = general_settings_interval.get("maximum_spend_logs_cleanup_cron") assert cleanup_cron_fallback is None # No cron configured # Simulate interval-based scheduling fallback - retention_interval = general_settings_interval.get( - "maximum_spend_logs_retention_interval", "1d" - ) + retention_interval = general_settings_interval.get("maximum_spend_logs_retention_interval", "1d") from litellm.litellm_core_utils.duration_parser import duration_in_seconds interval_seconds = duration_in_seconds(retention_interval) @@ -181,27 +178,19 @@ async def test_should_delete_spend_logs(): assert cleaner._should_delete_spend_logs() is False # Test case 2: Valid seconds string - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "3600s"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "3600s"}) assert cleaner._should_delete_spend_logs() is True # Test case 3: Valid days string - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "30d"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "30d"}) assert cleaner._should_delete_spend_logs() is True # Test case 4: Valid hours string - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "24h"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "24h"}) assert cleaner._should_delete_spend_logs() is True # Test case 5: Invalid format - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "invalid"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "invalid"}) assert cleaner._should_delete_spend_logs() is False @@ -288,9 +277,7 @@ async def test_cleanup_old_spend_logs_retention_period_cutoff(): # Verify the cutoff date is correct cutoff_date = mock_db.execute_raw.call_args[0][1] expected_cutoff = datetime.now(timezone.utc) - timedelta(seconds=86400) - assert ( - abs((cutoff_date - expected_cutoff).total_seconds()) < 1 - ) # Allow 1 second difference for test execution time + assert abs((cutoff_date - expected_cutoff).total_seconds()) < 1 # Allow 1 second difference for test execution time @pytest.mark.asyncio @@ -310,9 +297,7 @@ async def test_cleanup_drops_partitions_when_enabled_and_partitioned(): partition_manager = MagicMock() partition_manager.is_partitioned = AsyncMock(return_value=True) partition_manager.ensure_partitions = AsyncMock(return_value=["p1"]) - partition_manager.drop_partitions_older_than = AsyncMock( - return_value=["LiteLLM_SpendLogs_p20260601"] - ) + partition_manager.drop_partitions_older_than = AsyncMock(return_value=["LiteLLM_SpendLogs_p20260601"]) cleaner = SpendLogCleanup( general_settings={ @@ -450,9 +435,7 @@ async def test_integer_retention_treated_as_days(): An integer value for maximum_spend_logs_retention_period should be treated as days (e.g., 3 → '3d' → 259200 seconds). """ - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": 3} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": 3}) result = cleaner._should_delete_spend_logs() assert result is True assert cleaner.retention_seconds == 3 * 86400 # 3 days in seconds @@ -469,13 +452,11 @@ def test_string_retention_still_works(): ("2w", 2 * 604800), ] for setting, expected_seconds in cases: - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": setting} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": setting}) assert cleaner._should_delete_spend_logs() is True, f"Failed for {setting}" - assert ( - cleaner.retention_seconds == expected_seconds - ), f"Expected {expected_seconds} for {setting}, got {cleaner.retention_seconds}" + assert cleaner.retention_seconds == expected_seconds, ( + f"Expected {expected_seconds} for {setting}, got {cleaner.retention_seconds}" + ) @pytest.mark.asyncio @@ -489,9 +470,7 @@ async def test_delete_old_logs_aborts_on_non_int_execute_raw_return(): mock_db.execute_raw = AsyncMock(return_value=None) mock_prisma_client.db = mock_db - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -510,9 +489,7 @@ async def test_delete_old_logs_continues_on_valid_int_return(): mock_db.execute_raw = AsyncMock(side_effect=[500, 300, 0]) mock_prisma_client.db = mock_db - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -559,9 +536,7 @@ async def test_delete_old_tool_index_rows_deletes_on_composite_key(): mock_db.execute_raw = AsyncMock(side_effect=[5, 0]) mock_prisma_client.db = mock_db - cleaner = SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_tool_index_rows(mock_prisma_client, cutoff_date, _far_deadline()) @@ -581,9 +556,7 @@ async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch) import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module # Zero out the failure backoff so the test doesn't take ~0.5s of real sleep. - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 - ) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) @@ -591,14 +564,10 @@ async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch) _wire_tx(mock_db) # batch 1 succeeds, batch 2 raises (one-off DB timeout), batches 3-4 succeed, # batch 5 returns 0 → loop exits naturally. - mock_db.execute_raw = AsyncMock( - side_effect=[100, TimeoutError("simulated DB timeout"), 200, 50, 0] - ) + mock_db.execute_raw = AsyncMock(side_effect=[100, TimeoutError("simulated DB timeout"), 200, 50, 0]) mock_prisma_client.db = mock_db - cleaner = cleanup_module.SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -615,26 +584,18 @@ async def test_delete_old_logs_aborts_after_consecutive_failures(monkeypatch): import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module # Lower the threshold so the test is fast and deterministic. - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3 - ) - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 - ) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) mock_db = MagicMock() _wire_tx(mock_db) # Every batch raises — must abort after exactly 3 attempts, not loop forever. - mock_db.execute_raw = AsyncMock( - side_effect=ConnectionError("simulated persistent DB outage") - ) + mock_db.execute_raw = AsyncMock(side_effect=ConnectionError("simulated persistent DB outage")) mock_prisma_client.db = mock_db - cleaner = cleanup_module.SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -649,12 +610,8 @@ async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatc intermittent timeouts don't trip the abort threshold.""" import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3 - ) - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 - ) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) @@ -675,9 +632,7 @@ async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatc ) mock_prisma_client.db = mock_db - cleaner = cleanup_module.SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cutoff_date = datetime.now(timezone.utc) - timedelta(days=7) result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline()) @@ -698,9 +653,7 @@ async def test_cleanup_uses_logger_exception_for_full_traceback(monkeypatch): mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) # Force the outer try/except to fire by making _should_delete_spend_logs raise. - cleaner = cleanup_module.SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cleaner.pod_lock_manager = None def boom(): @@ -725,12 +678,8 @@ async def test_cleanup_releases_lock_after_persistent_batch_failures(monkeypatch must still be released so the next scheduled run isn't permanently blocked.""" import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 2 - ) - monkeypatch.setattr( - cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 - ) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 2) + monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0) mock_prisma_client = MagicMock() _wire_tx(mock_prisma_client.db) @@ -744,9 +693,7 @@ async def test_cleanup_releases_lock_after_persistent_batch_failures(monkeypatch mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) mock_pod_lock_manager.release_lock = AsyncMock() - cleaner = cleanup_module.SpendLogCleanup( - general_settings={"maximum_spend_logs_retention_period": "7d"} - ) + cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) cleaner.pod_lock_manager = mock_pod_lock_manager await cleaner.cleanup_old_spend_logs(mock_prisma_client) @@ -996,9 +943,7 @@ async def test_each_batch_carries_a_statement_and_lock_timeout(): } ) - await cleaner._delete_old_logs( - mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline() - ) + await cleaner._delete_old_logs(mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline()) assert "SET LOCAL statement_timeout = 12000" in recorded assert "SET LOCAL lock_timeout = 12000" in recorded @@ -1134,9 +1079,7 @@ async def test_remaining_rows_probe_is_capped_so_it_cannot_scan_the_table(): cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) - await cleaner._delete_old_logs( - mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline() - ) + await cleaner._delete_old_logs(mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline()) count_sql = mock_db.query_raw.call_args[0][0] assert "count(*)" in count_sql