From a678541da8f043ae32b2c970c09859c8b8369bd1 Mon Sep 17 00:00:00 2001 From: Taranum Wasu Date: Wed, 8 Jul 2026 03:53:45 +0530 Subject: [PATCH] test(proxy): cover non-LLM-route early return in master-key budget helper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Codecov flagged 2 missing lines on `user_api_key_auth.py` (92.59% patch coverage, 2 lines uncovered). They are the third guard inside `_maybe_enforce_master_key_end_user_model_max_budget`: ```python if not RouteChecks.is_llm_api_route(route=route): return ``` Existing tests exercise the flag-enabled/disabled and master-key-vs- virtual-key guards via the `_user_api_key_auth_builder` path, but none use a non-LLM route (the three builder tests all use `/v1/chat/completions`). New parametrised test `test_master_key_budget_early_return_for_non_llm_routes` calls `_maybe_enforce_master_key_end_user_model_max_budget` directly with `/health/liveliness`, `/health/readiness`, `/key/info`, and `/metrics` — all routes where `RouteChecks.is_llm_api_route` returns False — and asserts that `model_max_budget_limiter.is_end_user_within_model_budget` is never awaited. This brings the patch coverage to 100% on the modified file and pins the contract: budget enforcement is LLM-route-only, even when the master-key flag is enabled. Signed-off-by: Taranum Wasu Co-authored-by: Cursor --- ...t_end_user_model_max_budget_enforcement.py | 50 +++++++++++++++++++ 1 file changed, 50 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 3ff2904359a..c29ad8fec42 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 @@ -569,3 +569,53 @@ async def test_cached_proxy_admin_virtual_key_skips_master_key_budget_enforcemen 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 +@pytest.mark.parametrize( + "route", + [ + "/health/liveliness", + "/health/readiness", + "/key/info", + "/metrics", + ], +) +async def test_master_key_budget_early_return_for_non_llm_routes(route): + """Branch coverage for ``_maybe_enforce_master_key_end_user_model_max_budget``. + + The flag and master-key guards in the helper are exercised by the + builder-level tests above; this test pins the third guard — the + ``is_llm_api_route`` early return — so a future refactor that + accidentally runs the budget check on health/metrics routes fails + this test loudly. Hits line 2876 of ``user_api_key_auth.py``. + """ + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.auth.user_api_key_auth import ( + _maybe_enforce_master_key_end_user_model_max_budget, + ) + + flag_original = litellm.enforce_end_user_model_max_budget_on_master_key + litellm.enforce_end_user_model_max_budget_on_master_key = True + + try: + valid_token = UserAPIKeyAuth( + api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, + token=LITELLM_PROXY_MASTER_KEY_ALIAS, + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + with patch( + "litellm.proxy.proxy_server.model_max_budget_limiter.is_end_user_within_model_budget", + new_callable=AsyncMock, + ) as mock_check: + await _maybe_enforce_master_key_end_user_model_max_budget( + valid_token=valid_token, + request_data={"user": "customer-1", "model": MODEL}, + route=route, + request=MagicMock(), + ) + mock_check.assert_not_awaited() + finally: + litellm.enforce_end_user_model_max_budget_on_master_key = flag_original