mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-07 07:58:09 +00:00
fix(proxy): skip team member budget check on non-spending routes
This commit is contained in:
parent
d6c7d8cb6c
commit
9d2a7f1160
1 changed files with 108 additions and 0 deletions
|
|
@ -10593,6 +10593,114 @@ async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spe
|
|||
assert f"TeamMember={user_id}:{team_id}" in exc_info.value.message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"route, expect_blocked",
|
||||
[
|
||||
("/models", False),
|
||||
("/v1/models", False),
|
||||
("/model/info", False),
|
||||
("/key/info", False),
|
||||
("/user/info", False),
|
||||
("/v1/messages", True),
|
||||
("/health/services", True),
|
||||
],
|
||||
)
|
||||
async def test_cached_key_team_member_budget_skips_non_spending_routes(route, expect_blocked):
|
||||
"""An over-budget team member keeps access to model discovery and self-service
|
||||
info routes, while paid LLM routes and side-effectful health routes stay blocked."""
|
||||
from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj
|
||||
from litellm.proxy.common_utils.user_api_key_cache import team_membership_auth_cache_key
|
||||
from litellm.proxy.utils import hash_token
|
||||
|
||||
api_key = "sk-non-spending-routes"
|
||||
hashed_token = hash_token(api_key)
|
||||
team_id = "team-non-spending-routes"
|
||||
user_id = "user-non-spending-routes"
|
||||
team_member_spend = 2.4
|
||||
|
||||
user_api_key_cache = DualCache()
|
||||
await cache_key_object(
|
||||
hashed_token=hashed_token,
|
||||
user_api_key_obj=UserAPIKeyAuth(
|
||||
api_key=api_key,
|
||||
token=hashed_token,
|
||||
team_id=team_id,
|
||||
user_id=user_id,
|
||||
team_member_spend=team_member_spend,
|
||||
),
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=f"team_id:{team_id}",
|
||||
value=LiteLLM_TeamTableCachedObj(team_id=team_id),
|
||||
)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=user_id,
|
||||
value=LiteLLM_UserTable(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER),
|
||||
)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=team_membership_auth_cache_key(team_id=team_id, user_id=user_id),
|
||||
value=LiteLLM_TeamMembership(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
spend=team_member_spend,
|
||||
budget_id="budget-non-spending-routes",
|
||||
litellm_budget_table=LiteLLM_BudgetTable(max_budget=team_member_spend),
|
||||
),
|
||||
)
|
||||
|
||||
mock_request = MagicMock()
|
||||
mock_request.url.path = route
|
||||
mock_request.method = "GET"
|
||||
mock_request.headers = {"authorization": f"Bearer {api_key}"}
|
||||
mock_request.query_params = {}
|
||||
mock_request.state = SimpleNamespace()
|
||||
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.budget_alerts = AsyncMock()
|
||||
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
|
||||
|
||||
async def _auth():
|
||||
return await user_api_key_auth_builder(
|
||||
request=mock_request,
|
||||
api_key=f"Bearer {api_key}",
|
||||
azure_api_key_header="",
|
||||
anthropic_api_key_header=None,
|
||||
google_ai_studio_api_key_header=None,
|
||||
azure_apim_header=None,
|
||||
request_data={},
|
||||
)
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: the builder reads proxy settings from module globals, no injection seam
|
||||
"litellm.proxy.proxy_server.general_settings", {"disable_budget_reservation": True}
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.master_key", "sk-master"), # test-quality-ok: module-global proxy state
|
||||
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: module-global proxy state
|
||||
patch( # test-quality-ok: seed the cached key, team and membership without a DB
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache
|
||||
),
|
||||
patch( # test-quality-ok: module-global proxy state
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj
|
||||
),
|
||||
patch( # test-quality-ok: the live counter needs Redis or a DB; pin the spend the check compares
|
||||
"litellm.proxy.proxy_server.get_current_spend",
|
||||
new=AsyncMock(return_value=team_member_spend),
|
||||
),
|
||||
):
|
||||
if not expect_blocked:
|
||||
result = await _auth()
|
||||
assert result.team_member_spend == team_member_spend
|
||||
return
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _auth()
|
||||
|
||||
assert exc_info.value.type == ProxyErrorTypes.budget_exceeded
|
||||
assert f"TeamMember={user_id}:{team_id}" in exc_info.value.message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"expiry_offset, expect_blocked",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue