Spaces:
Configuration error
Configuration error
File size: 5,989 Bytes
447ebeb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 |
"""
Test cases for spend log cleanup functionality
"""
from datetime import UTC, datetime, timedelta, timezone
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup
@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 types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
# Setup Prisma client
mock_prisma_client = MagicMock()
mock_db = MagicMock()
# Mock spendlogs table
mock_spendlogs = MagicMock()
mock_spendlogs.find_many = AsyncMock()
mock_spendlogs.delete_many = AsyncMock()
# Create 1500 mocked logs with .request_id
mock_logs = [SimpleNamespace(request_id=f"req_{i}") for i in range(1500)]
mock_spendlogs.find_many.side_effect = [
mock_logs[:1000], # Batch 1
mock_logs[1000:], # Batch 2
[], # Done
]
# Wire up mocks
mock_db.litellm_spendlogs = mock_spendlogs
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
assert mock_spendlogs.find_many.call_count == 3
assert mock_spendlogs.delete_many.call_count == 2
mock_spendlogs.delete_many.assert_any_call(
where={"request_id": {"in": [f"req_{i}" for i in range(1000)]}}
)
mock_spendlogs.delete_many.assert_any_call(
where={"request_id": {"in": [f"req_{i}" for i in range(1000, 1500)]}}
)
@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_spendlogs = MagicMock()
mock_spendlogs.find_many = AsyncMock(return_value=[])
mock_spendlogs.delete_many = AsyncMock()
mock_db.litellm_spendlogs = mock_spendlogs
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_spendlogs.find_many.call_args[1]["where"]["startTime"]["lt"]
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.litellm_spendlogs.find_many = AsyncMock()
mock_prisma_client.db.litellm_spendlogs.delete = AsyncMock()
cleaner = SpendLogCleanup(general_settings={}) # no retention
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
mock_prisma_client.db.litellm_spendlogs.find_many.assert_not_called()
mock_prisma_client.db.litellm_spendlogs.delete.assert_not_called()
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)
|