From 5e255ae8f3cfafbc8430c4b25151c00a7bf2def7 Mon Sep 17 00:00:00 2001 From: Parinith Date: Mon, 11 May 2026 00:35:58 +0530 Subject: [PATCH] test(proxy/auth): add regression test for max_budget enforcement after spend reset (#27300) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds `test_max_budget_enforced_after_spend_reset` to verify that `_virtual_key_max_budget_check` correctly blocks requests after post-reset spend re-accumulates past `max_budget`. Covers the full lifecycle: 1. Pre-reset: exhausted key (spend == max_budget) is blocked. 2. Budget reset zeroes `valid_token.spend` and the spend counter. 3. First post-reset request is allowed (counter reads 0). 4. Post-reset spend accumulates past `max_budget` — blocked again. Co-Authored-By: Claude Sonnet 4.6 --- .../proxy/auth/test_auth_checks.py | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 26f04a4abcb..83f65d399c9 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -2267,6 +2267,73 @@ async def test_virtual_key_budget_check_fallback_no_counter(): assert exc_info.value.current_cost == 15.0 +@pytest.mark.asyncio +async def test_max_budget_enforced_after_spend_reset(): + """ + Regression test for issue #27300: max_budget must be re-enforced after a + monthly (or any-period) spend reset. + + Sequence: + 1. Key is exhausted (spend == max_budget) — pre-reset request is blocked. + 2. Budget reset zeroes the DB row and the spend counter. + 3. First post-reset request (cost=5.0) is allowed — counter reads 0. + 4. After accumulating 5.0 + 6.0 = 11.0 of post-reset spend the counter + exceeds max_budget and the next request is blocked again. + """ + from litellm.proxy.utils import ProxyLogging + + valid_token = UserAPIKeyAuth( + token="test-token-budget-reset-27300", + spend=10.0, + max_budget=10.0, + user_id="test-user", + ) + + proxy_logging_obj = ProxyLogging(user_api_key_cache=None) + proxy_logging_obj.budget_alerts = AsyncMock() + + # Step 1: budget exhausted — request blocked before reset + async def mock_spend_exhausted(counter_key, fallback_spend): + return 10.0 + + with patch("litellm.proxy.proxy_server.get_current_spend", mock_spend_exhausted): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _virtual_key_max_budget_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) + assert exc_info.value.current_cost == 10.0 + assert exc_info.value.max_budget == 10.0 + + # Step 2: simulate budget reset — DB row and counter both zeroed + valid_token.spend = 0.0 + + # Step 3: first post-reset request succeeds (counter reads 0) + async def mock_spend_post_reset(counter_key, fallback_spend): + return 0.0 + + with patch("litellm.proxy.proxy_server.get_current_spend", mock_spend_post_reset): + # Must NOT raise — budget period restarted + await _virtual_key_max_budget_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) + + # Step 4: post-reset spend accumulates past max_budget — blocked again + # 5.0 (first request) + 6.0 (second request) = 11.0 > max_budget=10.0 + async def mock_spend_over_limit(counter_key, fallback_spend): + return 11.0 + + with patch("litellm.proxy.proxy_server.get_current_spend", mock_spend_over_limit): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _virtual_key_max_budget_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) + assert exc_info.value.current_cost == 11.0 + assert exc_info.value.max_budget == 10.0 + + @pytest.mark.asyncio async def test_team_budget_check_reads_from_spend_counter(): """Team budget check should use get_current_spend when counter exists."""