mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
style(proxy): format cleanup shutdown tests
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng
parent
37ce0e47b6
commit
39a14f39e5
2 changed files with 37 additions and 94 deletions
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue