From 2ac8a2ee85cb7d11c56920b5a1eaff91188d9d4e Mon Sep 17 00:00:00 2001 From: Atharva Jaiswal <92455570+AtharvaJaiswal005@users.noreply.github.com> Date: Sat, 28 Mar 2026 21:24:39 +0530 Subject: [PATCH] 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 --- .../common_utils/test_reset_budget_job.py | 206 ++++++++++++++++++ 1 file changed, 206 insertions(+) diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index bc351829a5f..dae3d586f47 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -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)