From 33970a0a86559dc76d63b1ca36be913aca7e4315 Mon Sep 17 00:00:00 2001 From: Taranum Wasu Date: Sun, 5 Jul 2026 00:12:26 +0530 Subject: [PATCH] test(proxy): cover virtual-key and early-return budget auth paths Add integration tests for the DB lookup update_valid_token path and Check 5b to satisfy codecov patch coverage on user_api_key_auth.py. Co-authored-by: Cursor --- ...t_end_user_model_max_budget_enforcement.py | 176 ++++++++++++++++++ 1 file changed, 176 insertions(+) diff --git a/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py b/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py index e7d321306f0..aabc8bf904d 100644 --- a/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py +++ b/tests/proxy_unit_tests/test_end_user_model_max_budget_enforcement.py @@ -1,4 +1,5 @@ import pytest +from contextlib import contextmanager from unittest.mock import AsyncMock, MagicMock, patch import litellm @@ -45,6 +46,66 @@ def _proxy_server_attrs_for_master_key_auth(): }, limiter +def _end_user_with_model_budget(): + return LiteLLM_EndUserTable( + user_id="customer-1", + blocked=False, + spend=0.0, + litellm_budget_table=LiteLLM_BudgetTable(model_max_budget={MODEL: MODEL_BUDGET}), + ) + + +@contextmanager +def _virtual_key_builder_patches(*, resolved_token: UserAPIKeyAuth): + async def mock_resolve_key(self, hashed_token: str): + from litellm.proxy.auth.resolvers.store import KeyNotInCacheError + + if self._check_cache_only: + raise KeyNotInCacheError(hashed_token) + return resolved_token + + with ( + patch( + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new=mock_resolve_key, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.resolve_and_validate_end_user_id", + new_callable=AsyncMock, + return_value="customer-1", + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_end_user_object", + new_callable=AsyncMock, + return_value=_end_user_with_model_budget(), + ), + patch( + "litellm.proxy.auth.user_api_key_auth._get_model_from_request_context", + return_value=MODEL, + ), + patch( + "litellm.proxy.auth.user_api_key_auth._enforce_key_and_fallback_model_access", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.auth.user_api_key_auth._virtual_key_max_budget_alert_check", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.auth.user_api_key_auth._virtual_key_max_budget_check", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.auth.user_api_key_auth._virtual_key_soft_budget_check", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.auth.auth_exception_handler.seed_request_identity", + ), + ): + yield + + def test_update_valid_token_applies_end_user_model_max_budget_from_params(): valid_token = UserAPIKeyAuth(token="test-key") end_user_params = { @@ -127,6 +188,28 @@ async def test_enforce_end_user_model_max_budget_raises_when_over_budget(): mock_check.assert_awaited_once() +@pytest.mark.asyncio +async def test_enforce_end_user_model_max_budget_returns_early_when_unconfigured(): + from litellm.proxy.auth.user_api_key_auth import _enforce_end_user_model_max_budget_checks + + valid_token = UserAPIKeyAuth(token="test-key", end_user_id="customer-1") + request = MagicMock() + request_data = {"model": MODEL} + + with patch( + "litellm.proxy.proxy_server.model_max_budget_limiter.is_end_user_within_model_budget", + new_callable=AsyncMock, + ) as mock_check: + await _enforce_end_user_model_max_budget_checks( + valid_token=valid_token, + request_data=request_data, + route="/v1/chat/completions", + request=request, + ) + + mock_check.assert_not_awaited() + + @pytest.mark.asyncio async def test_master_key_auth_skips_end_user_model_budget_when_flag_disabled(): from fastapi import Request @@ -187,6 +270,99 @@ async def test_master_key_auth_skips_end_user_model_budget_when_flag_disabled(): setattr(proxy_server, k, v) +@pytest.mark.asyncio +async def test_master_key_auth_passes_when_flag_enabled_and_within_budget(): + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + + attrs, limiter = _proxy_server_attrs_for_master_key_auth() + originals = {k: getattr(proxy_server, k, None) for k in attrs} + flag_original = litellm.enforce_end_user_model_max_budget_on_master_key + litellm.enforce_end_user_model_max_budget_on_master_key = True + + try: + for k, v in attrs.items(): + setattr(proxy_server, k, v) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/v1/chat/completions") + + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.resolve_and_validate_end_user_id", + new_callable=AsyncMock, + return_value="customer-1", + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_end_user_object", + new_callable=AsyncMock, + return_value=_end_user_with_model_budget(), + ), + patch( + "litellm.proxy.auth.user_api_key_auth._get_model_from_request_context", + return_value=MODEL, + ), + ): + result = await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {attrs['master_key']}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"user": "customer-1", "model": MODEL}, + ) + + assert result.end_user_model_max_budget == {MODEL: MODEL_BUDGET} + limiter.is_end_user_within_model_budget.assert_awaited() + finally: + litellm.enforce_end_user_model_max_budget_on_master_key = flag_original + for k, v in originals.items(): + setattr(proxy_server, k, v) + + +@pytest.mark.asyncio +async def test_virtual_key_auth_applies_and_enforces_end_user_model_budget(): + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + + valid_token = UserAPIKeyAuth(api_key="sk-vk-test", token="hashed-valid") + attrs, limiter = _proxy_server_attrs_for_master_key_auth() + attrs["master_key"] = "sk-different-master" + originals = {k: getattr(proxy_server, k, None) for k in attrs} + + try: + for k, v in attrs.items(): + setattr(proxy_server, k, v) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/v1/chat/completions") + + with _virtual_key_builder_patches(resolved_token=valid_token): + result = await _user_api_key_auth_builder( + request=request, + api_key="Bearer sk-vk-test", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"user": "customer-1", "model": MODEL}, + ) + + assert result.end_user_id == "customer-1" + assert result.end_user_model_max_budget == {MODEL: MODEL_BUDGET} + limiter.is_end_user_within_model_budget.assert_awaited() + finally: + for k, v in originals.items(): + setattr(proxy_server, k, v) + + @pytest.mark.asyncio async def test_master_key_auth_enforces_end_user_model_budget_when_flag_enabled(): from fastapi import Request