mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
test(proxy): add comprehensive tests for enduser budget reset cache invalidation
- test_reset_budget_for_enduser_cache_failure_does_not_block_reset: verifies spend resets even when Redis is down - test_reset_budget_for_enduser_with_none_user_id: handles edge case of null user_id without crash - test_enduser_budget_reset_reproduces_stale_cache_bug: full reproduction of #24675 using real Pydantic models and DualCache - test_multiple_endusers_budget_reset_invalidates_all_caches: verifies batch reset invalidates all cache entries
This commit is contained in:
parent
efb25f2932
commit
2ac8a2ee85
1 changed files with 206 additions and 0 deletions
|
|
@ -299,6 +299,212 @@ def test_reset_budget_for_enduser_invalidates_cache():
|
|||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_reset_budget_for_enduser_cache_failure_does_not_block_reset():
|
||||
"""
|
||||
Test that if cache invalidation fails (e.g. Redis down), the spend reset
|
||||
still succeeds. The cache failure should be logged as a warning, not
|
||||
propagated as an exception.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
test_enduser = type(
|
||||
"LiteLLM_EndUserTable",
|
||||
(),
|
||||
{
|
||||
"spend": 30.0,
|
||||
"user_id": "enduser-cache-fail",
|
||||
},
|
||||
)
|
||||
|
||||
mock_cache = AsyncMock()
|
||||
mock_cache.async_delete_cache = AsyncMock(
|
||||
side_effect=ConnectionError("Redis connection refused")
|
||||
)
|
||||
|
||||
async def _run():
|
||||
with patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"litellm.proxy.proxy_server": type(
|
||||
"module", (), {"user_api_key_cache": mock_cache}
|
||||
)()
|
||||
},
|
||||
):
|
||||
result = await ResetBudgetJob._reset_budget_for_enduser(
|
||||
enduser=test_enduser
|
||||
)
|
||||
|
||||
# Spend should still be reset even though cache invalidation failed
|
||||
assert result is not None
|
||||
assert result.spend == 0.0
|
||||
# Cache delete was attempted
|
||||
mock_cache.async_delete_cache.assert_called_once()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_reset_budget_for_enduser_with_none_user_id():
|
||||
"""
|
||||
Test that _reset_budget_for_enduser handles end users with user_id=None
|
||||
gracefully - spend is reset but cache invalidation is skipped.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
test_enduser = type(
|
||||
"LiteLLM_EndUserTable",
|
||||
(),
|
||||
{
|
||||
"spend": 15.0,
|
||||
"user_id": None,
|
||||
},
|
||||
)
|
||||
|
||||
mock_cache = AsyncMock()
|
||||
|
||||
async def _run():
|
||||
with patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"litellm.proxy.proxy_server": type(
|
||||
"module", (), {"user_api_key_cache": mock_cache}
|
||||
)()
|
||||
},
|
||||
):
|
||||
result = await ResetBudgetJob._reset_budget_for_enduser(
|
||||
enduser=test_enduser
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.spend == 0.0
|
||||
# Cache delete should NOT be called when user_id is None
|
||||
mock_cache.async_delete_cache.assert_not_called()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_enduser_budget_reset_reproduces_stale_cache_bug():
|
||||
"""
|
||||
Reproduce the exact bug from issue #24675:
|
||||
1. End user has spend=45 cached, budget=50
|
||||
2. Cron job resets spend=0 in DB
|
||||
3. Without fix: auth check reads stale cache (spend=45), user appears near limit
|
||||
4. With fix: cache is invalidated, next read gets fresh spend=0
|
||||
|
||||
This test simulates the full flow using real LiteLLM_EndUserTable and
|
||||
LiteLLM_BudgetTable Pydantic objects.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_EndUserTable
|
||||
|
||||
# Step 1: Create an end user near their budget limit
|
||||
budget = LiteLLM_BudgetTable(
|
||||
budget_id="monthly-50",
|
||||
max_budget=50.0,
|
||||
budget_duration="1mo",
|
||||
)
|
||||
end_user = LiteLLM_EndUserTable(
|
||||
user_id="customer-123",
|
||||
blocked=False,
|
||||
spend=45.0,
|
||||
litellm_budget_table=budget,
|
||||
)
|
||||
|
||||
# Simulate the cache holding stale data
|
||||
from litellm.caching.caching import DualCache
|
||||
|
||||
cache = DualCache()
|
||||
|
||||
async def _run():
|
||||
# Put the end user in cache (simulating what get_end_user_object does)
|
||||
await cache.async_set_cache(
|
||||
key="end_user_id:customer-123",
|
||||
value=end_user.dict(),
|
||||
)
|
||||
|
||||
# Verify stale data is in cache
|
||||
cached = await cache.async_get_cache(key="end_user_id:customer-123")
|
||||
assert cached is not None
|
||||
assert cached["spend"] == 45.0 # stale!
|
||||
|
||||
# Step 2: Cron job runs _reset_budget_for_enduser WITH our fix
|
||||
with patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"litellm.proxy.proxy_server": type(
|
||||
"module", (), {"user_api_key_cache": cache}
|
||||
)()
|
||||
},
|
||||
):
|
||||
result = await ResetBudgetJob._reset_budget_for_enduser(enduser=end_user)
|
||||
|
||||
# Step 3: Verify spend is reset
|
||||
assert result.spend == 0.0
|
||||
|
||||
# Step 4: Verify cache was invalidated (the key fix)
|
||||
cached_after_reset = await cache.async_get_cache(key="end_user_id:customer-123")
|
||||
assert cached_after_reset is None, (
|
||||
"Cache should be invalidated after budget reset, but stale entry "
|
||||
f"still exists with spend={cached_after_reset.get('spend') if cached_after_reset else 'N/A'}"
|
||||
)
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_multiple_endusers_budget_reset_invalidates_all_caches():
|
||||
"""
|
||||
When multiple end users are linked to the same budget and the budget
|
||||
resets, ALL of their cache entries should be invalidated.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
mock_cache = AsyncMock()
|
||||
mock_cache.async_delete_cache = AsyncMock()
|
||||
deleted_keys = []
|
||||
|
||||
async def track_delete(key):
|
||||
deleted_keys.append(key)
|
||||
|
||||
mock_cache.async_delete_cache = AsyncMock(side_effect=track_delete)
|
||||
|
||||
endusers = []
|
||||
for i in range(5):
|
||||
endusers.append(
|
||||
type(
|
||||
"LiteLLM_EndUserTable",
|
||||
(),
|
||||
{
|
||||
"spend": 10.0 * (i + 1),
|
||||
"user_id": f"user-{i}",
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
async def _run():
|
||||
with patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"litellm.proxy.proxy_server": type(
|
||||
"module", (), {"user_api_key_cache": mock_cache}
|
||||
)()
|
||||
},
|
||||
):
|
||||
for enduser in endusers:
|
||||
await ResetBudgetJob._reset_budget_for_enduser(enduser=enduser)
|
||||
|
||||
# All 5 users should have spend reset
|
||||
for enduser in endusers:
|
||||
assert enduser.spend == 0.0
|
||||
|
||||
# All 5 cache entries should be invalidated
|
||||
assert len(deleted_keys) == 5
|
||||
expected_keys = [f"end_user_id:user-{i}" for i in range(5)]
|
||||
for key in expected_keys:
|
||||
assert key in deleted_keys, f"Cache key {key} was not invalidated"
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_reset_budget_all(reset_budget_job, mock_prisma_client):
|
||||
# Setup test data with timezone-aware datetime
|
||||
now = datetime.now(timezone.utc)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue