test(proxy): align retained-limit warning regression with existing fixtures

This commit is contained in:
leilei3167 2026-09-16 15:11:02 +00:00
parent 1c69fe21ca
commit 0db9be28b5

View file

@ -3398,34 +3398,30 @@ async def test_validate_update_key_data_returns_team_limit_warnings(monkeypatch)
async def test_validate_update_key_data_warns_on_retained_limits_team_change(monkeypatch):
"""Team reassignment without limit fields still warns on retained over-cap values."""
from unittest.mock import AsyncMock, MagicMock
from litellm.proxy._types import (
KeyTeamLimitWarning,
LiteLLM_TeamTableCachedObj,
LiteLLM_VerificationToken,
LitellmUserRoles,
UpdateKeyRequest,
UserAPIKeyAuth,
)
existing = LiteLLM_VerificationToken(
token="hashed-token",
team_id="team-old",
user_id="user-1",
models=[],
max_budget=100.0,
max_parallel_requests=50,
rpm_limit=None,
tpm_limit=None,
)
new_team = LiteLLM_TeamTableCachedObj(
team_id="team-new",
team_alias="team-new",
max_budget=10.0,
max_parallel_requests=5,
rpm_limit=None,
tpm_limit=None,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[],
)
data = UpdateKeyRequest(key="sk-test-key-123456", team_id="team-new")
user_api_key_dict = UserAPIKeyAuth(
api_key="sk-1234",
user_id="admin",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
data = UpdateKeyRequest(key="sk-test", team_id="team-new")
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
@ -3433,41 +3429,47 @@ async def test_validate_update_key_data_warns_on_retained_limits_team_change(mon
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints._check_team_key_limits",
AsyncMock(return_value=None),
AsyncMock(),
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
AsyncMock(return_value=None),
AsyncMock(),
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.common_key_access_checks",
MagicMock(),
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.enforce_member_can_assign_access_groups",
MagicMock(return_value=None),
MagicMock(),
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.validate_key_team_change",
AsyncMock(return_value=None),
)
# Avoid router requirement inside team-change validation path when mocked above
user = UserAPIKeyAuth(
user_id="user-1",
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
AsyncMock(),
)
warnings = await _validate_update_key_data(
data=data,
existing_key_row=existing,
user_api_key_dict=user,
user_api_key_dict=user_api_key_dict,
llm_router=MagicMock(),
premium_user=True,
prisma_client=MagicMock(),
user_api_key_cache=MagicMock(),
)
assert KeyTeamLimitWarning(
field="max_budget", requested=100.0, effective_team_cap=10.0
) in warnings
assert KeyTeamLimitWarning(
field="max_parallel_requests", requested=50, effective_team_cap=5
) in warnings
assert warnings == (
{
"field": "max_parallel_requests",
"requested": 50,
"effective_team_cap": 5,
},
{
"field": "max_budget",
"requested": 100.0,
"effective_team_cap": 10.0,
},
)