e079ee779f
- proxy_server.py: disable allow_credentials when allow_origins=['*'] (wildcard + credentials is a browser security misconfiguration). Add LITELLM_CORS_ORIGINS env var to configure explicit allowed origins. - create_views.py: narrow broad 'except Exception' to only catch genuine 'view does not exist' errors; re-raise all other DB errors (auth, connection, etc.) that were previously silently swallowed. - spend_log_cleanup.py: validate execute_raw() return type is int before using it as a deletion count; break loop safely on unexpected types to prevent infinite deletion loops.
353 lines
12 KiB
Python
353 lines
12 KiB
Python
"""
|
|
Test cases for spend log cleanup functionality
|
|
"""
|
|
|
|
from datetime import datetime, timedelta, timezone
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup
|
|
|
|
|
|
def test_spend_log_cleanup_cron_scheduling():
|
|
"""Test that cron expressions are correctly parsed for spend log cleanup scheduling"""
|
|
from apscheduler.triggers.cron import CronTrigger
|
|
|
|
# Valid cron expressions
|
|
cron_expr = "0 4 * * *" # 4:00 AM daily
|
|
trigger = CronTrigger.from_crontab(cron_expr)
|
|
assert trigger is not None
|
|
|
|
# Every minute (useful for testing)
|
|
trigger_minute = CronTrigger.from_crontab("*/1 * * * *")
|
|
assert trigger_minute is not None
|
|
|
|
# Specific day and hour
|
|
trigger_weekly = CronTrigger.from_crontab("0 3 * * 0") # 3 AM every Sunday
|
|
assert trigger_weekly is not None
|
|
|
|
# Invalid cron expression should raise ValueError
|
|
with pytest.raises(ValueError):
|
|
CronTrigger.from_crontab("invalid cron")
|
|
|
|
with pytest.raises(ValueError):
|
|
CronTrigger.from_crontab("60 25 * * *") # Invalid minute and hour
|
|
|
|
|
|
def test_spend_log_cleanup_cron_scheduler_integration():
|
|
"""
|
|
Integration test: Verify the proxy_server scheduler logic correctly adds
|
|
cron-based cleanup job when maximum_spend_logs_cleanup_cron is configured.
|
|
|
|
This tests the logic in proxy_server.py lines 4671-4717 without requiring
|
|
a real database connection.
|
|
"""
|
|
from unittest.mock import MagicMock
|
|
from apscheduler.triggers.cron import CronTrigger
|
|
|
|
# Mock scheduler
|
|
mock_scheduler = MagicMock()
|
|
mock_prisma_client = MagicMock()
|
|
mock_cleanup_instance = MagicMock()
|
|
|
|
# Test Case 1: Cron-based scheduling
|
|
general_settings_cron = {
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
"maximum_spend_logs_cleanup_cron": "0 4 * * *", # 4 AM daily
|
|
}
|
|
|
|
cleanup_cron = general_settings_cron.get("maximum_spend_logs_cleanup_cron")
|
|
assert cleanup_cron is not None
|
|
|
|
# Simulate the scheduler logic from proxy_server.py
|
|
cron_trigger = CronTrigger.from_crontab(cleanup_cron)
|
|
mock_scheduler.add_job(
|
|
mock_cleanup_instance.cleanup_old_spend_logs,
|
|
cron_trigger,
|
|
args=[mock_prisma_client],
|
|
id="spend_log_cleanup_job",
|
|
replace_existing=True,
|
|
misfire_grace_time=3600,
|
|
)
|
|
|
|
# Verify scheduler was called correctly
|
|
mock_scheduler.add_job.assert_called_once()
|
|
call_args = mock_scheduler.add_job.call_args
|
|
|
|
# Verify the trigger is a CronTrigger
|
|
assert isinstance(call_args[0][1], CronTrigger)
|
|
|
|
# Verify job ID
|
|
assert call_args[1]["id"] == "spend_log_cleanup_job"
|
|
assert call_args[1]["replace_existing"] is True
|
|
|
|
# Test Case 2: Interval-based scheduling (fallback)
|
|
mock_scheduler.reset_mock()
|
|
general_settings_interval = {
|
|
"maximum_spend_logs_retention_period": "7d",
|
|
# No cron, so it should fall back to interval
|
|
}
|
|
|
|
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"
|
|
)
|
|
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
|
|
|
interval_seconds = duration_in_seconds(retention_interval)
|
|
|
|
mock_scheduler.add_job(
|
|
mock_cleanup_instance.cleanup_old_spend_logs,
|
|
"interval",
|
|
seconds=interval_seconds,
|
|
args=[mock_prisma_client],
|
|
id="spend_log_cleanup_job",
|
|
replace_existing=True,
|
|
)
|
|
|
|
# Verify interval scheduling was called
|
|
mock_scheduler.add_job.assert_called_once()
|
|
interval_call_args = mock_scheduler.add_job.call_args
|
|
assert interval_call_args[0][1] == "interval"
|
|
assert interval_call_args[1]["seconds"] == 86400 # 1 day in seconds
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_delete_spend_logs():
|
|
# Test case 1: No retention set
|
|
cleaner = SpendLogCleanup(general_settings={})
|
|
assert cleaner._should_delete_spend_logs() is False
|
|
|
|
# Test case 2: Valid seconds string
|
|
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"}
|
|
)
|
|
assert cleaner._should_delete_spend_logs() is True
|
|
|
|
# Test case 4: Valid hours string
|
|
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"}
|
|
)
|
|
assert cleaner._should_delete_spend_logs() is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_old_spend_logs_batch_deletion():
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
# Setup Prisma client
|
|
mock_prisma_client = MagicMock()
|
|
mock_db = MagicMock()
|
|
|
|
# Mock execute_raw to return deleted counts
|
|
mock_db.execute_raw = AsyncMock(side_effect=[1000, 500, 0])
|
|
|
|
# Wire up mocks
|
|
mock_prisma_client.db = mock_db
|
|
|
|
# Mock Redis cache and pod_lock_manager
|
|
mock_redis_cache = MagicMock()
|
|
mock_pod_lock_manager = MagicMock()
|
|
mock_pod_lock_manager.redis_cache = mock_redis_cache
|
|
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
|
|
mock_pod_lock_manager.release_lock = AsyncMock()
|
|
|
|
# Run cleanup with mocked pod_lock_manager
|
|
test_settings = {"maximum_spend_logs_retention_period": "7d"}
|
|
cleaner = SpendLogCleanup(general_settings=test_settings)
|
|
cleaner.pod_lock_manager = mock_pod_lock_manager
|
|
assert cleaner._should_delete_spend_logs() is True
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
# Validate batching and deletion via raw SQL
|
|
assert mock_db.execute_raw.call_count == 3
|
|
|
|
# Check the first call argument
|
|
call_args_sql = mock_db.execute_raw.call_args_list[0][0][0]
|
|
assert 'DELETE FROM "LiteLLM_SpendLogs"' in call_args_sql
|
|
assert 'WHERE "request_id" IN' in call_args_sql
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_old_spend_logs_retention_period_cutoff():
|
|
"""
|
|
Test that logs are filtered using correct cutoff based on retention
|
|
"""
|
|
# Setup Prisma client
|
|
mock_prisma_client = MagicMock()
|
|
mock_db = MagicMock()
|
|
mock_db.execute_raw = AsyncMock(return_value=0)
|
|
mock_prisma_client.db = mock_db
|
|
|
|
# Mock Redis cache and pod_lock_manager
|
|
mock_redis_cache = MagicMock()
|
|
mock_pod_lock_manager = MagicMock()
|
|
mock_pod_lock_manager.redis_cache = mock_redis_cache
|
|
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
|
|
mock_pod_lock_manager.release_lock = AsyncMock()
|
|
|
|
# Run cleanup with mocked pod_lock_manager
|
|
test_settings = {"maximum_spend_logs_retention_period": "24h"}
|
|
cleaner = SpendLogCleanup(general_settings=test_settings)
|
|
cleaner.pod_lock_manager = mock_pod_lock_manager
|
|
assert cleaner._should_delete_spend_logs() is True
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
# 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
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_old_spend_logs_no_retention_period():
|
|
"""
|
|
Test that no logs are deleted when no retention period is set
|
|
"""
|
|
mock_prisma_client = MagicMock()
|
|
mock_prisma_client.db.execute_raw = AsyncMock()
|
|
|
|
cleaner = SpendLogCleanup(general_settings={}) # no retention
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
mock_prisma_client.db.execute_raw.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lock_not_released_when_not_acquired():
|
|
"""
|
|
Lock release should be skipped when _should_delete_spend_logs returns False
|
|
before the lock is ever acquired.
|
|
"""
|
|
mock_prisma_client = MagicMock()
|
|
mock_prisma_client.db.execute_raw = AsyncMock()
|
|
|
|
mock_redis_cache = MagicMock()
|
|
mock_pod_lock_manager = MagicMock()
|
|
mock_pod_lock_manager.redis_cache = mock_redis_cache
|
|
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
|
|
mock_pod_lock_manager.release_lock = AsyncMock()
|
|
|
|
# No retention setting → _should_delete_spend_logs() returns False before lock is acquired
|
|
cleaner = SpendLogCleanup(general_settings={})
|
|
cleaner.pod_lock_manager = mock_pod_lock_manager
|
|
|
|
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
|
|
|
|
mock_pod_lock_manager.acquire_lock.assert_not_called()
|
|
mock_pod_lock_manager.release_lock.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
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}
|
|
)
|
|
result = cleaner._should_delete_spend_logs()
|
|
assert result is True
|
|
assert cleaner.retention_seconds == 3 * 86400 # 3 days in seconds
|
|
|
|
|
|
def test_string_retention_still_works():
|
|
"""
|
|
String values like '3d', '24h', '3600s' should continue to parse correctly.
|
|
"""
|
|
cases = [
|
|
("3d", 3 * 86400),
|
|
("24h", 24 * 3600),
|
|
("3600s", 3600),
|
|
("2w", 2 * 604800),
|
|
]
|
|
for setting, expected_seconds in cases:
|
|
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}"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_old_logs_aborts_on_non_int_execute_raw_return():
|
|
"""should abort deletion loop immediately when execute_raw returns a non-int
|
|
(e.g. None or dict), preventing an infinite loop."""
|
|
mock_prisma_client = MagicMock()
|
|
mock_db = MagicMock()
|
|
mock_db.execute_raw = AsyncMock(return_value=None)
|
|
mock_prisma_client.db = mock_db
|
|
|
|
cleaner = SpendLogCleanup(
|
|
general_settings={"maximum_spend_logs_retention_period": "7d"}
|
|
)
|
|
|
|
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
|
|
total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date)
|
|
|
|
assert mock_db.execute_raw.call_count == 1
|
|
assert total_deleted == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_old_logs_continues_on_valid_int_return():
|
|
"""should continue deletion loop across batches when execute_raw returns valid int counts."""
|
|
mock_prisma_client = MagicMock()
|
|
mock_db = MagicMock()
|
|
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"}
|
|
)
|
|
|
|
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
|
|
total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date)
|
|
|
|
assert mock_db.execute_raw.call_count == 3
|
|
assert total_deleted == 800
|
|
|
|
|
|
def test_cleanup_batch_size_env_var(monkeypatch):
|
|
"""Ensure batch size is configurable via environment variable"""
|
|
import importlib
|
|
|
|
import litellm.constants as constants_module
|
|
import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
|
|
|
|
# Set env var and reload modules to pick up new value
|
|
monkeypatch.setenv("SPEND_LOG_CLEANUP_BATCH_SIZE", "25")
|
|
importlib.reload(constants_module)
|
|
importlib.reload(cleanup_module)
|
|
|
|
cleaner = cleanup_module.SpendLogCleanup(general_settings={})
|
|
assert cleaner.batch_size == 25
|
|
|
|
# Remove env var and reload to restore default for other tests
|
|
monkeypatch.delenv("SPEND_LOG_CLEANUP_BATCH_SIZE", raising=False)
|
|
importlib.reload(constants_module)
|
|
importlib.reload(cleanup_module)
|