mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
address PR review for bulk key updates
This commit is contained in:
Michael Riad Zaky
parent
ed06f3bcc0
commit
7c580c4d25
4 changed files with 521 additions and 379 deletions
|
|
@ -50,6 +50,7 @@ _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES = frozenset(
|
|||
KeyManagementRoutes.KEY_BLOCK.value,
|
||||
KeyManagementRoutes.KEY_UNBLOCK.value,
|
||||
KeyManagementRoutes.KEY_BULK_UPDATE.value,
|
||||
KeyManagementRoutes.TEAM_KEY_BULK_UPDATE.value,
|
||||
]
|
||||
)
|
||||
|
||||
|
|
@ -671,6 +672,7 @@ class RouteChecks:
|
|||
"/key/service-account/generate",
|
||||
"/key/block",
|
||||
"/key/unblock",
|
||||
"/team/key/bulk_update",
|
||||
]
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -2709,27 +2709,11 @@ async def bulk_update_team_keys(
|
|||
"""
|
||||
Apply one update payload to many keys inside a single team.
|
||||
|
||||
Scoped to a single team — pass `team_id` plus either `key_ids` (specific
|
||||
keys in that team) or `all_keys_in_team=True`. The same `update_fields`
|
||||
payload is applied to every selected key.
|
||||
Pass `team_id` plus either `key_ids` or `all_keys_in_team=True`. The
|
||||
`update_fields` payload is broadcast to every selected key. Per-key
|
||||
failures are returned in `failed_updates` rather than aborting the batch.
|
||||
|
||||
Callable by proxy admins, or by team admins with `KEY_UPDATE` permission
|
||||
for the target team.
|
||||
|
||||
Each key update is processed independently — partial failures are returned
|
||||
in `failed_updates` rather than aborting the whole batch.
|
||||
|
||||
Example request:
|
||||
```bash
|
||||
curl --location 'http://0.0.0.0:4000/team/key/bulk_update' \\
|
||||
--header 'Authorization: Bearer sk-1234' \\
|
||||
--header 'Content-Type: application/json' \\
|
||||
--data '{
|
||||
"team_id": "team-123",
|
||||
"all_keys_in_team": true,
|
||||
"update_fields": {"max_budget": 50.0, "budget_duration": "30d"}
|
||||
}'
|
||||
```
|
||||
Callable by proxy admins, or by team admins with `KEY_UPDATE` permission.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import (
|
||||
llm_router,
|
||||
|
|
@ -2760,10 +2744,19 @@ async def bulk_update_team_keys(
|
|||
},
|
||||
)
|
||||
|
||||
# Resolve which keys to update (single batched query, scoped to team).
|
||||
if data.all_keys_in_team:
|
||||
# "all" excludes blocked/expired — bulk refresh shouldn't revive a key an admin disabled.
|
||||
# `blocked` is Boolean? with no default; `/key/generate` writes NULL. Prisma's `NOT`
|
||||
# excludes NULLs, so explicitly OR `false` with `null` to include them.
|
||||
now = datetime.now(timezone.utc)
|
||||
existing_keys = await prisma_client.db.litellm_verificationtoken.find_many(
|
||||
where={"team_id": data.team_id},
|
||||
where={
|
||||
"team_id": data.team_id,
|
||||
"AND": [
|
||||
{"OR": [{"blocked": False}, {"blocked": None}]},
|
||||
{"OR": [{"expires": None}, {"expires": {"gt": now}}]},
|
||||
],
|
||||
},
|
||||
order={"token": "asc"},
|
||||
take=MAX_BATCH_SIZE + 1,
|
||||
)
|
||||
|
|
@ -2776,12 +2769,52 @@ async def bulk_update_team_keys(
|
|||
)
|
||||
requested_tokens = [row.token for row in existing_keys]
|
||||
else:
|
||||
# validator guarantees key_ids is set and non-empty when all_keys_in_team=False
|
||||
assert data.key_ids is not None
|
||||
if data.key_ids is None or len(data.key_ids) == 0:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": "key_ids must be provided when all_keys_in_team is False"
|
||||
},
|
||||
)
|
||||
# Dedupe by hashed form — duplicates collapse to one update.
|
||||
requested_tokens = []
|
||||
hashed_key_ids = []
|
||||
seen_hashes = set()
|
||||
for k in data.key_ids:
|
||||
h = _hash_token_if_needed(k)
|
||||
if h in seen_hashes:
|
||||
continue
|
||||
seen_hashes.add(h)
|
||||
requested_tokens.append(k)
|
||||
hashed_key_ids.append(h)
|
||||
existing_keys = await prisma_client.db.litellm_verificationtoken.find_many(
|
||||
where={"team_id": data.team_id, "token": {"in": data.key_ids}}
|
||||
where={"team_id": data.team_id, "token": {"in": hashed_key_ids}}
|
||||
)
|
||||
requested_tokens = list(data.key_ids)
|
||||
|
||||
# Anchor membership check on data.team_id (not existing_keys[0]); empty result must still gate non-admins.
|
||||
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
|
||||
auth_anchor = (
|
||||
existing_keys[0]
|
||||
if existing_keys
|
||||
else LiteLLM_VerificationToken(
|
||||
token="__team_scope_auth_check__",
|
||||
team_id=data.team_id,
|
||||
models=[],
|
||||
)
|
||||
)
|
||||
await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
route=KeyManagementRoutes.KEY_UPDATE,
|
||||
prisma_client=prisma_client,
|
||||
existing_key_row=auth_anchor,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
# Block metadata.allowed_passthrough_routes for non-admins — the runtime
|
||||
# route checker reads it from key/team metadata to grant passthrough.
|
||||
_check_passthrough_routes_caller_permission(
|
||||
data=data.update_fields, user_api_key_dict=user_api_key_dict
|
||||
)
|
||||
|
||||
if not requested_tokens:
|
||||
raise HTTPException(
|
||||
|
|
@ -2789,22 +2822,6 @@ async def bulk_update_team_keys(
|
|||
detail={"error": f"No keys found for team {data.team_id}"},
|
||||
)
|
||||
|
||||
# Fail-fast auth: if caller is not proxy admin, verify they have KEY_UPDATE
|
||||
# permission for this team once (using any in-team key as the auth subject).
|
||||
# Per-key checks inside _process_single_key_update will short-circuit on the
|
||||
# cached team object after this.
|
||||
if (
|
||||
user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
|
||||
and existing_keys
|
||||
):
|
||||
await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
route=KeyManagementRoutes.KEY_UPDATE,
|
||||
prisma_client=prisma_client,
|
||||
existing_key_row=existing_keys[0],
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
existing_by_token = {row.token: row for row in existing_keys}
|
||||
update_field_dict = data.update_fields.model_dump(exclude_unset=True)
|
||||
|
||||
|
|
@ -2812,15 +2829,18 @@ async def bulk_update_team_keys(
|
|||
failed_updates: List[FailedKeyUpdate] = []
|
||||
|
||||
for token in requested_tokens:
|
||||
db_token = _hash_token_if_needed(token)
|
||||
try:
|
||||
if token not in existing_by_token:
|
||||
if db_token not in existing_by_token:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={"error": f"Key not found in team {data.team_id}"},
|
||||
)
|
||||
|
||||
# team_id from validated scope, never user payload — drives _check_team_key_limits.
|
||||
update_key_request = UpdateKeyRequest(
|
||||
key=token,
|
||||
team_id=data.team_id,
|
||||
**update_field_dict,
|
||||
)
|
||||
updated_key_info = await _process_single_key_update(
|
||||
|
|
@ -2839,14 +2859,15 @@ async def bulk_update_team_keys(
|
|||
)
|
||||
|
||||
except Exception as e:
|
||||
# Log the hashed prefix — `token` may be a raw sk-... and ERROR logs persist.
|
||||
verbose_proxy_logger.exception(
|
||||
f"Failed to update key {token} in team {data.team_id}: {e}"
|
||||
f"Failed to update key {db_token[:12]}... in team {data.team_id}: {e}"
|
||||
)
|
||||
failed_updates.append(
|
||||
_build_failed_team_key_update(
|
||||
token=token,
|
||||
exception=e,
|
||||
existing_key_row=existing_by_token.get(token),
|
||||
existing_key_row=existing_by_token.get(db_token),
|
||||
)
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,7 @@
|
|||
from datetime import datetime
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any, Dict, List, Literal, Optional
|
||||
|
||||
from pydantic import BaseModel, model_validator
|
||||
|
||||
from litellm.proxy._types import KeyRequestBase
|
||||
from pydantic import BaseModel, ConfigDict, model_validator
|
||||
|
||||
|
||||
class BulkUpdateKeyRequestItem(BaseModel):
|
||||
|
|
@ -45,21 +43,41 @@ class BulkUpdateKeyResponse(BaseModel):
|
|||
failed_updates: List[FailedKeyUpdate]
|
||||
|
||||
|
||||
class KeyUpdateFields(KeyRequestBase):
|
||||
"""
|
||||
Mirror of UpdateKeyRequest minus per-key identifiers (`key`, `key_alias`)
|
||||
and the scope guard (`team_id`). Used as the broadcast payload in
|
||||
BulkUpdateTeamKeysRequest — one set of fields applied to many keys.
|
||||
"""
|
||||
class KeyUpdateFields(BaseModel):
|
||||
"""Allowlist of bulk-broadcastable fields for /team/key/bulk_update; `extra="forbid"` blocks RBAC/ownership/scope mutations even by team admins."""
|
||||
|
||||
duration: Optional[str] = None
|
||||
spend: Optional[float] = None
|
||||
metadata: Optional[dict] = None
|
||||
model_config = ConfigDict(extra="forbid", protected_namespaces=())
|
||||
|
||||
# Budgets
|
||||
max_budget: Optional[float] = None
|
||||
budget_id: Optional[str] = None
|
||||
budget_duration: Optional[str] = None
|
||||
budget_limits: Optional[List[Any]] = None
|
||||
model_max_budget: Optional[Dict[str, Any]] = None
|
||||
|
||||
# Rate limits
|
||||
tpm_limit: Optional[int] = None
|
||||
rpm_limit: Optional[int] = None
|
||||
model_tpm_limit: Optional[Dict[str, Any]] = None
|
||||
model_rpm_limit: Optional[Dict[str, Any]] = None
|
||||
max_parallel_requests: Optional[int] = None
|
||||
rpm_limit_type: Optional[
|
||||
Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"]
|
||||
] = None
|
||||
tpm_limit_type: Optional[
|
||||
Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"]
|
||||
] = None
|
||||
|
||||
# Temporary budget grants (auto-expire). `spend` deliberately omitted — bulk-zeroing it bypasses budget enforcement; admin-only via /key/update.
|
||||
temp_budget_increase: Optional[float] = None
|
||||
temp_budget_expiry: Optional[datetime] = None
|
||||
auto_rotate: Optional[bool] = None
|
||||
rotation_interval: Optional[str] = None
|
||||
organization_id: Optional[str] = None
|
||||
|
||||
# Expiry
|
||||
duration: Optional[str] = None
|
||||
|
||||
# Operational metadata
|
||||
tags: Optional[List[str]] = None
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_temp_budget(self) -> "KeyUpdateFields":
|
||||
|
|
@ -71,28 +89,15 @@ class KeyUpdateFields(KeyRequestBase):
|
|||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def reject_per_key_or_scope_fields(self) -> "KeyUpdateFields":
|
||||
forbidden = [
|
||||
name
|
||||
for name in ("key", "key_alias", "team_id")
|
||||
if getattr(self, name, None) is not None
|
||||
]
|
||||
if forbidden:
|
||||
raise ValueError(
|
||||
f"Fields not allowed in update_fields for bulk team key updates: "
|
||||
f"{forbidden}. `key`/`key_alias` are per-key identifiers; `team_id` "
|
||||
f"is the scope guard set at the request top level."
|
||||
)
|
||||
def require_at_least_one_field(self) -> "KeyUpdateFields":
|
||||
# Reject empty payload — would iterate every key with no-op writes.
|
||||
if not self.model_fields_set:
|
||||
raise ValueError("update_fields must specify at least one field to update.")
|
||||
return self
|
||||
|
||||
|
||||
class BulkUpdateTeamKeysRequest(BaseModel):
|
||||
"""
|
||||
Request for applying one update payload to many keys inside a single team.
|
||||
|
||||
Exactly one of `key_ids` (specific keys in the team) or `all_keys_in_team`
|
||||
(every key in the team) must be provided.
|
||||
"""
|
||||
"""Apply one update payload to many keys inside a team; provide either `key_ids` or `all_keys_in_team=True`."""
|
||||
|
||||
team_id: str
|
||||
key_ids: Optional[List[str]] = None
|
||||
|
|
|
|||
|
|
@ -10023,6 +10023,9 @@ async def test_execute_virtual_key_regeneration_cache_invalidation_with_token_ha
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
_BULK_PKG = "litellm.proxy.management_endpoints.key_management_endpoints"
|
||||
|
||||
|
||||
def _make_team_key(token: str, team_id: str = "team-abc") -> LiteLLM_VerificationToken:
|
||||
return LiteLLM_VerificationToken(
|
||||
token=token,
|
||||
|
|
@ -10033,222 +10036,228 @@ def _make_team_key(token: str, team_id: str = "team-abc") -> LiteLLM_Verificatio
|
|||
)
|
||||
|
||||
|
||||
def _patch_team_keys_helpers(monkeypatch, *, hash_lookup):
|
||||
"""Common monkeypatching for the bulk_update_team_keys path."""
|
||||
def _admin() -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin"
|
||||
)
|
||||
|
||||
|
||||
def _internal_user() -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-iu", user_id="iu"
|
||||
)
|
||||
|
||||
|
||||
def _updated(payload):
|
||||
m = MagicMock()
|
||||
m.model_dump.return_value = payload
|
||||
return m
|
||||
|
||||
|
||||
def _setup_team_keys_mocks(
|
||||
monkeypatch,
|
||||
*,
|
||||
find_many=None,
|
||||
find_unique=None,
|
||||
update_data=None,
|
||||
hash_identity=True,
|
||||
):
|
||||
"""Set up mocks for bulk_update_team_keys; returns mock_prisma."""
|
||||
mock_prisma = AsyncMock()
|
||||
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(
|
||||
return_value=[] if find_many is None else find_many
|
||||
)
|
||||
if find_unique is not None:
|
||||
mock_prisma.db.litellm_verificationtoken.find_unique = find_unique
|
||||
if update_data is not None:
|
||||
mock_prisma.update_data = update_data
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_update", None)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data",
|
||||
f"{_BULK_PKG}.prepare_key_update_data",
|
||||
AsyncMock(return_value={"max_budget": 50.0}),
|
||||
)
|
||||
monkeypatch.setattr(f"{_BULK_PKG}._delete_cache_key_object", AsyncMock())
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
|
||||
lambda token: hash_lookup[token],
|
||||
f"{_BULK_PKG}.KeyManagementEventHooks.async_key_updated_hook", AsyncMock()
|
||||
)
|
||||
monkeypatch.setattr(f"{_BULK_PKG}.get_team_object", AsyncMock(return_value=None))
|
||||
monkeypatch.setattr(f"{_BULK_PKG}._check_team_key_limits", AsyncMock())
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
|
||||
f"{_BULK_PKG}.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
|
||||
AsyncMock(),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook",
|
||||
AsyncMock(),
|
||||
if hash_identity:
|
||||
# Tests use already-hashed tokens; the raw-sk regression opts out.
|
||||
monkeypatch.setattr(f"{_BULK_PKG}._hash_token_if_needed", lambda token: token)
|
||||
return mock_prisma
|
||||
|
||||
|
||||
async def _call_as_admin(data):
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
bulk_update_team_keys,
|
||||
)
|
||||
|
||||
return await bulk_update_team_keys(
|
||||
data=data, user_api_key_dict=_admin(), litellm_changed_by=None
|
||||
)
|
||||
|
||||
|
||||
# ---- happy paths ----------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_update_team_keys_success_with_key_ids(monkeypatch):
|
||||
"""Proxy admin updates two explicitly listed keys; both succeed."""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
bulk_update_team_keys,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
|
||||
BulkUpdateTeamKeysRequest,
|
||||
KeyUpdateFields,
|
||||
)
|
||||
|
||||
key_a = _make_team_key("tok-a")
|
||||
key_b = _make_team_key("tok-b")
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
|
||||
return_value=[key_a, key_b]
|
||||
)
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
|
||||
side_effect=[key_a, key_b]
|
||||
)
|
||||
updated_a = MagicMock()
|
||||
updated_a.model_dump.return_value = {"max_budget": 50.0, "user_id": "user-123"}
|
||||
updated_b = MagicMock()
|
||||
updated_b.model_dump.return_value = {"max_budget": 50.0, "user_id": "user-123"}
|
||||
mock_prisma_client.update_data = AsyncMock(
|
||||
side_effect=[{"data": updated_a}, {"data": updated_b}]
|
||||
)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_update", None)
|
||||
_patch_team_keys_helpers(
|
||||
monkeypatch, hash_lookup={"tok-a": "hashed-a", "tok-b": "hashed-b"}
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
|
||||
AsyncMock(),
|
||||
)
|
||||
|
||||
request = BulkUpdateTeamKeysRequest(
|
||||
team_id="team-abc",
|
||||
key_ids=["tok-a", "tok-b"],
|
||||
update_fields=KeyUpdateFields(max_budget=50.0),
|
||||
)
|
||||
response = await bulk_update_team_keys(
|
||||
data=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-admin",
|
||||
user_id="admin",
|
||||
keys = [_make_team_key("tok-a"), _make_team_key("tok-b")]
|
||||
mock = _setup_team_keys_mocks(
|
||||
monkeypatch,
|
||||
find_many=keys,
|
||||
find_unique=AsyncMock(side_effect=keys),
|
||||
update_data=AsyncMock(
|
||||
side_effect=[{"data": _updated({"max_budget": 50.0})}] * 2
|
||||
),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert response.total_requested == 2
|
||||
response = await _call_as_admin(
|
||||
BulkUpdateTeamKeysRequest(
|
||||
team_id="team-abc",
|
||||
key_ids=["tok-a", "tok-b"],
|
||||
update_fields=KeyUpdateFields(max_budget=50.0),
|
||||
)
|
||||
)
|
||||
|
||||
assert len(response.successful_updates) == 2
|
||||
assert len(response.failed_updates) == 0
|
||||
assert {u.key for u in response.successful_updates} == {"tok-a", "tok-b"}
|
||||
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_many.assert_awaited_once()
|
||||
where_arg = (
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_many.await_args.kwargs[
|
||||
"where"
|
||||
]
|
||||
)
|
||||
assert where_arg["team_id"] == "team-abc"
|
||||
assert where_arg["token"] == {"in": ["tok-a", "tok-b"]}
|
||||
where = mock.db.litellm_verificationtoken.find_many.await_args.kwargs["where"]
|
||||
assert where["team_id"] == "team-abc"
|
||||
assert where["token"] == {"in": ["tok-a", "tok-b"]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_update_team_keys_success_all_keys_in_team(monkeypatch):
|
||||
"""`all_keys_in_team=True` resolves via find_many and updates every key."""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
bulk_update_team_keys,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
|
||||
BulkUpdateTeamKeysRequest,
|
||||
KeyUpdateFields,
|
||||
)
|
||||
|
||||
keys = [_make_team_key(f"tok-{i}") for i in range(3)]
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
|
||||
return_value=keys
|
||||
)
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
|
||||
side_effect=keys
|
||||
)
|
||||
updated_objs = [MagicMock() for _ in keys]
|
||||
for obj in updated_objs:
|
||||
obj.model_dump.return_value = {"max_budget": 50.0}
|
||||
mock_prisma_client.update_data = AsyncMock(
|
||||
side_effect=[{"data": obj} for obj in updated_objs]
|
||||
)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_update", None)
|
||||
_patch_team_keys_helpers(
|
||||
monkeypatch, hash_lookup={f"tok-{i}": f"hashed-{i}" for i in range(3)}
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
|
||||
AsyncMock(),
|
||||
)
|
||||
|
||||
request = BulkUpdateTeamKeysRequest(
|
||||
team_id="team-abc",
|
||||
all_keys_in_team=True,
|
||||
update_fields=KeyUpdateFields(max_budget=50.0),
|
||||
)
|
||||
response = await bulk_update_team_keys(
|
||||
data=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-admin",
|
||||
user_id="admin",
|
||||
mock = _setup_team_keys_mocks(
|
||||
monkeypatch,
|
||||
find_many=keys,
|
||||
find_unique=AsyncMock(side_effect=keys),
|
||||
update_data=AsyncMock(
|
||||
side_effect=[{"data": _updated({"max_budget": 50.0})}] * 3
|
||||
),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert response.total_requested == 3
|
||||
assert len(response.successful_updates) == 3
|
||||
assert len(response.failed_updates) == 0
|
||||
where_kwargs = (
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_many.await_args.kwargs
|
||||
response = await _call_as_admin(
|
||||
BulkUpdateTeamKeysRequest(
|
||||
team_id="team-abc",
|
||||
all_keys_in_team=True,
|
||||
update_fields=KeyUpdateFields(max_budget=50.0),
|
||||
)
|
||||
)
|
||||
|
||||
assert len(response.successful_updates) == 3
|
||||
where = mock.db.litellm_verificationtoken.find_many.await_args.kwargs["where"]
|
||||
# `blocked` is Boolean? with no default → /key/generate writes NULL. Prisma's
|
||||
# NOT excludes NULLs, so the filter has to OR `false` with `null` explicitly.
|
||||
blocked_or, expires_or = where["AND"][0]["OR"], where["AND"][1]["OR"]
|
||||
assert {"blocked": False} in blocked_or and {"blocked": None} in blocked_or
|
||||
assert {"expires": None} in expires_or
|
||||
assert any(
|
||||
"gt" in c.get("expires", {})
|
||||
for c in expires_or
|
||||
if isinstance(c.get("expires"), dict)
|
||||
)
|
||||
assert where_kwargs["take"] == 501
|
||||
assert where_kwargs["where"] == {"team_id": "team-abc"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_update_team_keys_key_not_in_team(monkeypatch):
|
||||
"""key_ids includes a token not in the team — that token fails, others succeed."""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
bulk_update_team_keys,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
|
||||
BulkUpdateTeamKeysRequest,
|
||||
KeyUpdateFields,
|
||||
)
|
||||
|
||||
in_team = _make_team_key("tok-a")
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
|
||||
return_value=[in_team]
|
||||
)
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
|
||||
return_value=in_team
|
||||
)
|
||||
updated = MagicMock()
|
||||
updated.model_dump.return_value = {"max_budget": 50.0}
|
||||
mock_prisma_client.update_data = AsyncMock(return_value={"data": updated})
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_update", None)
|
||||
_patch_team_keys_helpers(monkeypatch, hash_lookup={"tok-a": "hashed-a"})
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
|
||||
AsyncMock(),
|
||||
_setup_team_keys_mocks(
|
||||
monkeypatch,
|
||||
find_many=[in_team],
|
||||
find_unique=AsyncMock(return_value=in_team),
|
||||
update_data=AsyncMock(return_value={"data": _updated({"max_budget": 50.0})}),
|
||||
)
|
||||
|
||||
request = BulkUpdateTeamKeysRequest(
|
||||
team_id="team-abc",
|
||||
key_ids=["tok-a", "tok-foreign"],
|
||||
update_fields=KeyUpdateFields(max_budget=50.0),
|
||||
response = await _call_as_admin(
|
||||
BulkUpdateTeamKeysRequest(
|
||||
team_id="team-abc",
|
||||
key_ids=["tok-a", "tok-foreign"],
|
||||
update_fields=KeyUpdateFields(max_budget=50.0),
|
||||
)
|
||||
)
|
||||
response = await bulk_update_team_keys(
|
||||
data=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-admin",
|
||||
user_id="admin",
|
||||
),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert response.total_requested == 2
|
||||
assert [u.key for u in response.successful_updates] == ["tok-a"]
|
||||
assert [u.key for u in response.failed_updates] == ["tok-foreign"]
|
||||
assert "not found in team" in response.failed_updates[0].failed_reason
|
||||
|
||||
|
||||
# ---- error paths ----------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_update_team_keys_team_admin_authorized(monkeypatch):
|
||||
"""Non-admin caller passes the upfront team-permission check and succeeds."""
|
||||
async def test_bulk_update_team_keys_batch_size_cap(monkeypatch):
|
||||
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
|
||||
BulkUpdateTeamKeysRequest,
|
||||
KeyUpdateFields,
|
||||
)
|
||||
|
||||
_setup_team_keys_mocks(
|
||||
monkeypatch,
|
||||
find_many=[_make_team_key(f"tok-{i}") for i in range(501)],
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _call_as_admin(
|
||||
BulkUpdateTeamKeysRequest(
|
||||
team_id="team-abc",
|
||||
all_keys_in_team=True,
|
||||
update_fields=KeyUpdateFields(max_budget=50.0),
|
||||
)
|
||||
)
|
||||
assert exc.value.status_code == 400
|
||||
assert "more than 500" in exc.value.detail["error"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_update_team_keys_empty_team_returns_404(monkeypatch):
|
||||
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
|
||||
BulkUpdateTeamKeysRequest,
|
||||
KeyUpdateFields,
|
||||
)
|
||||
|
||||
_setup_team_keys_mocks(monkeypatch, find_many=[])
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _call_as_admin(
|
||||
BulkUpdateTeamKeysRequest(
|
||||
team_id="team-empty",
|
||||
all_keys_in_team=True,
|
||||
update_fields=KeyUpdateFields(max_budget=50.0),
|
||||
)
|
||||
)
|
||||
assert exc.value.status_code == 404
|
||||
|
||||
|
||||
# ---- auth -----------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_update_team_keys_team_member_with_permission(monkeypatch):
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
bulk_update_team_keys,
|
||||
)
|
||||
|
|
@ -10258,53 +10267,34 @@ async def test_bulk_update_team_keys_team_admin_authorized(monkeypatch):
|
|||
)
|
||||
|
||||
key_a = _make_team_key("tok-a")
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
|
||||
return_value=[key_a]
|
||||
_setup_team_keys_mocks(
|
||||
monkeypatch,
|
||||
find_many=[key_a],
|
||||
find_unique=AsyncMock(return_value=key_a),
|
||||
update_data=AsyncMock(return_value={"data": _updated({"max_budget": 50.0})}),
|
||||
)
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
|
||||
return_value=key_a
|
||||
)
|
||||
updated = MagicMock()
|
||||
updated.model_dump.return_value = {"max_budget": 50.0}
|
||||
mock_prisma_client.update_data = AsyncMock(return_value={"data": updated})
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_update", None)
|
||||
_patch_team_keys_helpers(monkeypatch, hash_lookup={"tok-a": "hashed-a"})
|
||||
|
||||
auth_check = AsyncMock()
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
|
||||
f"{_BULK_PKG}.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
|
||||
auth_check,
|
||||
)
|
||||
|
||||
request = BulkUpdateTeamKeysRequest(
|
||||
team_id="team-abc",
|
||||
all_keys_in_team=True,
|
||||
update_fields=KeyUpdateFields(max_budget=50.0),
|
||||
)
|
||||
response = await bulk_update_team_keys(
|
||||
data=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
api_key="sk-team-admin",
|
||||
user_id="team-admin-user",
|
||||
data=BulkUpdateTeamKeysRequest(
|
||||
team_id="team-abc",
|
||||
all_keys_in_team=True,
|
||||
update_fields=KeyUpdateFields(max_budget=50.0),
|
||||
),
|
||||
user_api_key_dict=_internal_user(),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert len(response.successful_updates) == 1
|
||||
# upfront check (1) + per-key check inside _process_single_key_update (1)
|
||||
# Upfront check + per-key check inside _process_single_key_update
|
||||
assert auth_check.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_update_team_keys_team_admin_no_permission(monkeypatch):
|
||||
"""Non-admin without KEY_UPDATE permission raises ProxyException upfront (no partial result)."""
|
||||
async def test_bulk_update_team_keys_team_member_no_permission(monkeypatch):
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
bulk_update_team_keys,
|
||||
)
|
||||
|
|
@ -10313,20 +10303,9 @@ async def test_bulk_update_team_keys_team_admin_no_permission(monkeypatch):
|
|||
KeyUpdateFields,
|
||||
)
|
||||
|
||||
key_a = _make_team_key("tok-a")
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
|
||||
return_value=[key_a]
|
||||
)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_update", None)
|
||||
|
||||
mock = _setup_team_keys_mocks(monkeypatch, find_many=[_make_team_key("tok-a")])
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
|
||||
f"{_BULK_PKG}.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
|
||||
AsyncMock(
|
||||
side_effect=ProxyException(
|
||||
message="not in team",
|
||||
|
|
@ -10337,67 +10316,255 @@ async def test_bulk_update_team_keys_team_admin_no_permission(monkeypatch):
|
|||
),
|
||||
)
|
||||
|
||||
request = BulkUpdateTeamKeysRequest(
|
||||
team_id="team-abc",
|
||||
all_keys_in_team=True,
|
||||
update_fields=KeyUpdateFields(max_budget=50.0),
|
||||
with pytest.raises(ProxyException):
|
||||
await bulk_update_team_keys(
|
||||
data=BulkUpdateTeamKeysRequest(
|
||||
team_id="team-abc",
|
||||
all_keys_in_team=True,
|
||||
update_fields=KeyUpdateFields(max_budget=1.0),
|
||||
),
|
||||
user_api_key_dict=_internal_user(),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
mock.update_data.assert_not_called()
|
||||
|
||||
|
||||
# ---- pydantic-layer validation -------------------------------------------
|
||||
|
||||
|
||||
def test_bulk_update_team_keys_request_validation():
|
||||
"""Allowlist (extra='forbid'), empty-payload rejection, and selection XOR."""
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
|
||||
BulkUpdateTeamKeysRequest,
|
||||
KeyUpdateFields,
|
||||
)
|
||||
|
||||
forbidden = [
|
||||
"key",
|
||||
"key_alias",
|
||||
"team_id",
|
||||
"allowed_routes",
|
||||
"allowed_passthrough_routes",
|
||||
"permissions",
|
||||
"object_permission",
|
||||
"access_group_ids",
|
||||
"user_id",
|
||||
"organization_id",
|
||||
"blocked",
|
||||
"key_type",
|
||||
"models",
|
||||
"config",
|
||||
"router_settings",
|
||||
"spend",
|
||||
]
|
||||
for f in forbidden:
|
||||
with pytest.raises(ValidationError, match=f):
|
||||
KeyUpdateFields(**{f: True})
|
||||
|
||||
with pytest.raises(ValidationError, match="at least one"):
|
||||
KeyUpdateFields()
|
||||
|
||||
assert KeyUpdateFields(max_budget=50.0, tags=["x"]).max_budget == 50.0
|
||||
|
||||
valid = KeyUpdateFields(max_budget=10)
|
||||
with pytest.raises(ValidationError):
|
||||
BulkUpdateTeamKeysRequest(
|
||||
team_id="t", key_ids=["k"], all_keys_in_team=True, update_fields=valid
|
||||
)
|
||||
with pytest.raises(ValidationError):
|
||||
BulkUpdateTeamKeysRequest(team_id="t", update_fields=valid)
|
||||
|
||||
|
||||
# ---- security regressions ------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_update_team_keys_hashes_raw_sk_key_ids(monkeypatch):
|
||||
"""Regression: raw sk-... key_ids must be hashed before the find_many lookup."""
|
||||
from litellm.proxy._types import hash_token
|
||||
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
|
||||
BulkUpdateTeamKeysRequest,
|
||||
KeyUpdateFields,
|
||||
)
|
||||
|
||||
raw_sk = "sk-rawkey1234567890"
|
||||
hashed = hash_token(raw_sk)
|
||||
row = LiteLLM_VerificationToken(
|
||||
token=hashed, user_id="u", models=[], team_id="team-abc", max_budget=None
|
||||
)
|
||||
mock = _setup_team_keys_mocks(
|
||||
monkeypatch,
|
||||
find_many=[row],
|
||||
find_unique=AsyncMock(return_value=row),
|
||||
update_data=AsyncMock(return_value={"data": _updated({"max_budget": 50.0})}),
|
||||
hash_identity=False,
|
||||
)
|
||||
|
||||
response = await _call_as_admin(
|
||||
BulkUpdateTeamKeysRequest(
|
||||
team_id="team-abc",
|
||||
key_ids=[raw_sk],
|
||||
update_fields=KeyUpdateFields(max_budget=50.0),
|
||||
)
|
||||
)
|
||||
where = mock.db.litellm_verificationtoken.find_many.await_args.kwargs["where"]
|
||||
assert where["token"] == {"in": [hashed]}
|
||||
# Response reports the user-supplied form, not the hash.
|
||||
assert response.successful_updates[0].key == raw_sk
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_update_team_keys_auth_check_runs_when_no_keys_match(monkeypatch):
|
||||
"""Regression: non-admin with bogus key_ids must still hit the membership gate."""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
bulk_update_team_keys,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
|
||||
BulkUpdateTeamKeysRequest,
|
||||
KeyUpdateFields,
|
||||
)
|
||||
|
||||
mock = _setup_team_keys_mocks(monkeypatch, find_many=[])
|
||||
auth_check = AsyncMock(
|
||||
side_effect=ProxyException(
|
||||
message="not in team",
|
||||
type="team_member_permission_error",
|
||||
param="/key/update",
|
||||
code=401,
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
f"{_BULK_PKG}.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
|
||||
auth_check,
|
||||
)
|
||||
|
||||
with pytest.raises(ProxyException):
|
||||
await bulk_update_team_keys(
|
||||
data=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
api_key="sk-outsider",
|
||||
user_id="outsider",
|
||||
data=BulkUpdateTeamKeysRequest(
|
||||
team_id="victim-team",
|
||||
key_ids=["bogus-1", "bogus-2"],
|
||||
update_fields=KeyUpdateFields(max_budget=1.0),
|
||||
),
|
||||
user_api_key_dict=_internal_user(),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
mock_prisma_client.update_data.assert_not_called()
|
||||
# Anchored on data.team_id, not existing_keys[0].
|
||||
assert auth_check.await_args.kwargs["existing_key_row"].team_id == "victim-team"
|
||||
mock.update_data.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_update_team_keys_xor_validation():
|
||||
"""Both selection modes set, or neither set, raises at request construction."""
|
||||
async def test_bulk_update_team_keys_does_not_log_raw_sk_token_on_failure(
|
||||
monkeypatch, caplog
|
||||
):
|
||||
"""Regression: per-key failure must not log the raw sk-... (ERROR-level logs persist)."""
|
||||
import logging
|
||||
|
||||
from litellm.proxy._types import hash_token
|
||||
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
|
||||
BulkUpdateTeamKeysRequest,
|
||||
KeyUpdateFields,
|
||||
)
|
||||
|
||||
with pytest.raises(Exception):
|
||||
BulkUpdateTeamKeysRequest(
|
||||
team_id="t",
|
||||
key_ids=["k"],
|
||||
all_keys_in_team=True,
|
||||
update_fields=KeyUpdateFields(max_budget=10.0),
|
||||
)
|
||||
raw_sk = "sk-supersecret1234567890"
|
||||
row = LiteLLM_VerificationToken(
|
||||
token=hash_token(raw_sk),
|
||||
user_id="u",
|
||||
models=[],
|
||||
team_id="team-abc",
|
||||
max_budget=None,
|
||||
)
|
||||
_setup_team_keys_mocks(
|
||||
monkeypatch,
|
||||
find_many=[row],
|
||||
# Force the per-key fetch to raise so the exception logger fires.
|
||||
find_unique=AsyncMock(side_effect=RuntimeError("boom")),
|
||||
hash_identity=False,
|
||||
)
|
||||
|
||||
with pytest.raises(Exception):
|
||||
BulkUpdateTeamKeysRequest(
|
||||
team_id="t",
|
||||
update_fields=KeyUpdateFields(max_budget=10.0),
|
||||
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
|
||||
response = await _call_as_admin(
|
||||
BulkUpdateTeamKeysRequest(
|
||||
team_id="team-abc",
|
||||
key_ids=[raw_sk],
|
||||
update_fields=KeyUpdateFields(max_budget=50.0),
|
||||
)
|
||||
)
|
||||
assert len(response.failed_updates) == 1
|
||||
log_text = "\n".join(r.getMessage() for r in caplog.records)
|
||||
assert raw_sk not in log_text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_update_team_keys_forbidden_fields():
|
||||
"""update_fields.team_id / .key / .key_alias are rejected on construction."""
|
||||
async def test_bulk_update_team_keys_propagates_team_id_to_per_key_request(monkeypatch):
|
||||
"""Regression: per-key UpdateKeyRequest carries data.team_id (gates _check_team_key_limits)."""
|
||||
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
|
||||
BulkUpdateTeamKeysRequest,
|
||||
KeyUpdateFields,
|
||||
)
|
||||
|
||||
with pytest.raises(Exception):
|
||||
KeyUpdateFields(team_id="other-team")
|
||||
with pytest.raises(Exception):
|
||||
KeyUpdateFields(key="sk-abc")
|
||||
with pytest.raises(Exception):
|
||||
KeyUpdateFields(key_alias="alias")
|
||||
_setup_team_keys_mocks(monkeypatch, find_many=[_make_team_key("tok-a")])
|
||||
captured = []
|
||||
|
||||
async def fake_process(*, update_key_request, **kw):
|
||||
captured.append(update_key_request)
|
||||
return {"max_budget": update_key_request.max_budget}
|
||||
|
||||
monkeypatch.setattr(f"{_BULK_PKG}._process_single_key_update", fake_process)
|
||||
|
||||
await _call_as_admin(
|
||||
BulkUpdateTeamKeysRequest(
|
||||
team_id="team-abc",
|
||||
key_ids=["tok-a"],
|
||||
update_fields=KeyUpdateFields(
|
||||
tpm_limit=10_000, tpm_limit_type="guaranteed_throughput"
|
||||
),
|
||||
)
|
||||
)
|
||||
assert captured[0].team_id == "team-abc"
|
||||
assert captured[0].tpm_limit_type == "guaranteed_throughput"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_update_team_keys_batch_size_cap(monkeypatch):
|
||||
"""find_many returning > MAX_BATCH_SIZE rows raises 400."""
|
||||
async def test_bulk_update_team_keys_dedupes_key_ids(monkeypatch):
|
||||
"""Duplicate key_ids collapse to a single update (no redundant DB writes, no inflated counts)."""
|
||||
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
|
||||
BulkUpdateTeamKeysRequest,
|
||||
KeyUpdateFields,
|
||||
)
|
||||
|
||||
key_a = _make_team_key("tok-a")
|
||||
update_data = AsyncMock(return_value={"data": _updated({"max_budget": 50.0})})
|
||||
_setup_team_keys_mocks(
|
||||
monkeypatch,
|
||||
find_many=[key_a],
|
||||
find_unique=AsyncMock(return_value=key_a),
|
||||
update_data=update_data,
|
||||
)
|
||||
|
||||
response = await _call_as_admin(
|
||||
BulkUpdateTeamKeysRequest(
|
||||
team_id="team-abc",
|
||||
key_ids=["tok-a", "tok-a", "tok-a"],
|
||||
update_fields=KeyUpdateFields(max_budget=50.0),
|
||||
)
|
||||
)
|
||||
|
||||
assert response.total_requested == 1
|
||||
assert len(response.successful_updates) == 1
|
||||
assert len(response.failed_updates) == 0
|
||||
update_data.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_update_team_keys_blocks_metadata_allowed_passthrough_routes(
|
||||
monkeypatch,
|
||||
):
|
||||
"""Non-admin can't grant passthrough access by smuggling allowed_passthrough_routes through metadata."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
bulk_update_team_keys,
|
||||
)
|
||||
|
|
@ -10406,76 +10573,23 @@ async def test_bulk_update_team_keys_batch_size_cap(monkeypatch):
|
|||
KeyUpdateFields,
|
||||
)
|
||||
|
||||
too_many = [_make_team_key(f"tok-{i}") for i in range(501)]
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
|
||||
return_value=too_many
|
||||
)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_update", None)
|
||||
mock = _setup_team_keys_mocks(monkeypatch, find_many=[_make_team_key("tok-a")])
|
||||
|
||||
request = BulkUpdateTeamKeysRequest(
|
||||
team_id="team-abc",
|
||||
all_keys_in_team=True,
|
||||
update_fields=KeyUpdateFields(max_budget=50.0),
|
||||
update_fields=KeyUpdateFields(
|
||||
metadata={"allowed_passthrough_routes": ["/admin/*"]}
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await bulk_update_team_keys(
|
||||
data=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-admin",
|
||||
user_id="admin",
|
||||
),
|
||||
user_api_key_dict=_internal_user(),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert exc.value.status_code == 400
|
||||
assert "more than 500" in exc.value.detail["error"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_update_team_keys_no_keys_found(monkeypatch):
|
||||
"""find_many returns nothing → 404 (clear signal, not silent empty success)."""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
bulk_update_team_keys,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
|
||||
BulkUpdateTeamKeysRequest,
|
||||
KeyUpdateFields,
|
||||
)
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
|
||||
return_value=[]
|
||||
)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_update", None)
|
||||
|
||||
request = BulkUpdateTeamKeysRequest(
|
||||
team_id="team-empty",
|
||||
all_keys_in_team=True,
|
||||
update_fields=KeyUpdateFields(max_budget=50.0),
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await bulk_update_team_keys(
|
||||
data=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-admin",
|
||||
user_id="admin",
|
||||
),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert exc.value.status_code == 404
|
||||
assert exc.value.status_code == 403
|
||||
assert "allowed_passthrough_routes" in str(exc.value.detail)
|
||||
mock.update_data.assert_not_called()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue