mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
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 <cursoragent@cursor.com>
This commit is contained in:
Taranum Wasu
parent
0938b65367
commit
33970a0a86
1 changed files with 176 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue