scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-12 03:48:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 57

test: migrate tests/proxy_behavior into integration and unit tiers (#45939)

Browse source 
* test: migrate tests/proxy_behavior into integration and unit tiers (partial)

* test(integration): run settings-dependent management tests on an owned proxy

* test(integration): assert stored values in the rewritten management rows

* test(integration): pin sibling team allocations against the org guaranteed throughput cap

* test(integration): keep Prisma clients out of the shared integration conftest import path

---------

Co-authored-by: yuneng <yuneng@berri.ai>
This commit is contained in:
devin-ai-integration[bot] 2026-10-11 08:28:42 -07:00 • committed by GitHub
parent 21cfa0eacd
commit 780c504e00
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
75 changed files with 2705 additions and 2980 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
.circleci/config.yml
View file

@ -1732,13 +1732,6 @@ workflows:
shards: 2
requires: [integration-database-replica]
filters: *integration_default_run
- postgres_suite:
name: proxy-behavior
test_path: tests/proxy_behavior
seed: true
coverage_flag: lens-postgres
timeout_minutes: 25
filters: *integration_default_run
- postgres_suite:
name: proxy-security
test_path: tests/proxy_security_tests

2
codecov.yaml
View file

@ -71,6 +71,8 @@ flag_management:
carryforward: false
- name: lens-python-310
carryforward: false
- name: lens-postgres
carryforward: false
- name: proxy-db-key-generation
carryforward: false
- name: proxy-db-auth-checks

443
tests/integration/_support/actors.py Normal file
View file

@ -0,0 +1,443 @@
import uuid
from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from types import MappingProxyType
from typing import Final
import httpx
import yaml
from prisma import Prisma
from pydantic import JsonValue
from litellm.proxy.utils import PrismaClient, hash_token
from tests.integration._support.client import Gateway, Scenario, string_value
from tests.integration._support.prisma_clients import connected_proxy_prisma
from tests.integration._support.process import owned_proxy
TEAM_ALPHA: Final = "behavior-pin-team-alpha"
TEAM_BETA: Final = "behavior-pin-team-beta"
TEAM_GAMMA: Final = "behavior-pin-team-gamma"
ORG_A: Final = "behavior-pin-org-a"
ORG_B: Final = "behavior-pin-org-b"
class Actor(str, Enum):
PROXY_ADMIN = "proxy_admin"
ORG_ADMIN = "org_admin"
TEAM_ADMIN = "team_admin"
INTERNAL_USER = "internal_user"
OWNER = "owner"
UNRELATED_SAME_ORG = "unrelated_same_org"
CROSS_ORG_USER = "cross_org_user"
SERVICE_ACCOUNT = "service_account"
ORG_B_ADMIN = "org_b_admin"
@dataclass(frozen=True, slots=True)
class SeededKey:
user_id: str
cleartext: str
hashed: str
@dataclass(frozen=True, slots=True)
class ActorSet:
org_a_id: str
org_b_id: str
team_alpha_id: str
team_beta_id: str
team_gamma_id: str
keys: Mapping[Actor, SeededKey]
@dataclass(frozen=True, slots=True)
class ScratchResources:
prefix: str
def tag(self, suffix: str = "") -> str:
return f"{self.prefix}-{suffix}" if suffix else self.prefix
@dataclass(frozen=True, slots=True)
class GatewayClient:
gateway: Gateway
replacements: Mapping[str, str]
async def request(
self,
method: str,
path: str,
*,
headers: Mapping[str, str] | None = None,
json: Mapping[str, JsonValue] | None = None,
params: Mapping[str, str] | None = None,
) -> httpx.Response:
resolved_path: Final = _replace_text(path, self.replacements)
resolved_body: Final = _replace_body(json, self.replacements) if json is not None else None
return self.gateway.request(method, resolved_path, resolved_body, headers=headers, params=params)
async def post(
self,
path: str,
*,
headers: Mapping[str, str] | None = None,
json: Mapping[str, JsonValue] | None = None,
) -> httpx.Response:
return await self.request("POST", path, headers=headers, json=json)
async def get(
self,
path: str,
*,
headers: Mapping[str, str] | None = None,
params: Mapping[str, str] | None = None,
) -> httpx.Response:
return await self.request("GET", path, headers=headers, params=params)
async def put(
self,
path: str,
*,
headers: Mapping[str, str] | None = None,
json: Mapping[str, JsonValue] | None = None,
) -> httpx.Response:
return await self.request("PUT", path, headers=headers, json=json)
async def delete(
self,
path: str,
*,
headers: Mapping[str, str] | None = None,
json: Mapping[str, JsonValue] | None = None,
) -> httpx.Response:
return await self.request("DELETE", path, headers=headers, json=json)
@contextmanager
def owned_settings_client(
client: GatewayClient, directory: Path, litellm_settings: Mapping[str, JsonValue]
) -> Iterator[GatewayClient]:
base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config: Final = {**base, "litellm_settings": {**base["litellm_settings"], **litellm_settings}}
path: Final = directory / "owned-settings.yaml"
path.write_text(yaml.safe_dump(config))
with owned_proxy(client.gateway, directory, {}, config=path) as candidate:
yield GatewayClient(candidate, client.replacements)
@dataclass(frozen=True, slots=True)
class ManagementDatabase:
prisma: PrismaClient
gateway: Gateway
@property
def db(self) -> Prisma:
return self.prisma.db
def _replace_text(value: str, replacements: Mapping[str, str]) -> str:
return value.replace(TEAM_ALPHA, replacements[TEAM_ALPHA]).replace(
TEAM_BETA, replacements[TEAM_BETA]
).replace(TEAM_GAMMA, replacements[TEAM_GAMMA]).replace(ORG_A, replacements[ORG_A]).replace(
ORG_B, replacements[ORG_B]
)
def _replace_body(
value: JsonValue, replacements: Mapping[str, str]
) -> JsonValue:
if isinstance(value, str):
return _replace_text(value, replacements)
if isinstance(value, list):
return [_replace_body(item, replacements) for item in value]
if isinstance(value, dict):
return {key: _replace_body(item, replacements) for key, item in value.items()}
return value
def _seeded_key(
scenario: Scenario,
actor: Actor,
user_id: str,
*,
team_id: str | None = None,
organization_id: str | None = None,
) -> SeededKey:
token: Final = scenario.key(
key_alias=f"integration-{actor.value}-{uuid.uuid4().hex}",
user_id=user_id,
team_id=team_id,
organization_id=organization_id,
)
return SeededKey(user_id=user_id, cleartext=token, hashed=hash_token(token))
async def create_actor_set(
gateway: Gateway, scenario: Scenario, db: ManagementDatabase
) -> ActorSet:
run_id: Final = uuid.uuid4().hex
org_a_id: Final = scenario.organization(organization_alias=f"alpha-{run_id}")
org_b_id: Final = scenario.organization(organization_alias=f"beta-{run_id}")
proxy_admin_id: Final = scenario.user(user_role="proxy_admin")
org_admin_a_id: Final = scenario.org_member(org_a_id, "org_admin")
org_admin_b_id: Final = scenario.org_member(org_b_id, "org_admin")
await db.db.litellm_usertable.update(
where={"user_id": org_admin_a_id}, data={"user_role": "org_admin"}
)
await db.db.litellm_usertable.update(
where={"user_id": org_admin_b_id}, data={"user_role": "org_admin"}
)
team_alpha_id: Final = scenario.team(
team_id=f"integration-{run_id}-alpha",
team_alias=f"alpha-{run_id}",
organization_id=org_a_id,
team_member_permissions=["/key/health"],
)
team_beta_id: Final = scenario.team(
team_id=f"integration-{run_id}-beta",
team_alias=f"beta-{run_id}",
organization_id=org_b_id,
team_member_permissions=["/key/health"],
)
team_gamma_id: Final = scenario.team(
team_id=f"integration-{run_id}-gamma",
team_alias=f"gamma-{run_id}",
organization_id=org_a_id,
team_member_permissions=["/key/health"],
)
team_admin_id: Final = scenario.member(team_alpha_id, role="admin")
internal_user_id: Final = scenario.member(team_alpha_id)
owner_id: Final = scenario.member(team_alpha_id)
unrelated_id: Final = scenario.member(team_alpha_id)
service_account_id: Final = scenario.member(team_alpha_id)
cross_org_id: Final = scenario.member(team_beta_id)
keys: Final[Mapping[Actor, SeededKey]] = MappingProxyType(
{
Actor.PROXY_ADMIN: _seeded_key(scenario, Actor.PROXY_ADMIN, proxy_admin_id),
Actor.ORG_ADMIN: _seeded_key(
scenario, Actor.ORG_ADMIN, org_admin_a_id, organization_id=org_a_id
),
Actor.TEAM_ADMIN: _seeded_key(
scenario, Actor.TEAM_ADMIN, team_admin_id, team_id=team_alpha_id, organization_id=org_a_id
),
Actor.INTERNAL_USER: _seeded_key(
scenario, Actor.INTERNAL_USER, internal_user_id, team_id=team_alpha_id, organization_id=org_a_id
),
Actor.OWNER: _seeded_key(
scenario, Actor.OWNER, owner_id, team_id=team_alpha_id, organization_id=org_a_id
),
Actor.UNRELATED_SAME_ORG: _seeded_key(
scenario, Actor.UNRELATED_SAME_ORG, unrelated_id, team_id=team_alpha_id, organization_id=org_a_id
),
Actor.CROSS_ORG_USER: _seeded_key(
scenario, Actor.CROSS_ORG_USER, cross_org_id, team_id=team_beta_id, organization_id=org_b_id
),
Actor.SERVICE_ACCOUNT: _seeded_key(
scenario, Actor.SERVICE_ACCOUNT, service_account_id, team_id=team_alpha_id, organization_id=org_a_id
),
Actor.ORG_B_ADMIN: _seeded_key(
scenario, Actor.ORG_B_ADMIN, org_admin_b_id, organization_id=org_b_id
),
}
)
return ActorSet(org_a_id, org_b_id, team_alpha_id, team_beta_id, team_gamma_id, keys)
def create_gateway_client(gateway: Gateway, actors: ActorSet) -> GatewayClient:
replacements: Final[Mapping[str, str]] = MappingProxyType(
{
TEAM_ALPHA: actors.team_alpha_id,
TEAM_BETA: actors.team_beta_id,
TEAM_GAMMA: actors.team_gamma_id,
ORG_A: actors.org_a_id,
ORG_B: actors.org_b_id,
}
)
return GatewayClient(gateway, replacements)
async def create_scratch_key(
gateway_client: GatewayClient,
seeder_cleartext: str,
scratch_prefix: str,
*,
user_id: str,
team_id: str | None = None,
organization_id: str | None = None,
key_alias: str | None = None,
) -> str:
body: Final[dict[str, JsonValue]] = {
"key_alias": key_alias or scratch_prefix,
"user_id": user_id,
**({"team_id": team_id} if team_id is not None else {}),
**({"organization_id": organization_id} if organization_id is not None else {}),
}
created: Final = gateway_client.gateway.post(
"/key/generate", body, key=seeder_cleartext
)
return string_value(created["key"])
async def create_scratch_team(
db: ManagementDatabase,
team_id: str,
*,
organization_id: str | None = None,
admin_user_ids: Sequence[str] = (),
member_user_ids: Sequence[str] = (),
team_member_permissions: Sequence[str] | None = None,
models: Sequence[str] | None = None,
max_budget: float | None = None,
tpm_limit: int | None = None,
rpm_limit: int | None = None,
metadata: Mapping[str, JsonValue] | None = None,
) -> str:
admins: Final = list(admin_user_ids)
members: Final = list(member_user_ids)
members_with_roles: Final = [
{"user_id": identity, "role": "admin"} for identity in admins
] + [{"user_id": identity, "role": "user"} for identity in members]
fields: Final[dict[str, JsonValue]] = {
"team_id": team_id,
"team_alias": team_id,
"admins": admins,
"members": admins + members,
"members_with_roles": members_with_roles,
**({"organization_id": organization_id} if organization_id is not None else {}),
**(
{"team_member_permissions": list(team_member_permissions)}
if team_member_permissions is not None
else {}
),
**({"models": list(models)} if models is not None else {}),
**({"max_budget": max_budget} if max_budget is not None else {}),
**({"tpm_limit": tpm_limit} if tpm_limit is not None else {}),
**({"rpm_limit": rpm_limit} if rpm_limit is not None else {}),
**({"metadata": dict(metadata)} if metadata is not None else {}),
}
created: Final = db.gateway.post("/team/new", fields)
return string_value(created["team_id"])
async def create_scratch_user(
db: ManagementDatabase,
scratch_prefix: str,
*,
suffix: str,
user_email: str | None = None,
) -> str:
user_id: Final = f"{scratch_prefix}-{suffix}"
fields: Final[dict[str, JsonValue]] = {
"user_id": user_id,
"user_role": "internal_user",
"auto_create_key": False,
**({"user_email": user_email} if user_email is not None else {}),
}
created: Final = db.gateway.post("/user/new", fields)
return string_value(created["user_id"])
async def create_scratch_org(
db: ManagementDatabase,
scratch_prefix: str,
*,
max_budget: float | None = None,
tpm_limit: int | None = None,
rpm_limit: int | None = None,
models: Sequence[str] | None = None,
metadata: Mapping[str, JsonValue] | None = None,
suffix: str = "org",
) -> str:
organization_id: Final = f"{scratch_prefix}-{suffix}"
budget_fields: Final[dict[str, JsonValue]] = {
"budget_id": f"{organization_id}-budget",
**({"max_budget": max_budget} if max_budget is not None else {}),
**({"tpm_limit": tpm_limit} if tpm_limit is not None else {}),
**({"rpm_limit": rpm_limit} if rpm_limit is not None else {}),
}
budget: Final = db.gateway.post("/budget/new", budget_fields)
budget_id: Final = string_value(budget["budget_id"])
fields: Final[dict[str, JsonValue]] = {
"organization_id": organization_id,
"organization_alias": organization_id,
"budget_id": budget_id,
**({"models": list(models)} if models is not None else {}),
**({"metadata": dict(metadata)} if metadata is not None else {}),
}
created: Final = db.gateway.post("/organization/new", fields)
return string_value(created["organization_id"])
async def create_scratch_actor(
db: ManagementDatabase,
scratch_prefix: str,
*,
user_role: str,
org_admin_of: Sequence[str] = (),
organization_id: str | None = None,
suffix: str = "actor",
) -> SeededKey:
api_user_role: Final = (
user_role
if user_role in {"proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer"}
else "internal_user"
)
created_user: Final = db.gateway.post(
"/user/new",
{
"user_id": f"{scratch_prefix}-{suffix}",
"user_role": api_user_role,
"auto_create_key": False,
**({"organization_id": organization_id} if organization_id is not None else {}),
},
)
user_id: Final = string_value(created_user["user_id"])
token: Final = string_value(
db.gateway.post(
"/key/generate",
{"user_id": user_id, "key_alias": f"{scratch_prefix}-{suffix}-key"},
)["key"]
)
tuple(
db.gateway.post(
"/organization/member_add",
{"organization_id": org_id, "member": {"user_id": user_id, "role": "org_admin"}},
)
for org_id in org_admin_of
)
return SeededKey(user_id, token, hash_token(token))
@asynccontextmanager
async def management_database(gateway: Gateway) -> AsyncIterator[ManagementDatabase]:
async with connected_proxy_prisma() as client:
yield ManagementDatabase(client, gateway)
@asynccontextmanager
async def scratch_resources(db: ManagementDatabase) -> AsyncIterator[ScratchResources]:
resource: Final = ScratchResources(f"integration-{uuid.uuid4().hex}")
yield resource
prefix: Final = resource.prefix
await db.db.litellm_verificationtoken.delete_many(
where={"OR": [{"key_alias": {"startswith": prefix}}, {"key_name": {"startswith": prefix}}]}
)
await db.db.litellm_teammembership.delete_many(where={"team_id": {"startswith": prefix}})
await db.db.litellm_organizationmembership.delete_many(where={"user_id": {"startswith": prefix}})
await db.db.litellm_teamtable.delete_many(where={"team_id": {"startswith": prefix}})
await db.db.litellm_usertable.delete_many(
where={"OR": [{"user_id": {"startswith": prefix}}, {"user_email": {"startswith": prefix}}]}
)
await db.db.litellm_organizationtable.delete_many(where={"organization_id": {"startswith": prefix}})
await db.db.litellm_budgettable.delete_many(where={"budget_id": {"startswith": prefix}})
users: Final = await db.db.litellm_usertable.find_many(where={"teams": {"isEmpty": False}})
for user in users:
teams: Final = [team_id for team_id in user.teams if not team_id.startswith(prefix)]
if teams != list(user.teams):
await db.db.litellm_usertable.update(
where={"user_id": user.user_id},
data={"teams": {"set": teams}},
)

35
tests/integration/_support/prisma_clients.py Normal file
View file

@ -0,0 +1,35 @@
import os
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Final
from prisma import Prisma
from prisma.types import DatasourceOverride
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.utils import PrismaClient, ProxyLogging
@asynccontextmanager
async def connected_prisma(database_url: str | None = None) -> AsyncIterator[Prisma]:
client: Final = Prisma(
datasource=DatasourceOverride(url=database_url)
) if database_url is not None else Prisma()
await client.connect()
try:
yield client
finally:
await client.disconnect()
@asynccontextmanager
async def connected_proxy_prisma(database_url: str | None = None) -> AsyncIterator[PrismaClient]:
client: Final = PrismaClient(
database_url=database_url or os.environ["DATABASE_URL"],
proxy_logging_obj=ProxyLogging(UserApiKeyCache()),
)
await client.connect()
try:
yield client
finally:
await client.disconnect()

0
tests/proxy_behavior/__init__.py → tests/integration/authorization/__init__.py
View file

10
tests/proxy_behavior/auth/test_auth_object_prefetch.py → tests/integration/authorization/test_auth_object_prefetch.py
View file

@ -3,10 +3,12 @@ team and hand the getters rows they validate. The per-regime round-trip counts a
tests/unit/proxy/auth/test_auth_object_prefetch.py."""
import json
from collections.abc import AsyncIterator
from unittest.mock import AsyncMock, MagicMock
from uuid import uuid4
import pytest
import pytest_asyncio
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy.auth.auth_checks import (
@ -18,10 +20,18 @@ from litellm.proxy.auth.auth_checks import (
from litellm.proxy.auth.auth_object_prefetch import AuthObjectRefs, prefetch_auth_objects
from litellm.proxy.auth.team_grants import team_model_aliases
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.utils import PrismaClient
from tests.integration._support.prisma_clients import connected_proxy_prisma
pytestmark = pytest.mark.asyncio(loop_scope="session")
@pytest_asyncio.fixture(scope="module", loop_scope="session")
async def prisma() -> AsyncIterator[PrismaClient]:
async with connected_proxy_prisma() as client:
yield client
def _frozen_cache() -> UserApiKeyCache:
"""The org entries carry a 5s TTL; a frozen clock keeps a slow first call from expiring them mid-test."""
return UserApiKeyCache(in_memory_cache=InMemoryCache(clock=lambda: 1_000_000.0), redis_cache=None)

59
tests/integration/conftest.py
View file

@ -4,20 +4,21 @@ import hashlib
import json
import os
import re
from collections.abc import Generator, Iterator, Mapping, Sequence
from collections.abc import AsyncIterator, Generator, Iterator, Mapping, Sequence
from contextlib import ExitStack
from copy import copy
from importlib.metadata import version
from pathlib import Path
from typing import Final, Protocol, runtime_checkable
from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable
import httpx
import pytest
import pytest_asyncio
from execnet import XSpec
from pydantic import TypeAdapter
from redis import Redis
from tests.integration._support.client import Gateway, eventually, gateway_from_environment
from tests.integration._support.client import Gateway, Scenario, eventually, gateway_from_environment
from tests.integration._support.generation import LIFECYCLE_SETTINGS
from tests.integration._support.manifest import OWNED_DIRECTORIES
from tests.integration._support.ordering import collected_case, fixture_groups, scheduling_groups
@ -26,6 +27,9 @@ from tests.integration._support.routing import RoutingPlugin
from tests.integration._support.workers import worker_endpoints, worker_gateway
from tests.integration.run import GITHUB_FILES
if TYPE_CHECKING:
from tests.integration._support.actors import ActorSet, GatewayClient, ManagementDatabase, ScratchResources
COLLECTED: Final = pytest.StashKey[tuple[str, ...]]()
INVENTORY: Final = pytest.StashKey[tuple[str, ...]]()
REPORTS: Final = pytest.StashKey[list[pytest.TestReport]]()
@ -293,6 +297,55 @@ def gateway() -> Iterator[Gateway]:
yield value
@pytest.fixture(scope="session")
def session_gateway() -> Iterator[Gateway]:
with gateway_from_environment() as value:
yield value
@pytest.fixture
def scenario(gateway: Gateway) -> Iterator[Scenario]:
value: Final[Scenario]
with gateway.scenario() as value:
yield value
@pytest_asyncio.fixture(scope="session", loop_scope="session")
async def actors(
session_gateway: Gateway, db: ManagementDatabase
) -> AsyncIterator[ActorSet]:
with session_gateway.scenario() as scenario:
from tests.integration._support.actors import create_actor_set
value: Final[ActorSet] = await create_actor_set(session_gateway, scenario, db)
yield value
@pytest.fixture(scope="session")
def gateway_client(session_gateway: Gateway, actors: ActorSet) -> GatewayClient:
from tests.integration._support.actors import create_gateway_client
return create_gateway_client(session_gateway, actors)
@pytest_asyncio.fixture(scope="session", loop_scope="session")
async def db(session_gateway: Gateway) -> AsyncIterator[ManagementDatabase]:
from tests.integration._support.actors import management_database
value: Final[ManagementDatabase]
async with management_database(session_gateway) as value:
yield value
@pytest_asyncio.fixture(scope="function", loop_scope="session")
async def resources(db: ManagementDatabase) -> AsyncIterator[ScratchResources]:
from tests.integration._support.actors import scratch_resources
value: Final[ScratchResources]
async with scratch_resources(db) as value:
yield value
@pytest.fixture(scope="session")
def shared_provider_server() -> Iterator[SharedProvider]:
if os.environ.get("PYTEST_XDIST_WORKER") and os.environ.get("INTEGRATION_WORKER_ISOLATION") != "1":

0
tests/proxy_behavior/auth/__init__.py → tests/integration/management/__init__.py
View file

29
tests/proxy_behavior/management/test_credential_migration_endpoint.py → tests/integration/management/test_credential_migration_endpoint.py
View file

@ -13,16 +13,15 @@ the ``user_api_key_auth`` layer (401) before the endpoint's own admin guard runs
import pytest
from .conftest import MASTER_KEY
pytestmark = pytest.mark.asyncio(loop_scope="session")
async def test_migrate_encryption_check_as_admin_is_read_only(proxy_client):
async def test_migrate_encryption_check_as_admin_is_read_only(gateway_client):
"""GET /credentials/migrate-encryption/check returns a residual report (no writes)."""
resp = await proxy_client.get(
resp = await gateway_client.get(
"/credentials/migrate-encryption/check",
headers={"Authorization": f"Bearer {MASTER_KEY}"},
headers={"Authorization": f"Bearer {gateway_client.gateway.key}"},
)
assert resp.status_code == 200, resp.text
body = resp.json()
@ -30,41 +29,41 @@ async def test_migrate_encryption_check_as_admin_is_read_only(proxy_client):
assert "residual_legacy" in body["report"]
async def test_migrate_encryption_check_requires_admin(proxy_client, scratch):
async def test_migrate_encryption_check_requires_admin(gateway_client, resources):
"""A non-admin key cannot reach the residual scan (auth layer rejects, 401)."""
gen = await proxy_client.post(
gen = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {MASTER_KEY}"},
json={"key_alias": scratch.tag("check"), "user_id": scratch.tag("check-user")},
headers={"Authorization": f"Bearer {gateway_client.gateway.key}"},
json={"key_alias": resources.tag("check"), "user_id": resources.tag("check-user")},
)
assert gen.status_code == 200, gen.text
nonadmin_key = gen.json()["key"]
resp = await proxy_client.get(
resp = await gateway_client.get(
"/credentials/migrate-encryption/check",
headers={"Authorization": f"Bearer {nonadmin_key}"},
)
assert resp.status_code == 401, resp.text
async def test_migrate_encryption_requires_admin(proxy_client, scratch):
async def test_migrate_encryption_requires_admin(gateway_client, resources):
"""A non-admin key cannot trigger the migration (auth layer rejects, 401).
Rejection happens before any write: the admin-only route check fires in
``user_api_key_auth``, ahead of the endpoint body.
"""
gen = await proxy_client.post(
gen = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {MASTER_KEY}"},
headers={"Authorization": f"Bearer {gateway_client.gateway.key}"},
json={
"key_alias": scratch.tag("migrate"),
"user_id": scratch.tag("migrate-user"),
"key_alias": resources.tag("migrate"),
"user_id": resources.tag("migrate-user"),
},
)
assert gen.status_code == 200, gen.text
nonadmin_key = gen.json()["key"]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/credentials/migrate-encryption",
headers={"Authorization": f"Bearer {nonadmin_key}"},
)

199
tests/proxy_behavior/management/test_f7_coverage_closeout.py → tests/integration/management/test_f7_coverage_closeout.py
View file

@ -1,3 +1,4 @@
from typing import Final
"""Phase 4 F7 — coverage gap-closer scenarios picked from the un-covered
ranges left after F1–F6 landed.
@ -27,8 +28,8 @@ from prisma import Json
from litellm.proxy.utils import hash_token
from .actors import TEAM_ALPHA, Actor
from .conftest import create_scratch_org, create_scratch_team
from tests.integration._support.actors import TEAM_ALPHA, Actor
from tests.integration._support.actors import create_scratch_org, create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -43,11 +44,11 @@ pytestmark = pytest.mark.asyncio(loop_scope="session")
async def test_org_team_guaranteed_throughput_model_over_bound_rejected(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
org_id = await create_scratch_org(
prisma,
scratch.prefix,
db,
resources.prefix,
models=["gpt-4"],
metadata={"model_rpm_limit": {"gpt-4": 30}},
)
@ -55,13 +56,13 @@ async def test_org_team_guaranteed_throughput_model_over_bound_rejected(
# forces _check_team_model_specific_limits's `model_specific_rpm_limit`
# accumulator to actually accumulate (covers lines 468–478).
await create_scratch_team(
prisma,
team_id=scratch.tag("sibling"),
db,
team_id=resources.tag("sibling"),
organization_id=org_id,
metadata={"model_rpm_limit": {"gpt-4": 20}},
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
team_id = scratch.tag("new")
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
team_id = resources.tag("new")
body = {
"team_id": team_id,
"team_alias": team_id,
@ -70,7 +71,7 @@ async def test_org_team_guaranteed_throughput_model_over_bound_rejected(
"model_rpm_limit": {"gpt-4": 100},
"rpm_limit_type": "guaranteed_throughput",
}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
@ -78,30 +79,30 @@ async def test_org_team_guaranteed_throughput_model_over_bound_rejected(
# 20 (sibling) + 100 (new) > 30 (org cap) → guard fires.
assert resp.status_code == 400, resp.text
assert "RPM" in resp.text, resp.text
rows = await prisma.db.litellm_teamtable.find_many(where={"team_id": team_id})
rows = await db.db.litellm_teamtable.find_many(where={"team_id": team_id})
assert rows == []
async def test_org_team_guaranteed_throughput_model_tpm_over_bound_rejected(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""Mirror of the rpm scenario but for the model_tpm side — pins
_check_team_model_specific_limits's tpm branch (lines 503–521) which
the rpm scenario doesn't exercise."""
org_id = await create_scratch_org(
prisma,
scratch.prefix,
db,
resources.prefix,
models=["gpt-4"],
metadata={"model_tpm_limit": {"gpt-4": 500}},
)
await create_scratch_team(
prisma,
team_id=scratch.tag("sibling"),
db,
team_id=resources.tag("sibling"),
organization_id=org_id,
metadata={"model_tpm_limit": {"gpt-4": 200}},
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
team_id = scratch.tag("new")
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
team_id = resources.tag("new")
body = {
"team_id": team_id,
"team_alias": team_id,
@ -110,51 +111,48 @@ async def test_org_team_guaranteed_throughput_model_tpm_over_bound_rejected(
"model_tpm_limit": {"gpt-4": 5000},
"tpm_limit_type": "guaranteed_throughput",
}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
)
assert resp.status_code == 400, resp.text
assert "TPM" in resp.text, resp.text
rows = await prisma.db.litellm_teamtable.find_many(where={"team_id": team_id})
rows = await db.db.litellm_teamtable.find_many(where={"team_id": team_id})
assert rows == []
async def test_org_team_guaranteed_throughput_aggregate_runs(
proxy_client, prisma, scratch, world
):
"""Aggregate guard's no-op path (line 549-566 — entity_rpm_limit is None
because include_budget_table=False at the call site). The helper still
executes its `allocated_tpm = sum(...)` and `allocated_rpm = sum(...)`
lines, which is the coverage target."""
org_id = await create_scratch_org(prisma, scratch.prefix, models=["m"])
await create_scratch_team(
prisma,
team_id=scratch.tag("sibling"),
organization_id=org_id,
tpm_limit=100,
rpm_limit=10,
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
team_id = scratch.tag("new")
body = {
"team_id": team_id,
"team_alias": team_id,
"organization_id": org_id,
"models": ["m"],
"tpm_limit": 50,
"rpm_limit": 5,
"tpm_limit_type": "guaranteed_throughput",
}
resp = await proxy_client.post(
"/team/new",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
)
# No org budget table loaded → check is no-op → 200.
assert resp.status_code == 200, resp.text
@pytest.mark.parametrize("requested_tpm,expected_status", [(50, 400), (20, 200)], ids=["over", "within"])
async def test_org_team_guaranteed_throughput_counts_sibling_allocations(
requested_tpm: int, expected_status: int, gateway_client, db, resources, actors
):
org_id = await create_scratch_org(db, resources.prefix, models=["m"], tpm_limit=120)
await create_scratch_team(
db,
team_id=resources.tag("sibling"),
organization_id=org_id,
models=["m"],
tpm_limit=100,
)
team_id = resources.tag("new")
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={
"team_id": team_id,
"team_alias": team_id,
"organization_id": org_id,
"models": ["m"],
"tpm_limit": requested_tpm,
"tpm_limit_type": "guaranteed_throughput",
},
)
assert resp.status_code == expected_status, resp.text
rows = await db.db.litellm_teamtable.find_many(where={"team_id": team_id})
assert len(rows) == (1 if expected_status == 200 else 0)
if expected_status == 400:
assert "Allocated TPM limit=100" in resp.text, resp.text
# ---------------------------------------------------------------------------
# /key/generate — _check_project_key_limits project-not-found branch
@ -164,14 +162,14 @@ async def test_org_team_guaranteed_throughput_aggregate_runs(
async def test_key_generate_with_unknown_project_id_rejected(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
body = {
"key_alias": scratch.prefix,
"project_id": f"{scratch.prefix}-ghost-project",
"key_alias": resources.prefix,
"project_id": f"{resources.prefix}-ghost-project",
}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
@ -180,8 +178,8 @@ async def test_key_generate_with_unknown_project_id_rejected(
# The route's exception handler may wrap it; pin both shapes.
assert resp.status_code in (400, 404), resp.text
assert "project" in resp.text.lower() or "not found" in resp.text.lower(), resp.text
rows = await prisma.db.litellm_verificationtoken.find_many(
where={"key_alias": scratch.prefix}
rows = await db.db.litellm_verificationtoken.find_many(
where={"key_alias": resources.prefix}
)
assert rows == [], "rejected key leaked a row"
@ -195,10 +193,10 @@ async def test_key_generate_with_unknown_project_id_rejected(
async def _seed_key_for_relocation(
prisma, scratch_prefix: str, *, user_id: str, team_id: str
db, scratch_prefix: str, *, user_id: str, team_id: str
) -> str:
cleartext = "sk-" + uuid.uuid4().hex
await prisma.db.litellm_verificationtoken.create(
await db.db.litellm_verificationtoken.create(
data={
"token": hash_token(cleartext),
"key_alias": f"{scratch_prefix}-key",
@ -212,39 +210,39 @@ async def _seed_key_for_relocation(
async def test_team_new_with_team_member_budget_creates_budget_row(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""/team/new with `team_member_budget` routes through
TeamMemberBudgetHandler.create_team_member_budget_table (lines 196–248).
Observable end-state: a litellm_budgettable row is created and the
team's metadata.team_member_budget_id points at it."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
team_id = scratch.tag("team-with-mbudget")
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
team_id = resources.tag("team-with-mbudget")
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {seeder}"},
json={
"team_id": team_id,
"team_alias": scratch.tag("alias"),
"team_alias": resources.tag("alias"),
"team_member_budget": 10.0,
"team_member_rpm_limit": 100,
"team_member_tpm_limit": 1000,
},
)
assert resp.status_code == 200, resp.text
team_row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
team_row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert team_row is not None
budget_id = (team_row.metadata or {}).get("team_member_budget_id")
assert (
budget_id is not None
), f"team_member_budget_id not stamped on team metadata: {team_row.metadata!r}"
# The handler writes the per-member budget row under a non-scratch id
# The handler writes the per-member budget row under a non-resources id
# pattern (`team-<alias>-budget-<uuid>`), so the prefix sweep can't
# reclaim it. Cleanup must run even when a downstream assertion fires;
# otherwise orphan rows accumulate across CI re-runs.
try:
budget_row = await prisma.db.litellm_budgettable.find_unique(
budget_row = await db.db.litellm_budgettable.find_unique(
where={"budget_id": budget_id}
)
assert budget_row is not None, "team_member_budget row was not created"
@ -252,23 +250,19 @@ async def test_team_new_with_team_member_budget_creates_budget_row(
assert budget_row.rpm_limit == 100
assert budget_row.tpm_limit == 1000
finally:
await prisma.db.litellm_teamtable.update(
await db.db.litellm_teamtable.update(
where={"team_id": team_id},
data={"metadata": Json({})},
)
await prisma.db.litellm_budgettable.delete(where={"budget_id": budget_id})
await db.db.litellm_budgettable.delete(where={"budget_id": budget_id})
async def test_team_update_team_member_budget_upserts(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""/team/update with `team_member_budget` against a team that has no
pre-existing team_member_budget_id routes through
TeamMemberBudgetHandler.upsert_team_member_budget_table's else-branch
(lines 294–303), which in turn calls create_team_member_budget_table."""
team_id = await create_scratch_team(prisma, team_id=scratch.tag("team"))
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
team_id = await create_scratch_team(db, team_id=resources.tag("team"))
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {seeder}"},
json={
@ -277,69 +271,70 @@ async def test_team_update_team_member_budget_upserts(
},
)
assert resp.status_code == 200, resp.text
team_row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
team_row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert team_row is not None
budget_id = (team_row.metadata or {}).get("team_member_budget_id")
assert budget_id is not None, "team_member_budget_id not upserted"
budget_row = await db.db.litellm_budgettable.find_unique(where={"budget_id": budget_id})
assert budget_row is not None and budget_row.max_budget == 5.0, budget_row
# See comment on the sibling test — non-prefixed budget row, cleanup
# must survive an assertion failure to avoid orphan accumulation.
try:
# (No further assertions today, but the try/finally keeps the
# cleanup contract uniform with the sibling test and is robust to
# future asserts being added here.)
pass
budget_row: Final = await db.db.litellm_budgettable.find_unique(
where={"budget_id": budget_id}
)
assert budget_row is not None, "team_member_budget row was not created"
assert budget_row.max_budget == 5.0
finally:
await prisma.db.litellm_teamtable.update(
await db.db.litellm_teamtable.update(
where={"team_id": team_id},
data={"metadata": Json({})},
)
await prisma.db.litellm_budgettable.delete(where={"budget_id": budget_id})
await db.db.litellm_budgettable.delete(where={"budget_id": budget_id})
async def test_key_team_change_accepted_by_target_team_admin(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""Caller is admin of the destination team AND the key's owner — the
common_key_access_checks gate requires user_id-match for non-proxy-admin
callers, so the team-admin-accepts branch of validate_key_team_change
can only be reached when the team admin is also the key holder.
Hits validate_key_team_change line 3007–3011."""
actor_user_id = f"{scratch.prefix}-self-admin"
actor_user_id = f"{resources.prefix}-self-admin"
actor_cleartext = "sk-" + uuid.uuid4().hex
await prisma.db.litellm_usertable.create(
await db.db.litellm_usertable.create(
data={"user_id": actor_user_id, "user_role": "internal_user"}
)
await prisma.db.litellm_verificationtoken.create(
await db.db.litellm_verificationtoken.create(
data={
"token": hash_token(actor_cleartext),
"key_alias": f"{scratch.prefix}-actor-key",
"key_name": f"{scratch.prefix}-actor-key",
"key_alias": f"{resources.prefix}-actor-key",
"key_name": f"{resources.prefix}-actor-key",
"user_id": actor_user_id,
"models": [],
"allowed_routes": ["/key/update"],
}
)
source_team = await create_scratch_team(
prisma,
team_id=scratch.tag("source"),
db,
team_id=resources.tag("source"),
admin_user_ids=[actor_user_id],
)
target_team = await create_scratch_team(
prisma,
team_id=scratch.tag("target"),
db,
team_id=resources.tag("target"),
admin_user_ids=[actor_user_id],
)
key_cleartext = await _seed_key_for_relocation(
prisma, scratch.prefix, user_id=actor_user_id, team_id=source_team
db, resources.prefix, user_id=actor_user_id, team_id=source_team
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {actor_cleartext}"},
json={"key": key_cleartext, "team_id": target_team},
)
assert resp.status_code == 200, resp.text
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(key_cleartext)}
)
assert row is not None

550
tests/proxy_behavior/management/test_f7_key_coverage_push.py → tests/integration/management/test_f7_key_coverage_push.py
View file

@ -1,38 +1,20 @@
"""Phase 4 F7-extension — additional payload pins pushing
`key_management_endpoints.py` past the 70 % stretch.
Same pattern as `test_f7_coverage_closeout.py`: each scenario cites the
file:line range it pins and asserts observable end-state, not response-body
snapshots. Targets the largest non-deferred un-covered ranges left after
the first F7 pass:
* 5942–6063 /key/health logging-metadata path (test_key_logging body)
* 4630–4692 /key/reset_spend happy path + 404 + admin gate
* 4565–4596 _validate_reset_spend_value branches
* 4421–4476 /key/regenerate ghost-key 404 + premium gate
* 6118–6133 _enforce_unique_key_alias duplicate-alias rejection
* 6148–6169 validate_model_max_budget malformed payload rejection
* 4708–4789 validate_key_list_check user/team/org/key_hash branches
Excluded: `_rotate_master_key` (lines 3997–4123) — deferred per plan §6.
"""
import uuid
from typing import Any, Dict, Optional
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, Final, Optional
import pytest
from prisma import Json
from litellm.proxy.utils import hash_token
from .actors import TEAM_ALPHA, TEAM_BETA, Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, Actor
from tests.integration._support.actors import create_scratch_key, create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
async def _seed_token(
prisma,
db,
scratch_prefix: str,
*,
suffix: str = "tok",
@ -57,7 +39,7 @@ async def _seed_token(
data["metadata"] = Json(metadata)
if team_id is not None:
data["team_id"] = team_id
await prisma.db.litellm_verificationtoken.create(data=data)
await db.db.litellm_verificationtoken.create(data=data)
return cleartext
@ -71,19 +53,19 @@ async def _seed_token(
async def test_key_health_with_logging_metadata_runs_test_logging(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.INTERNAL_USER].user_id,
metadata={
"logging": [
{"callback_name": "noop-scratch-callback"},
{"callback_name": "noop-resources-callback"},
]
},
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/health",
headers={"Authorization": f"Bearer {cleartext}"},
)
@ -92,21 +74,21 @@ async def test_key_health_with_logging_metadata_runs_test_logging(
# Either healthy or unhealthy — both pin the logging branch.
assert body["key"] in ("healthy", "unhealthy")
assert "logging_callbacks" in body
assert body["logging_callbacks"]["callbacks"] == ["noop-scratch-callback"]
assert body["logging_callbacks"]["callbacks"] == ["noop-resources-callback"]
async def test_key_health_with_missing_callback_name_rejected(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""test_key_logging raises ValueError if a callback dict lacks
callback_name — wrapped by the outer try/except into a 500."""
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
metadata={"logging": [{"not_callback_name": "x"}]},
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/health",
headers={"Authorization": f"Bearer {cleartext}"},
)
@ -124,25 +106,25 @@ async def test_key_health_with_missing_callback_name_rejected(
# ---------------------------------------------------------------------------
async def test_reset_spend_ghost_key_404(proxy_client, scratch, world):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
f"/key/sk-{scratch.prefix}-ghost/reset_spend",
async def test_reset_spend_ghost_key_404(gateway_client, resources, actors):
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
f"/key/sk-{resources.prefix}-ghost/reset_spend",
headers={"Authorization": f"Bearer {seeder}"},
json={"reset_to": 0.0},
)
assert resp.status_code == 404, resp.text
async def test_reset_spend_happy_path(proxy_client, prisma, scratch, world):
async def test_reset_spend_happy_path(gateway_client, db, resources, actors):
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
spend=5.0,
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
f"/key/{cleartext}/reset_spend",
headers={"Authorization": f"Bearer {seeder}"},
json={"reset_to": 1.0},
@ -151,7 +133,7 @@ async def test_reset_spend_happy_path(proxy_client, prisma, scratch, world):
assert resp.json()["spend"] == 1.0
assert resp.json()["previous_spend"] == 5.0
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(cleartext)}
)
assert row is not None
@ -171,16 +153,16 @@ async def test_reset_spend_happy_path(proxy_client, prisma, scratch, world):
ids=["negative", "above_current_spend"],
)
async def test_reset_spend_validate_value_branches(
reset_to, expected_detail: str, proxy_client, prisma, scratch, world
reset_to, expected_detail: str, gateway_client, db, resources, actors
):
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
spend=5.0,
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
f"/key/{cleartext}/reset_spend",
headers={"Authorization": f"Bearer {seeder}"},
json={"reset_to": reset_to},
@ -188,50 +170,28 @@ async def test_reset_spend_validate_value_branches(
assert resp.status_code == 400, resp.text
assert expected_detail in resp.text, resp.text
# Row spend must be unchanged.
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(cleartext)}
)
assert row.spend == 5.0, "spend mutated despite validation rejection"
async def test_reset_spend_non_numeric_caught_by_pydantic(
proxy_client, prisma, scratch, world
):
"""Pin: non-float reset_to is rejected at the Pydantic layer with 422
before reaching _validate_reset_spend_value. This documents that the
helper's `isinstance(reset_to, (int, float))` guard at line 4568 is
structurally unreachable via HTTP."""
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
spend=5.0,
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
f"/key/{cleartext}/reset_spend",
headers={"Authorization": f"Bearer {seeder}"},
json={"reset_to": "not-a-number"},
)
assert resp.status_code == 422, resp.text
async def test_reset_spend_non_admin_caller_rejected(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""_check_proxy_or_team_admin_for_key raises 403 when caller is neither
proxy admin nor admin of the key's team."""
# Seed a key against TEAM_BETA (no actor in our world is admin of beta).
# Seed a key against actors.team_beta_id (no actor in our actors is admin of beta).
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.CROSS_ORG_USER].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.CROSS_ORG_USER].user_id,
spend=5.0,
team_id=TEAM_BETA,
team_id=actors.team_beta_id,
)
# Caller is a member of TEAM_ALPHA but not admin of TEAM_BETA.
initiator = world.keys[Actor.INTERNAL_USER].cleartext
resp = await proxy_client.post(
# Caller is a member of actors.team_alpha_id but not admin of actors.team_beta_id.
initiator = actors.keys[Actor.INTERNAL_USER].cleartext
resp = await gateway_client.post(
f"/key/{cleartext}/reset_spend",
headers={"Authorization": f"Bearer {initiator}"},
json={"reset_to": 1.0},
@ -239,7 +199,7 @@ async def test_reset_spend_non_admin_caller_rejected(
# Route-level admin gate may fire first (401) or the helper's own 403 —
# both prove the path is guarded; pin either as rejection.
assert resp.status_code in (401, 403), resp.text
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(cleartext)}
)
assert row.spend == 5.0, "spend mutated despite rejection"
@ -252,37 +212,38 @@ async def test_reset_spend_non_admin_caller_rejected(
# ---------------------------------------------------------------------------
async def test_regenerate_ghost_key_404(proxy_client, scratch, world):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
f"/key/sk-{scratch.prefix}-ghost/regenerate",
async def test_regenerate_ghost_key_404(gateway_client, resources, actors):
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
f"/key/sk-{resources.prefix}-ghost/regenerate",
headers={"Authorization": f"Bearer {seeder}"},
json={},
)
assert resp.status_code == 404, resp.text
async def test_regenerate_no_key_supplied_400(proxy_client, world):
"""POST /key/regenerate with no key in path AND no key in body."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
async def test_regenerate_no_key_supplied_400(gateway_client, actors):
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/key/regenerate",
headers={"Authorization": f"Bearer {seeder}"},
json={},
)
assert resp.status_code == 400, resp.text
assert "No key passed in" in resp.text or "key" in resp.text.lower()
error = resp.json()["error"]
assert error["code"] == "400", resp.text
assert "No key passed in." in error["message"], resp.text
async def test_regenerate_happy_path(proxy_client, prisma, scratch, world):
async def test_regenerate_happy_path(gateway_client, db, resources, actors):
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
)
old_hash = hash_token(cleartext)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
f"/key/{cleartext}/regenerate",
headers={"Authorization": f"Bearer {seeder}"},
json={},
@ -294,35 +255,35 @@ async def test_regenerate_happy_path(proxy_client, prisma, scratch, world):
assert new_hash != old_hash
# Old token should no longer be in active tokens (deleted by regenerate).
old_row = await prisma.db.litellm_verificationtoken.find_unique(
old_row = await db.db.litellm_verificationtoken.find_unique(
where={"token": old_hash}
)
assert old_row is None, "old token still present after regenerate"
# New token should be active.
new_row = await prisma.db.litellm_verificationtoken.find_unique(
new_row = await db.db.litellm_verificationtoken.find_unique(
where={"token": new_hash}
)
assert new_row is not None, "new token not written after regenerate"
async def test_regenerate_with_explicit_new_key(proxy_client, prisma, scratch, world):
async def test_regenerate_with_explicit_new_key(gateway_client, db, resources, actors):
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
suffix="explicit",
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
new_key = "sk-" + uuid.uuid4().hex
resp = await proxy_client.post(
resp = await gateway_client.post(
f"/key/{cleartext}/regenerate",
headers={"Authorization": f"Bearer {seeder}"},
json={"new_key": new_key},
)
assert resp.status_code == 200, resp.text
assert resp.json()["key"] == new_key
new_row = await prisma.db.litellm_verificationtoken.find_unique(
new_row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(new_key)}
)
assert new_row is not None
@ -334,16 +295,16 @@ async def test_regenerate_with_explicit_new_key(proxy_client, prisma, scratch, w
# ---------------------------------------------------------------------------
async def test_generate_duplicate_alias_rejected(proxy_client, prisma, scratch, world):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
alias = scratch.prefix + "-shared-alias"
first = await proxy_client.post(
async def test_generate_duplicate_alias_rejected(gateway_client, db, resources, actors):
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
alias = resources.prefix + "-shared-alias"
first = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json={"key_alias": alias},
)
assert first.status_code == 200, first.text
second = await proxy_client.post(
second = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json={"key_alias": alias},
@ -359,14 +320,14 @@ async def test_generate_duplicate_alias_rejected(proxy_client, prisma, scratch,
async def test_generate_with_invalid_model_max_budget_rejected(
proxy_client, scratch, world
gateway_client, resources, actors
):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json={
"key_alias": scratch.prefix,
"key_alias": resources.prefix,
# Wrong shape — budget_limit should be numeric; passing a dict
# for the model value bypasses the BudgetConfig parse and trips
# the validator's exception wrap.
@ -401,14 +362,14 @@ async def test_generate_with_invalid_model_max_budget_rejected(
ids=["user_mismatch", "team_mismatch", "org_mismatch"],
)
async def test_key_list_non_admin_authz_branches(
filter_kwarg, expected_substring: str, proxy_client, world
filter_kwarg, expected_substring: str, gateway_client, actors
):
"""Non-admin caller hits validate_key_list_check's three rejection
branches. INTERNAL_USER (Org A, TEAM_ALPHA member) is the caller; each
branches. INTERNAL_USER (Org A, actors.team_alpha_id member) is the caller; each
filter targets a foreign user/team/org and trips the matching guard."""
caller = world.keys[Actor.INTERNAL_USER].cleartext
caller = actors.keys[Actor.INTERNAL_USER].cleartext
qs = "&".join(f"{k}={v}" for k, v in filter_kwarg.items())
resp = await proxy_client.get(
resp = await gateway_client.get(
f"/key/list?{qs}",
headers={"Authorization": f"Bearer {caller}"},
)
@ -424,22 +385,22 @@ async def test_key_list_non_admin_authz_branches(
async def test_key_bulk_update_mixed_success_and_failure(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
existing = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
suffix="bulk1",
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={
"keys": [
{"key": existing, "max_budget": 5.0},
{"key": "sk-ghost-" + scratch.prefix, "max_budget": 5.0},
{"key": "sk-ghost-" + resources.prefix, "max_budget": 5.0},
]
},
)
@ -449,16 +410,16 @@ async def test_key_bulk_update_mixed_success_and_failure(
assert len(body["successful_updates"]) == 1
assert len(body["failed_updates"]) == 1
# Re-read confirms the successful key was actually updated.
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(existing)}
)
assert row.max_budget == 5.0
async def test_key_bulk_update_non_admin_rejected(proxy_client, world):
async def test_key_bulk_update_non_admin_rejected(gateway_client, actors):
"""Lines 2631–2635 — admin-only role gate."""
caller = world.keys[Actor.INTERNAL_USER].cleartext
resp = await proxy_client.post(
caller = actors.keys[Actor.INTERNAL_USER].cleartext
resp = await gateway_client.post(
"/key/bulk_update",
headers={"Authorization": f"Bearer {caller}"},
json={"keys": [{"key": "sk-x", "max_budget": 1.0}]},
@ -466,10 +427,10 @@ async def test_key_bulk_update_non_admin_rejected(proxy_client, world):
assert resp.status_code in (401, 403), resp.text
async def test_key_bulk_update_empty_keys_rejected(proxy_client, world):
async def test_key_bulk_update_empty_keys_rejected(gateway_client, actors):
"""Line 2643–2647 — empty keys list rejected."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/key/bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={"keys": []},
@ -478,12 +439,12 @@ async def test_key_bulk_update_empty_keys_rejected(proxy_client, world):
assert "No keys" in resp.text
async def test_key_bulk_update_exceeds_max_batch_rejected(proxy_client, world):
async def test_key_bulk_update_exceeds_max_batch_rejected(gateway_client, actors):
"""Lines 2649–2656 — over-batch-size rejection. 501 ghost keys are fine
here because validation fires before any update runs."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
over_batch = [{"key": f"sk-{i}", "max_budget": 1.0} for i in range(501)]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={"keys": over_batch},
@ -501,15 +462,15 @@ async def test_key_bulk_update_exceeds_max_batch_rejected(proxy_client, world):
async def test_key_update_with_duration_and_budget_duration(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {seeder}"},
json={
@ -519,35 +480,48 @@ async def test_key_update_with_duration_and_budget_duration(
},
)
assert resp.status_code == 200, resp.text
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(cleartext)}
)
assert row.expires is not None, "expires not stamped from duration"
assert row.budget_reset_at is not None, "budget_reset_at not stamped"
async def test_key_update_with_clear_duration(proxy_client, prisma, scratch, world):
"""`duration: -1` clears expires (line 1796–1798)."""
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
assert row is not None
assert row.budget_duration == "30d"
now = datetime.now(timezone.utc)
assert timedelta(minutes=55) < row.expires - now <= timedelta(hours=1), row.expires
assert now < row.budget_reset_at <= now + timedelta(days=31), row.budget_reset_at
info = await gateway_client.get(
f"/key/info?key={hash_token(cleartext)}",
headers={"Authorization": f"Bearer {seeder}"},
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
# First set an expiry
await proxy_client.post(
assert info.status_code == 200, info.text
assert info.json()["info"]["expires"] == row.expires.isoformat()
assert info.json()["info"]["budget_reset_at"] == row.budget_reset_at.isoformat()
async def test_key_update_with_clear_duration(gateway_client, db, resources, actors):
cleartext = await _seed_token(
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
)
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
expiring = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {seeder}"},
json={"key": cleartext, "duration": "1h"},
)
# Then clear it
resp = await proxy_client.post(
assert expiring.status_code == 200, expiring.text
expiring_row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(cleartext)}
)
assert expiring_row.expires is not None, "duration=1h did not set expires"
resp = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {seeder}"},
json={"key": cleartext, "duration": "-1"},
)
assert resp.status_code == 200, resp.text
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(cleartext)}
)
assert row.expires is None, "expires not cleared by duration=-1"
@ -564,27 +538,15 @@ async def test_key_update_with_clear_duration(proxy_client, prisma, scratch, wor
# ---------------------------------------------------------------------------
async def test_team_key_bulk_update_missing_team_id_rejected(proxy_client, world):
"""Pydantic catches missing team_id at the request layer → 422 before
the handler's own `if not data.team_id` guard at line 2803."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
"/team/key/bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={"update_fields": {"max_budget": 5.0}},
)
assert resp.status_code == 422, resp.text
async def test_team_key_bulk_update_no_selector_rejected(proxy_client, world):
async def test_team_key_bulk_update_no_selector_rejected(gateway_client, actors):
"""Pydantic root-validator catches missing key_ids/all_keys_in_team
at the request layer → 422."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/key/bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={
"team_id": TEAM_ALPHA,
"team_id": actors.team_alpha_id,
"update_fields": {"max_budget": 5.0},
},
)
@ -593,31 +555,31 @@ async def test_team_key_bulk_update_no_selector_rejected(proxy_client, world):
async def test_team_key_bulk_update_all_keys_in_team(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""Pin the all_keys_in_team branch (lines 2818–2841) plus the
per-key success path. Seed two scratch keys against the scratch team."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
per-key success path. Seed two resources keys against the resources team."""
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
team_id = await create_scratch_team(
prisma,
team_id=scratch.tag("team"),
admin_user_ids=[world.keys[Actor.PROXY_ADMIN].user_id],
db,
team_id=resources.tag("team"),
admin_user_ids=[actors.keys[Actor.PROXY_ADMIN].user_id],
)
k1 = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
suffix="t1",
team_id=team_id,
)
k2 = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
suffix="t2",
team_id=team_id,
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/key/bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={
@ -631,36 +593,36 @@ async def test_team_key_bulk_update_all_keys_in_team(
assert body["total_requested"] == 2
# Re-read both keys; max_budget should be set.
for k in (k1, k2):
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(k)}
)
assert row.max_budget == 7.0
async def test_team_key_bulk_update_explicit_key_ids_mixed(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""Pin the explicit-key_ids dedupe + per-key 404 path (lines 2842–2944).
Send a real key + a ghost key under the same team_id."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
team_id = await create_scratch_team(
prisma,
team_id=scratch.tag("team"),
admin_user_ids=[world.keys[Actor.PROXY_ADMIN].user_id],
db,
team_id=resources.tag("team"),
admin_user_ids=[actors.keys[Actor.PROXY_ADMIN].user_id],
)
real_key = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
suffix="real",
team_id=team_id,
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/key/bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={
"team_id": team_id,
"key_ids": [real_key, real_key, "sk-ghost-" + scratch.prefix],
"key_ids": [real_key, real_key, "sk-ghost-" + resources.prefix],
"update_fields": {"max_budget": 9.0},
},
)
@ -686,24 +648,24 @@ async def test_team_key_bulk_update_explicit_key_ids_mixed(
# ---------------------------------------------------------------------------
async def test_key_info_ghost_key_404(proxy_client, scratch, world):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.get(
f"/key/info?key=sk-{scratch.prefix}-ghost",
async def test_key_info_ghost_key_404(gateway_client, resources, actors):
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.get(
f"/key/info?key=sk-{resources.prefix}-ghost",
headers={"Authorization": f"Bearer {seeder}"},
)
assert resp.status_code == 404, resp.text
async def test_key_info_explicit_existing_key(proxy_client, prisma, scratch, world):
async def test_key_info_explicit_existing_key(gateway_client, db, resources, actors):
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
suffix="info-test",
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.get(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.get(
f"/key/info?key={cleartext}",
headers={"Authorization": f"Bearer {seeder}"},
)
@ -715,15 +677,18 @@ async def test_key_info_explicit_existing_key(proxy_client, prisma, scratch, wor
assert "token" not in body["info"]
async def test_key_info_no_key_uses_auth_header(proxy_client, world):
"""Pin line 3260: `key = key or user_api_key_dict.api_key` — caller's
own key info is returned when no `?key=` is supplied."""
caller = world.keys[Actor.INTERNAL_USER].cleartext
resp = await proxy_client.get(
async def test_key_info_no_key_uses_auth_header(gateway_client, actors):
caller = actors.keys[Actor.INTERNAL_USER].cleartext
resp = await gateway_client.get(
"/key/info",
headers={"Authorization": f"Bearer {caller}"},
)
assert resp.status_code == 200, resp.text
body: Final = resp.json()
assert body["key"] == hash_token(caller)
assert body["info"]["user_id"] == actors.keys[Actor.INTERNAL_USER].user_id
assert body["info"]["status"] == "active"
assert "token" not in body["info"]
# ---------------------------------------------------------------------------
@ -732,13 +697,13 @@ async def test_key_info_no_key_uses_auth_header(proxy_client, world):
# ---------------------------------------------------------------------------
async def test_key_generate_with_budget_limits(proxy_client, prisma, scratch, world):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
async def test_key_generate_with_budget_limits(gateway_client, db, resources, actors):
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json={
"key_alias": scratch.prefix,
"key_alias": resources.prefix,
"budget_limits": [
{"max_budget": 5.0, "budget_duration": "1d"},
{"max_budget": 20.0, "budget_duration": "30d"},
@ -746,17 +711,27 @@ async def test_key_generate_with_budget_limits(proxy_client, prisma, scratch, wo
},
)
assert resp.status_code == 200, resp.text
cleartext: Final = resp.json()["key"]
info: Final = await gateway_client.get(
f"/key/info?key={hash_token(cleartext)}",
headers={"Authorization": f"Bearer {seeder}"},
)
assert info.status_code == 200, info.text
windows: Final = info.json()["info"]["budget_limits"]
assert [
(window["max_budget"], window["budget_duration"]) for window in windows
] == [(5.0, "1d"), (20.0, "30d")]
async def test_key_aliases_proxy_admin_unscoped(proxy_client, prisma, scratch, world):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
alias = scratch.prefix + "-aliases-test"
await proxy_client.post(
async def test_key_aliases_proxy_admin_unscoped(gateway_client, db, resources, actors):
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
alias = resources.prefix + "-aliases-test"
await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json={"key_alias": alias},
)
resp = await proxy_client.get(
resp = await gateway_client.get(
"/key/aliases?size=10",
headers={"Authorization": f"Bearer {seeder}"},
)
@ -771,36 +746,23 @@ async def test_key_aliases_proxy_admin_unscoped(proxy_client, prisma, scratch, w
), f"newly created alias not in list: {body['aliases']}"
async def test_key_aliases_with_search_filter(proxy_client, prisma, scratch, world):
async def test_key_aliases_with_search_filter(gateway_client, db, resources, actors):
"""Pin the `search` ILIKE branch (lines 5166–5168)."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
unique_alias = scratch.prefix + "-search-unique-tag"
await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
unique_alias = resources.prefix + "-search-unique-tag"
await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json={"key_alias": unique_alias},
)
resp = await proxy_client.get(
f"/key/aliases?search={scratch.prefix}-search",
resp = await gateway_client.get(
f"/key/aliases?search={resources.prefix}-search",
headers={"Authorization": f"Bearer {seeder}"},
)
assert resp.status_code == 200, resp.text
assert unique_alias in resp.json()["aliases"]
async def test_key_aliases_non_admin_scoped(proxy_client, world):
"""Non-admin caller routes through _apply_non_admin_alias_scope (line
5161–5164). The exact alias visibility depends on team membership; the
pin is that the call succeeds with a scoped result."""
caller = world.keys[Actor.INTERNAL_USER].cleartext
resp = await proxy_client.get(
"/key/aliases?size=10",
headers={"Authorization": f"Bearer {caller}"},
)
assert resp.status_code == 200, resp.text
assert "aliases" in resp.json()
# ---------------------------------------------------------------------------
# /key/list with extended filters — pins _build_filter_conditions branches
# at lines 5280–5388. Each scenario varies one filter so a different
@ -809,23 +771,51 @@ async def test_key_aliases_non_admin_scoped(proxy_client, world):
@pytest.mark.parametrize(
"query",
"filter_name",
[
"include_created_by_keys=true",
"include_team_keys=true",
"return_full_object=true",
"sort_by=created_at&sort_order=asc",
"key_alias=behavior",
"created_by",
"team_keys",
"full_object",
"sort",
"alias_substring",
],
ids=["created_by", "team_keys", "full_object", "sort", "alias_substring"],
)
async def test_key_list_filter_branches(query: str, proxy_client, world):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.get(
f"/key/list?{query}",
headers={"Authorization": f"Bearer {seeder}"},
async def test_key_list_filter_branches(filter_name: str, gateway_client, resources, actors):
admin: Final = actors.keys[Actor.PROXY_ADMIN]
key_alias: Final = resources.tag("filter-branch")
cleartext: Final = await create_scratch_key(
gateway_client,
admin.cleartext,
resources.prefix,
user_id=admin.user_id,
key_alias=key_alias,
team_id=actors.team_alpha_id,
)
query: Final = {
"created_by": "include_created_by_keys=true",
"team_keys": "include_team_keys=true",
"full_object": "return_full_object=true",
"sort": "sort_by=created_at&sort_order=asc",
"alias_substring": f"key_alias={key_alias[:-5]}&substring_matching=true",
}[filter_name]
list_query: Final = (
query
if filter_name == "alias_substring"
else f"{query}&key_alias={key_alias}"
)
resp: Final = await gateway_client.get(
f"/key/list?{list_query}",
headers={"Authorization": f"Bearer {admin.cleartext}"},
)
assert resp.status_code == 200, resp.text
response_body: Final = resp.json()
entries: Final = response_body["keys"]
assert response_body["total_count"] == 1
returned_hashes: Final = frozenset(
entry["token"] if isinstance(entry, dict) else entry for entry in entries
)
assert returned_hashes == {hash_token(cleartext)}
# ---------------------------------------------------------------------------
@ -836,17 +826,17 @@ async def test_key_list_filter_branches(query: str, proxy_client, world):
async def test_regenerate_with_grace_period_inserts_deprecated_row(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
suffix="grace",
)
old_hash = hash_token(cleartext)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
f"/key/{cleartext}/regenerate",
headers={"Authorization": f"Bearer {seeder}"},
json={"grace_period": "1h"},
@ -854,45 +844,45 @@ async def test_regenerate_with_grace_period_inserts_deprecated_row(
assert resp.status_code == 200, resp.text
new_key = resp.json()["key"]
# Old token should now be in the deprecated table.
deprecated_row = await prisma.db.litellm_deprecatedverificationtoken.find_unique(
deprecated_row = await db.db.litellm_deprecatedverificationtoken.find_unique(
where={"token": old_hash}
)
assert deprecated_row is not None, "old token not retained in deprecated table"
assert deprecated_row.active_token_id == hash_token(new_key)
assert deprecated_row.revoke_at is not None
# Manual cleanup — scratch prefix sweep doesn't cover this table.
await prisma.db.litellm_deprecatedverificationtoken.delete(
# Manual cleanup — resources prefix sweep doesn't cover this table.
await db.db.litellm_deprecatedverificationtoken.delete(
where={"token": old_hash}
)
async def test_regenerate_with_invalid_grace_period_format(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""Invalid grace_period format falls through silently (line 4170–4175);
regenerate still succeeds but no deprecated row is inserted."""
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.PROXY_ADMIN].user_id,
db,
resources.prefix,
user_id=actors.keys[Actor.PROXY_ADMIN].user_id,
suffix="grace-bad",
)
old_hash = hash_token(cleartext)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
f"/key/{cleartext}/regenerate",
headers={"Authorization": f"Bearer {seeder}"},
json={"grace_period": "totally-not-a-duration"},
)
assert resp.status_code == 200, resp.text
deprecated_row = await prisma.db.litellm_deprecatedverificationtoken.find_unique(
deprecated_row = await db.db.litellm_deprecatedverificationtoken.find_unique(
where={"token": old_hash}
)
assert deprecated_row is None, "deprecated row created despite invalid grace_period"
async def test_key_list_key_hash_filter_unauthorized(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""validate_key_list_check's key_hash branch (lines 4766–4789): a
cross-tenant non-admin caller asks for a key_hash they don't own → 403.
@ -903,13 +893,13 @@ async def test_key_list_key_hash_filter_unauthorized(
the key owner, team member, nor admin — i.e. CROSS_ORG_USER.
"""
cleartext = await _seed_token(
prisma,
scratch.prefix,
user_id=world.keys[Actor.OWNER].user_id,
team_id=TEAM_ALPHA,
db,
resources.prefix,
user_id=actors.keys[Actor.OWNER].user_id,
team_id=actors.team_alpha_id,
)
caller = world.keys[Actor.CROSS_ORG_USER].cleartext
resp = await proxy_client.get(
caller = actors.keys[Actor.CROSS_ORG_USER].cleartext
resp = await gateway_client.get(
f"/key/list?key_hash={hash_token(cleartext)}",
headers={"Authorization": f"Bearer {caller}"},
)

50
tests/proxy_behavior/management/test_key_aliases.py → tests/integration/management/test_key_aliases.py
View file

@ -5,7 +5,7 @@ import pytest
from litellm.proxy.utils import hash_token
from .actors import TEAM_ALPHA, TEAM_BETA, Actor
from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -16,11 +16,11 @@ pytestmark = pytest.mark.asyncio(loop_scope="session")
# own — owned by INTERNAL_USER, no team -> user_id scope only
# alpha — owned by OWNER, team TEAM_ALPHA -> team scope for alpha members
# beta — owned by CROSS_ORG_USER, TEAM_BETA
async def _seed_alias_keys(prisma, prefix: str, world) -> dict:
async def _seed_alias_keys(db, prefix: str, actors) -> dict:
spec = {
"own": (Actor.INTERNAL_USER, None),
"alpha": (Actor.OWNER, TEAM_ALPHA),
"beta": (Actor.CROSS_ORG_USER, TEAM_BETA),
"alpha": (Actor.OWNER, actors.team_alpha_id),
"beta": (Actor.CROSS_ORG_USER, actors.team_beta_id),
}
out = {}
for tag, (owner, team_id) in spec.items():
@ -29,18 +29,18 @@ async def _seed_alias_keys(prisma, prefix: str, world) -> dict:
"token": hash_token("sk-" + uuid.uuid4().hex),
"key_name": f"{prefix}-{tag}-key",
"key_alias": alias,
"user_id": world.keys[owner].user_id,
"user_id": actors.keys[owner].user_id,
"models": [],
}
if team_id is not None:
data["team_id"] = team_id
await prisma.db.litellm_verificationtoken.create(data=data)
await db.db.litellm_verificationtoken.create(data=data)
out[tag] = alias
return out
async def _fetch_aliases(proxy_client, caller_cleartext: str, query: str) -> set:
resp = await proxy_client.get(
async def _fetch_aliases(gateway_client, caller_cleartext: str, query: str) -> set:
resp = await gateway_client.get(
f"/key/aliases?{query}&size=100",
headers={"Authorization": f"Bearer {caller_cleartext}"},
)
@ -73,17 +73,17 @@ async def test_key_aliases_visibility(
actor: Actor,
expected_status: int,
expected_tags: FrozenSet[str],
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
aliases = await _seed_alias_keys(prisma, scratch.prefix, world)
aliases = await _seed_alias_keys(db, resources.prefix, actors)
known = {v: k for k, v in aliases.items()}
resp = await proxy_client.get(
f"/key/aliases?search={scratch.prefix}&size=100",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
resp = await gateway_client.get(
f"/key/aliases?search={resources.prefix}&size=100",
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
)
assert (
resp.status_code == expected_status
@ -97,23 +97,23 @@ async def test_key_aliases_visibility(
), f"{actor.value}: expected {sorted(expected_tags)}, got {sorted(visible)}"
async def test_key_aliases_team_id_filter(proxy_client, prisma, scratch, world):
async def test_key_aliases_team_id_filter(gateway_client, db, resources, actors):
"""team_id filter narrows the result to keys of that team."""
aliases = await _seed_alias_keys(prisma, scratch.prefix, world)
aliases = await _seed_alias_keys(db, resources.prefix, actors)
returned = await _fetch_aliases(
proxy_client,
world.keys[Actor.PROXY_ADMIN].cleartext,
f"search={scratch.prefix}&team_id={TEAM_ALPHA}",
gateway_client,
actors.keys[Actor.PROXY_ADMIN].cleartext,
f"search={resources.prefix}&team_id={actors.team_alpha_id}",
)
assert returned & set(aliases.values()) == {aliases["alpha"]}
async def test_key_aliases_search_filter(proxy_client, prisma, scratch, world):
async def test_key_aliases_search_filter(gateway_client, db, resources, actors):
"""search is a case-insensitive substring match on key_alias."""
aliases = await _seed_alias_keys(prisma, scratch.prefix, world)
aliases = await _seed_alias_keys(db, resources.prefix, actors)
returned = await _fetch_aliases(
proxy_client,
world.keys[Actor.PROXY_ADMIN].cleartext,
gateway_client,
actors.keys[Actor.PROXY_ADMIN].cleartext,
f"search={aliases['beta']}",
)
assert returned & set(aliases.values()) == {aliases["beta"]}

60
tests/proxy_behavior/management/test_key_block_unblock.py → tests/integration/management/test_key_block_unblock.py
View file

@ -4,8 +4,8 @@ import pytest
from litellm.proxy.utils import hash_token
from .actors import TEAM_ALPHA, TEAM_BETA, Actor
from .conftest import create_scratch_key
from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, Actor
from tests.integration._support.actors import create_scratch_key
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -40,26 +40,26 @@ _SCENARIOS = [
]
async def _seed_target(proxy_client, seeder, scratch_prefix, world, shape, caller):
async def _seed_target(gateway_client, seeder, scratch_prefix, actors, shape, caller):
if shape == "self":
return await create_scratch_key(
proxy_client, seeder, scratch_prefix, user_id=caller.user_id
gateway_client, seeder, scratch_prefix, user_id=caller.user_id
)
if shape == "owner":
return await create_scratch_key(
proxy_client,
gateway_client,
seeder,
scratch_prefix,
user_id=world.keys[Actor.OWNER].user_id,
team_id=TEAM_ALPHA,
user_id=actors.keys[Actor.OWNER].user_id,
team_id=actors.team_alpha_id,
)
if shape == "cross_org":
return await create_scratch_key(
proxy_client,
gateway_client,
seeder,
scratch_prefix,
user_id=world.keys[Actor.CROSS_ORG_USER].user_id,
team_id=TEAM_BETA,
user_id=actors.keys[Actor.CROSS_ORG_USER].user_id,
team_id=actors.team_beta_id,
)
pytest.fail(f"unknown shape={shape}") # pragma: no cover
@ -75,25 +75,25 @@ async def test_key_block_unblock_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
caller = world.keys[actor]
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
caller = actors.keys[actor]
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
target_cleartext = await _seed_target(
proxy_client, seeder, scratch.prefix, world, shape, caller
gateway_client, seeder, resources.prefix, actors, shape, caller
)
target_hashed = hash_token(target_cleartext)
# /unblock starts from a blocked row so a 200 is observable as True->False.
if route == "unblock":
await prisma.db.litellm_verificationtoken.update(
await db.db.litellm_verificationtoken.update(
where={"token": target_hashed}, data={"blocked": True}
)
resp = await proxy_client.post(
resp = await gateway_client.post(
f"/key/{route}",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"key": target_cleartext},
@ -102,7 +102,7 @@ async def test_key_block_unblock_authz_matrix(
resp.status_code == expected_status
), f"{route} {actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": target_hashed}
)
assert row is not None
@ -114,27 +114,27 @@ async def test_key_block_unblock_authz_matrix(
assert bool(row.blocked) is (route == "unblock"), "denied but blocked mutated"
async def test_key_block_unblock_round_trip(proxy_client, prisma, scratch, world):
async def test_key_block_unblock_round_trip(gateway_client, db, resources, actors):
"""PROXY_ADMIN block then unblock flips the blocked column True then False."""
admin = world.keys[Actor.PROXY_ADMIN]
admin = actors.keys[Actor.PROXY_ADMIN]
target = await create_scratch_key(
proxy_client, admin.cleartext, scratch.prefix, user_id=admin.user_id
gateway_client, admin.cleartext, resources.prefix, user_id=admin.user_id
)
hashed = hash_token(target)
headers = {"Authorization": f"Bearer {admin.cleartext}"}
blocked = await proxy_client.post(
blocked = await gateway_client.post(
"/key/block", headers=headers, json={"key": target}
)
assert blocked.status_code == 200, blocked.text
row = await prisma.db.litellm_verificationtoken.find_unique(where={"token": hashed})
row = await db.db.litellm_verificationtoken.find_unique(where={"token": hashed})
assert row is not None and row.blocked is True
unblocked = await proxy_client.post(
unblocked = await gateway_client.post(
"/key/unblock", headers=headers, json={"key": target}
)
assert unblocked.status_code == 200, unblocked.text
row = await prisma.db.litellm_verificationtoken.find_unique(where={"token": hashed})
row = await db.db.litellm_verificationtoken.find_unique(where={"token": hashed})
assert row is not None and row.blocked is False
@ -143,13 +143,13 @@ async def test_key_block_unblock_round_trip(proxy_client, prisma, scratch, world
"actor", [Actor.PROXY_ADMIN, Actor.TEAM_ADMIN], ids=["proxy_admin", "team_admin"]
)
async def test_key_block_unblock_missing_key_returns_404(
route: str, actor: Actor, proxy_client, world
route: str, actor: Actor, gateway_client, actors
):
"""A well-formed but unseeded key is 404 — not 401/403 — for both the
PROXY_ADMIN existence check and the non-admin _check_key_admin_access path."""
caller = world.keys[actor]
caller = actors.keys[actor]
missing = "sk-" + uuid.uuid4().hex
resp = await proxy_client.post(
resp = await gateway_client.post(
f"/key/{route}",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"key": missing},

88
tests/proxy_behavior/management/test_key_budget_limits.py → tests/integration/management/test_key_budget_limits.py
View file

@ -26,8 +26,8 @@ from typing import Any, Dict
import pytest
from .actors import Actor
from .conftest import create_scratch_org, create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_org, create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -93,24 +93,24 @@ async def test_check_team_key_limits_aggregate(
body_extras: Dict[str, Any],
expected_status: int,
detail_substring,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
team_id = await create_scratch_team(
prisma,
team_id=scratch.tag("team"),
db,
team_id=resources.tag("team"),
tpm_limit=team_tpm,
rpm_limit=team_rpm,
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
body: Dict[str, Any] = {
"key_alias": scratch.prefix,
"key_alias": resources.prefix,
"team_id": team_id,
**body_extras,
}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
@ -121,8 +121,8 @@ async def test_check_team_key_limits_aggregate(
if detail_substring is not None:
assert detail_substring in resp.text, resp.text
rows = await prisma.db.litellm_verificationtoken.find_many(
where={"key_alias": scratch.prefix}
rows = await db.db.litellm_verificationtoken.find_many(
where={"key_alias": resources.prefix}
)
if expected_status == 200:
assert len(rows) == 1
@ -178,23 +178,23 @@ async def test_check_team_key_limits_model_specific(
body_extras: Dict[str, Any],
expected_status: int,
detail_substring,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
team_id = await create_scratch_team(
prisma,
team_id=scratch.tag("team"),
db,
team_id=resources.tag("team"),
metadata=team_metadata,
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
body: Dict[str, Any] = {
"key_alias": scratch.prefix,
"key_alias": resources.prefix,
"team_id": team_id,
**body_extras,
}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
@ -204,8 +204,8 @@ async def test_check_team_key_limits_model_specific(
), f"{body!r} → {resp.status_code}: {resp.text}"
if detail_substring is not None:
assert detail_substring in resp.text, resp.text
rows = await prisma.db.litellm_verificationtoken.find_many(
where={"key_alias": scratch.prefix}
rows = await db.db.litellm_verificationtoken.find_many(
where={"key_alias": resources.prefix}
)
assert len(rows) == (1 if expected_status == 200 else 0)
@ -273,19 +273,19 @@ async def test_check_org_key_limits_aggregate(
body_extras: Dict[str, Any],
expected_status: int,
detail_substring,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
org_id = await create_scratch_org(prisma, scratch.prefix, **org_budget)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
org_id = await create_scratch_org(db, resources.prefix, **org_budget)
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
body: Dict[str, Any] = {
"key_alias": scratch.prefix,
"key_alias": resources.prefix,
"organization_id": org_id,
**body_extras,
}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
@ -295,8 +295,8 @@ async def test_check_org_key_limits_aggregate(
), f"{body!r} → {resp.status_code}: {resp.text}"
if detail_substring is not None:
assert detail_substring in resp.text, resp.text
rows = await prisma.db.litellm_verificationtoken.find_many(
where={"key_alias": scratch.prefix}
rows = await db.db.litellm_verificationtoken.find_many(
where={"key_alias": resources.prefix}
)
assert len(rows) == (1 if expected_status == 200 else 0)
@ -347,19 +347,19 @@ async def test_check_org_key_limits_model_specific(
body_extras: Dict[str, Any],
expected_status: int,
detail_substring,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
org_id = await create_scratch_org(prisma, scratch.prefix, metadata=org_metadata)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
org_id = await create_scratch_org(db, resources.prefix, metadata=org_metadata)
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
body: Dict[str, Any] = {
"key_alias": scratch.prefix,
"key_alias": resources.prefix,
"organization_id": org_id,
**body_extras,
}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
@ -369,8 +369,8 @@ async def test_check_org_key_limits_model_specific(
), f"{body!r} → {resp.status_code}: {resp.text}"
if detail_substring is not None:
assert detail_substring in resp.text, resp.text
rows = await prisma.db.litellm_verificationtoken.find_many(
where={"key_alias": scratch.prefix}
rows = await db.db.litellm_verificationtoken.find_many(
where={"key_alias": resources.prefix}
)
assert len(rows) == (1 if expected_status == 200 else 0)

44
tests/proxy_behavior/management/test_key_bulk_update.py → tests/integration/management/test_key_bulk_update.py
View file

@ -4,8 +4,8 @@ import pytest
from litellm.proxy.utils import hash_token
from .actors import Actor
from .conftest import create_scratch_key
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_key
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -37,19 +37,19 @@ _MATRIX = [
async def test_key_bulk_update_authz_matrix(
actor: Actor,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
caller = world.keys[actor]
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
caller = actors.keys[actor]
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
target = await create_scratch_key(
proxy_client, seeder, scratch.prefix, user_id=caller.user_id
gateway_client, seeder, resources.prefix, user_id=caller.user_id
)
hashed = hash_token(target)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/bulk_update",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"keys": [{"key": target, "max_budget": _MARKER_BUDGET}]},
@ -58,7 +58,7 @@ async def test_key_bulk_update_authz_matrix(
resp.status_code == expected_status
), f"{actor.value}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_verificationtoken.find_unique(where={"token": hashed})
row = await db.db.litellm_verificationtoken.find_unique(where={"token": hashed})
assert row is not None
if expected_status == 200:
body = resp.json()
@ -69,39 +69,39 @@ async def test_key_bulk_update_authz_matrix(
assert row.max_budget != _MARKER_BUDGET, "denied but key mutated"
async def test_key_bulk_update_empty_keys_is_400(proxy_client, world):
async def test_key_bulk_update_empty_keys_is_400(gateway_client, actors):
"""An empty batch is rejected 400 before any per-key processing."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/bulk_update",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"keys": []},
)
assert resp.status_code == 400, resp.text
async def test_key_bulk_update_over_max_batch_is_400(proxy_client, world):
async def test_key_bulk_update_over_max_batch_is_400(gateway_client, actors):
"""A batch larger than the 500-key cap is rejected 400."""
items = [{"key": "sk-" + uuid.uuid4().hex} for _ in range(501)]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/bulk_update",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"keys": items},
)
assert resp.status_code == 400, resp.text
async def test_key_bulk_update_per_key_failure_is_isolated(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""One bad key in the batch does not abort the others — it lands in
failed_updates while the valid key is still updated."""
admin = world.keys[Actor.PROXY_ADMIN]
admin = actors.keys[Actor.PROXY_ADMIN]
valid = await create_scratch_key(
proxy_client, admin.cleartext, scratch.prefix, user_id=admin.user_id
gateway_client, admin.cleartext, resources.prefix, user_id=admin.user_id
)
missing = "sk-" + uuid.uuid4().hex
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/bulk_update",
headers={"Authorization": f"Bearer {admin.cleartext}"},
json={
@ -117,7 +117,7 @@ async def test_key_bulk_update_per_key_failure_is_isolated(
assert len(body["successful_updates"]) == 1
assert len(body["failed_updates"]) == 1
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(valid)}
)
assert row is not None and row.max_budget == _MARKER_BUDGET

46
tests/proxy_behavior/management/test_key_delete.py → tests/integration/management/test_key_delete.py
View file

@ -4,8 +4,8 @@ import pytest
from litellm.proxy.utils import hash_token
from .actors import TEAM_ALPHA, TEAM_BETA, Actor
from .conftest import create_scratch_key
from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, Actor
from tests.integration._support.actors import create_scratch_key
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -46,40 +46,40 @@ async def test_key_delete_authz_matrix(
actor: Actor,
target_shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
caller = world.keys[actor]
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
caller = actors.keys[actor]
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
if target_shape == "self":
target_cleartext = await create_scratch_key(
proxy_client, seeder, scratch.prefix, user_id=caller.user_id
gateway_client, seeder, resources.prefix, user_id=caller.user_id
)
elif target_shape == "owner":
target_cleartext = await create_scratch_key(
proxy_client,
gateway_client,
seeder,
scratch.prefix,
user_id=world.keys[Actor.OWNER].user_id,
team_id=TEAM_ALPHA,
resources.prefix,
user_id=actors.keys[Actor.OWNER].user_id,
team_id=actors.team_alpha_id,
)
elif target_shape == "cross_org":
target_cleartext = await create_scratch_key(
proxy_client,
gateway_client,
seeder,
scratch.prefix,
user_id=world.keys[Actor.CROSS_ORG_USER].user_id,
team_id=TEAM_BETA,
resources.prefix,
user_id=actors.keys[Actor.CROSS_ORG_USER].user_id,
team_id=actors.team_beta_id,
)
else:
pytest.fail(f"unknown target_shape={target_shape}")
target_hashed = hash_token(target_cleartext)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/delete",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"keys": [target_cleartext]},
@ -88,10 +88,10 @@ async def test_key_delete_authz_matrix(
resp.status_code == expected_status
), f"{actor.value} {target_shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": target_hashed}
)
auth_check = await proxy_client.get(
auth_check = await gateway_client.get(
"/key/info", headers={"Authorization": f"Bearer {target_cleartext}"}
)
@ -103,11 +103,11 @@ async def test_key_delete_authz_matrix(
assert auth_check.status_code == 200
async def test_key_delete_missing_key_is_404(proxy_client, world):
async def test_key_delete_missing_key_is_404(gateway_client, actors):
"""Deleting a key absent from the DB is a 404 — not 401/403."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/delete",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"keys": ["sk-" + uuid.uuid4().hex]},
)
assert resp.status_code == 404, resp.text

20
tests/proxy_behavior/management/test_key_generate.py → tests/integration/management/test_key_generate.py
View file

@ -2,7 +2,7 @@ from typing import Any, Dict
import pytest
from .actors import TEAM_ALPHA, TEAM_BETA, Actor
from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -42,15 +42,15 @@ async def test_key_generate_authz_matrix(
actor: Actor,
body_extras: Dict[str, Any],
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
seeded = world.keys[actor]
body: Dict[str, Any] = {"key_alias": scratch.prefix, **body_extras}
seeded = actors.keys[actor]
body: Dict[str, Any] = {"key_alias": resources.prefix, **body_extras}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeded.cleartext}"},
json=body,
@ -59,8 +59,8 @@ async def test_key_generate_authz_matrix(
resp.status_code == expected_status
), f"{actor.value} {body!r} → {resp.status_code}: {resp.text}"
rows = await prisma.db.litellm_verificationtoken.find_many(
where={"key_alias": scratch.prefix}
rows = await db.db.litellm_verificationtoken.find_many(
where={"key_alias": resources.prefix}
)
if expected_status == 200:
cleartext = resp.json()["key"]

14
tests/proxy_behavior/management/test_key_health.py → tests/integration/management/test_key_health.py
View file

@ -1,17 +1,17 @@
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
# POST /key/health has no role gate — it reflects the caller's OWN key logging
# metadata. The world keys carry no "logging" metadata, so every authenticated
# metadata. The actors keys carry no "logging" metadata, so every authenticated
# actor gets 200 with key="healthy". This pins auth-required + route coverage.
@pytest.mark.parametrize("actor", list(Actor), ids=[a.value for a in Actor])
async def test_key_health_each_actor_is_healthy(actor: Actor, proxy_client, world):
caller = world.keys[actor]
resp = await proxy_client.post(
async def test_key_health_each_actor_is_healthy(actor: Actor, gateway_client, actors):
caller = actors.keys[actor]
resp = await gateway_client.post(
"/key/health",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)
@ -19,6 +19,6 @@ async def test_key_health_each_actor_is_healthy(actor: Actor, proxy_client, worl
assert resp.json()["key"] == "healthy"
async def test_key_health_requires_auth(proxy_client):
resp = await proxy_client.post("/key/health")
async def test_key_health_requires_auth(gateway_client):
resp = await gateway_client.post("/key/health", headers={"Authorization": "Bearer invalid"})
assert resp.status_code == 401, resp.text

10
tests/proxy_behavior/management/test_key_info.py → tests/integration/management/test_key_info.py
View file

@ -1,6 +1,6 @@
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -53,12 +53,12 @@ _SCENARIOS = [
ids=[s[0] for s in _SCENARIOS],
)
async def test_key_info_authz_matrix(
actor: Actor, target_actor: Actor, expected_status: int, proxy_client, world
actor: Actor, target_actor: Actor, expected_status: int, gateway_client, actors
):
caller = world.keys[actor]
target = world.keys[target_actor]
caller = actors.keys[actor]
target = actors.keys[target_actor]
resp = await proxy_client.get(
resp = await gateway_client.get(
f"/key/info?key={target.cleartext}",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)

28
tests/proxy_behavior/management/test_key_info_v2.py → tests/integration/management/test_key_info_v2.py
View file

@ -2,16 +2,16 @@ import uuid
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
# POST /v2/key/info resolves the posted keys, then drops any key the caller
# cannot see via _can_user_query_key_info — silently, no 403. A non-admin sees
# a key it owns (user_id match) or a key whose team it belongs to. The world's
# a key it owns (user_id match) or a key whose team it belongs to. The actors's
# TEAM_ALPHA members all see each other's keys; CROSS_ORG_USER and the org
# admins see only their own. The request is posted with every world key, and
# admins see only their own. The request is posted with every actors key, and
# the returned info set is asserted to equal the visible subset.
_ALPHA_KEYS = frozenset(
{
@ -40,14 +40,14 @@ _VISIBILITY = {
list(_VISIBILITY.items()),
ids=[a.value for a in _VISIBILITY],
)
async def test_key_info_v2_visibility(actor, expected_visible, proxy_client, world):
caller = world.keys[actor]
user_id_to_actor = {world.keys[a].user_id: a for a in Actor}
async def test_key_info_v2_visibility(actor, expected_visible, gateway_client, actors):
caller = actors.keys[actor]
user_id_to_actor = {actors.keys[a].user_id: a for a in Actor}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/v2/key/info",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"keys": [world.keys[a].cleartext for a in Actor]},
json={"keys": [actors.keys[a].cleartext for a in Actor]},
)
assert resp.status_code == 200, f"{actor.value}: {resp.status_code} {resp.text}"
@ -62,20 +62,20 @@ async def test_key_info_v2_visibility(actor, expected_visible, proxy_client, wor
)
async def test_key_info_v2_no_body_is_422(proxy_client, world):
async def test_key_info_v2_no_body_is_422(gateway_client, actors):
"""A request with no body is a 422 — the handler has no keys to resolve."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/v2/key/info",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 422, resp.text
async def test_key_info_v2_unknown_key_returns_empty_info(proxy_client, world):
async def test_key_info_v2_unknown_key_returns_empty_info(gateway_client, actors):
"""Keys that resolve to no rows yield an empty info list, not an error."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/v2/key/info",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"keys": ["sk-" + uuid.uuid4().hex]},
)
assert resp.status_code == 200, resp.text

104
tests/proxy_behavior/management/test_key_list.py → tests/integration/management/test_key_list.py
View file

@ -4,8 +4,8 @@ import pytest
from litellm.proxy.utils import hash_token
from .actors import TEAM_ALPHA, Actor
from .conftest import create_scratch_key
from tests.integration._support.actors import TEAM_ALPHA, Actor
from tests.integration._support.actors import create_scratch_key
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -24,13 +24,13 @@ _VISIBILITY = {
}
async def _all_visible_hashes(proxy_client, caller_cleartext) -> set:
async def _all_visible_hashes(gateway_client, caller_cleartext) -> set:
"""Walk every /key/list page — size is capped at 100 by the endpoint, so a
single request can truncate PROXY_ADMIN's view on a non-fresh DB."""
hashes: set = set()
page = 1
while True:
resp = await proxy_client.get(
resp = await gateway_client.get(
f"/key/list?page={page}&size=100",
headers={"Authorization": f"Bearer {caller_cleartext}"},
)
@ -51,12 +51,12 @@ async def _all_visible_hashes(proxy_client, caller_cleartext) -> set:
ids=[a.value for a in _VISIBILITY],
)
async def test_key_list_visibility(
actor: Actor, expected_visible: FrozenSet[Actor], proxy_client, world
actor: Actor, expected_visible: FrozenSet[Actor], gateway_client, actors
):
caller = world.keys[actor]
hashed_to_actor = {world.keys[a].hashed: a for a in Actor}
caller = actors.keys[actor]
hashed_to_actor = {actors.keys[a].hashed: a for a in Actor}
returned_hashes = await _all_visible_hashes(proxy_client, caller.cleartext)
returned_hashes = await _all_visible_hashes(gateway_client, caller.cleartext)
visible_seeded = {
hashed_to_actor[h] for h in returned_hashes if h in hashed_to_actor
}
@ -66,8 +66,8 @@ async def test_key_list_visibility(
)
async def _list_hashes(proxy_client, caller_cleartext: str, query: str) -> set:
resp = await proxy_client.get(
async def _list_hashes(gateway_client, caller_cleartext: str, query: str) -> set:
resp = await gateway_client.get(
f"/key/list?{query}&size=100",
headers={"Authorization": f"Bearer {caller_cleartext}"},
)
@ -80,125 +80,125 @@ async def _list_hashes(proxy_client, caller_cleartext: str, query: str) -> set:
return hashes
async def test_key_list_admin_key_alias_substring_match(proxy_client, scratch, world):
async def test_key_list_admin_key_alias_substring_match(gateway_client, resources, actors):
"""A PROXY_ADMIN's key_alias filter is a case-insensitive substring match
when substring_matching=true is requested (the dashboard search box); a
narrower fragment selects the subset whose alias contains it. Substring
matching is opt-in: without the flag the filter is exact (see
test_key_list_admin_key_alias_exact_without_substring_flag)."""
admin = world.keys[Actor.PROXY_ADMIN]
admin = actors.keys[Actor.PROXY_ADMIN]
a = await create_scratch_key(
proxy_client,
gateway_client,
admin.cleartext,
scratch.prefix,
resources.prefix,
user_id=admin.user_id,
key_alias=f"{scratch.prefix}-sub-a",
key_alias=f"{resources.prefix}-sub-a",
)
b = await create_scratch_key(
proxy_client,
gateway_client,
admin.cleartext,
scratch.prefix,
resources.prefix,
user_id=admin.user_id,
key_alias=f"{scratch.prefix}-sub-b",
key_alias=f"{resources.prefix}-sub-b",
)
seeded = {hash_token(a), hash_token(b)}
broad = await _list_hashes(
proxy_client,
gateway_client,
admin.cleartext,
f"key_alias={scratch.prefix}-sub&substring_matching=true",
f"key_alias={resources.prefix}-sub&substring_matching=true",
)
assert broad & seeded == seeded
narrow = await _list_hashes(
proxy_client,
gateway_client,
admin.cleartext,
f"key_alias={scratch.prefix}-sub-a&substring_matching=true",
f"key_alias={resources.prefix}-sub-a&substring_matching=true",
)
assert narrow & seeded == {hash_token(a)}
async def test_key_list_admin_key_alias_exact_without_substring_flag(
proxy_client, scratch, world
gateway_client, resources, actors
):
"""Regression guard for the prior exact-match contract: without
substring_matching, even a PROXY_ADMIN's key_alias filter is exact, so a
fragment of a seeded alias does not select it."""
admin = world.keys[Actor.PROXY_ADMIN]
full_alias = f"{scratch.prefix}-exactflag"
admin = actors.keys[Actor.PROXY_ADMIN]
full_alias = f"{resources.prefix}-exactflag"
key = await create_scratch_key(
proxy_client,
gateway_client,
admin.cleartext,
scratch.prefix,
resources.prefix,
user_id=admin.user_id,
key_alias=full_alias,
)
key_hash = hash_token(key)
exact = await _list_hashes(proxy_client, admin.cleartext, f"key_alias={full_alias}")
exact = await _list_hashes(gateway_client, admin.cleartext, f"key_alias={full_alias}")
assert key_hash in exact
fragment = await _list_hashes(
proxy_client, admin.cleartext, f"key_alias={scratch.prefix}-exactfla"
gateway_client, admin.cleartext, f"key_alias={resources.prefix}-exactfla"
)
assert key_hash not in fragment
async def test_key_list_non_admin_key_alias_is_exact_match(
proxy_client, scratch, world
gateway_client, resources, actors
):
"""A non-admin's key_alias filter is exact-match only — substring filtering
is restricted to admins. The full alias matches; a fragment does not."""
caller = world.keys[Actor.INTERNAL_USER]
alias = f"{scratch.prefix}-exact"
caller = actors.keys[Actor.INTERNAL_USER]
alias = f"{resources.prefix}-exact"
key = await create_scratch_key(
proxy_client,
world.keys[Actor.PROXY_ADMIN].cleartext,
scratch.prefix,
gateway_client,
actors.keys[Actor.PROXY_ADMIN].cleartext,
resources.prefix,
user_id=caller.user_id,
key_alias=alias,
)
key_hash = hash_token(key)
exact = await _list_hashes(proxy_client, caller.cleartext, f"key_alias={alias}")
exact = await _list_hashes(gateway_client, caller.cleartext, f"key_alias={alias}")
assert key_hash in exact
fragment = await _list_hashes(
proxy_client, caller.cleartext, f"key_alias={scratch.prefix}-exac"
gateway_client, caller.cleartext, f"key_alias={resources.prefix}-exac"
)
assert key_hash not in fragment
async def test_key_list_team_id_filter(proxy_client, scratch, world):
async def test_key_list_team_id_filter(gateway_client, resources, actors):
"""A team_id filter narrows the listing to keys of that team."""
admin = world.keys[Actor.PROXY_ADMIN]
admin = actors.keys[Actor.PROXY_ADMIN]
team_key = await create_scratch_key(
proxy_client,
gateway_client,
admin.cleartext,
scratch.prefix,
user_id=world.keys[Actor.OWNER].user_id,
team_id=TEAM_ALPHA,
key_alias=f"{scratch.prefix}-team",
resources.prefix,
user_id=actors.keys[Actor.OWNER].user_id,
team_id=actors.team_alpha_id,
key_alias=f"{resources.prefix}-team",
)
no_team_key = await create_scratch_key(
proxy_client,
gateway_client,
admin.cleartext,
scratch.prefix,
resources.prefix,
user_id=admin.user_id,
key_alias=f"{scratch.prefix}-noteam",
key_alias=f"{resources.prefix}-noteam",
)
hashes = await _list_hashes(proxy_client, admin.cleartext, f"team_id={TEAM_ALPHA}")
hashes = await _list_hashes(gateway_client, admin.cleartext, f"team_id={actors.team_alpha_id}")
assert hash_token(team_key) in hashes
assert hash_token(no_team_key) not in hashes
async def test_key_list_non_admin_cannot_filter_other_team(proxy_client, world):
async def test_key_list_non_admin_cannot_filter_other_team(gateway_client, actors):
"""A non-admin filtering by a team it does not belong to is rejected 403."""
resp = await proxy_client.get(
f"/key/list?team_id={world.team_beta_id}",
resp = await gateway_client.get(
f"/key/list?team_id={actors.team_beta_id}",
headers={
"Authorization": f"Bearer {world.keys[Actor.INTERNAL_USER].cleartext}"
"Authorization": f"Bearer {actors.keys[Actor.INTERNAL_USER].cleartext}"
},
)
assert resp.status_code == 403, resp.text

107
tests/proxy_behavior/management/test_key_regenerate.py → tests/integration/management/test_key_regenerate.py
View file

@ -1,12 +1,7 @@
import litellm
import pytest
from litellm.types.proxy.management_endpoints.ui_sso import (
LiteLLM_UpperboundKeyGenerateParams,
)
from .actors import TEAM_ALPHA, TEAM_BETA, Actor
from .conftest import create_scratch_key
from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, Actor
from tests.integration._support.actors import create_scratch_key, owned_settings_client
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -39,8 +34,8 @@ _SCENARIOS = [
]
async def _info(proxy_client, cleartext: str):
return await proxy_client.get(
async def _info(gateway_client, cleartext: str):
return await gateway_client.get(
"/key/info", headers={"Authorization": f"Bearer {cleartext}"}
)
@ -54,37 +49,37 @@ async def test_key_regenerate_authz_matrix(
actor: Actor,
target_shape: str,
expected_status: int,
proxy_client,
scratch,
world,
gateway_client,
resources,
actors,
):
caller = world.keys[actor]
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
caller = actors.keys[actor]
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
if target_shape == "self":
target_cleartext = await create_scratch_key(
proxy_client, seeder, scratch.prefix, user_id=caller.user_id
gateway_client, seeder, resources.prefix, user_id=caller.user_id
)
elif target_shape == "owner":
target_cleartext = await create_scratch_key(
proxy_client,
gateway_client,
seeder,
scratch.prefix,
user_id=world.keys[Actor.OWNER].user_id,
team_id=TEAM_ALPHA,
resources.prefix,
user_id=actors.keys[Actor.OWNER].user_id,
team_id=actors.team_alpha_id,
)
elif target_shape == "cross_org":
target_cleartext = await create_scratch_key(
proxy_client,
gateway_client,
seeder,
scratch.prefix,
user_id=world.keys[Actor.CROSS_ORG_USER].user_id,
team_id=TEAM_BETA,
resources.prefix,
user_id=actors.keys[Actor.CROSS_ORG_USER].user_id,
team_id=actors.team_beta_id,
)
else:
pytest.fail(f"unknown target_shape={target_shape}")
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/regenerate",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"key": target_cleartext},
@ -96,21 +91,21 @@ async def test_key_regenerate_authz_matrix(
if expected_status == 200:
new_cleartext = resp.json()["key"]
assert new_cleartext.startswith("sk-") and new_cleartext != target_cleartext
assert (await _info(proxy_client, target_cleartext)).status_code == 401
assert (await _info(proxy_client, new_cleartext)).status_code == 200
assert (await _info(gateway_client, target_cleartext)).status_code == 401
assert (await _info(gateway_client, new_cleartext)).status_code == 200
else:
# Denied: rotation must not have leaked — old cleartext still works.
assert (await _info(proxy_client, target_cleartext)).status_code == 200
assert (await _info(gateway_client, target_cleartext)).status_code == 200
async def test_key_path_regenerate_smoke(proxy_client, scratch, world):
async def test_key_path_regenerate_smoke(gateway_client, resources, actors):
"""Pins that POST /key/{key:path}/regenerate shares the same handler."""
caller = world.keys[Actor.PROXY_ADMIN]
caller = actors.keys[Actor.PROXY_ADMIN]
target_cleartext = await create_scratch_key(
proxy_client, caller.cleartext, scratch.prefix, user_id=caller.user_id
gateway_client, caller.cleartext, resources.prefix, user_id=caller.user_id
)
resp = await proxy_client.post(
resp = await gateway_client.post(
f"/key/{target_cleartext}/regenerate",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={},
@ -118,48 +113,40 @@ async def test_key_path_regenerate_smoke(proxy_client, scratch, world):
assert resp.status_code == 200, resp.text
new_cleartext = resp.json()["key"]
assert new_cleartext.startswith("sk-") and new_cleartext != target_cleartext
assert (await _info(proxy_client, target_cleartext)).status_code == 401
assert (await _info(proxy_client, new_cleartext)).status_code == 200
assert (await _info(gateway_client, target_cleartext)).status_code == 401
assert (await _info(gateway_client, new_cleartext)).status_code == 200
async def test_key_regenerate_enforces_upperbound_key_params(
proxy_client, scratch, world, monkeypatch
gateway_client, resources, actors, tmp_path
):
"""Regenerate runs _enforce_upperbound_key_params: a max_budget above
litellm.upperbound_key_generate_params is rejected 400, a value within the
bound is accepted. Pins #26340 (db8ef44323) — regenerate previously
bypassed the upperbound. upperbound_key_generate_params is module-level
litellm.* state, so monkeypatch save/restores it."""
admin = world.keys[Actor.PROXY_ADMIN]
admin = actors.keys[Actor.PROXY_ADMIN]
over_key = await create_scratch_key(
proxy_client,
gateway_client,
admin.cleartext,
scratch.prefix,
resources.prefix,
user_id=admin.user_id,
key_alias=f"{scratch.prefix}-over",
key_alias=f"{resources.prefix}-over",
)
within_key = await create_scratch_key(
proxy_client,
gateway_client,
admin.cleartext,
scratch.prefix,
resources.prefix,
user_id=admin.user_id,
key_alias=f"{scratch.prefix}-within",
)
monkeypatch.setattr(
litellm,
"upperbound_key_generate_params",
LiteLLM_UpperboundKeyGenerateParams(max_budget=100.0),
key_alias=f"{resources.prefix}-within",
)
headers = {"Authorization": f"Bearer {admin.cleartext}"}
over = await proxy_client.post(
"/key/regenerate", headers=headers, json={"key": over_key, "max_budget": 500.0}
)
with owned_settings_client(
gateway_client, tmp_path, {"upperbound_key_generate_params": {"max_budget": 100.0}}
) as client:
over = await client.post(
"/key/regenerate", headers=headers, json={"key": over_key, "max_budget": 500.0}
)
within = await client.post(
"/key/regenerate",
headers=headers,
json={"key": within_key, "max_budget": 50.0},
)
assert over.status_code == 400, over.text
within = await proxy_client.post(
"/key/regenerate",
headers=headers,
json={"key": within_key, "max_budget": 50.0},
)
assert within.status_code == 200, within.text

54
tests/proxy_behavior/management/test_key_reset_spend.py → tests/integration/management/test_key_reset_spend.py
View file

@ -4,8 +4,8 @@ import pytest
from litellm.proxy.utils import hash_token
from .actors import TEAM_ALPHA, TEAM_BETA, Actor
from .conftest import create_scratch_key
from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, Actor
from tests.integration._support.actors import create_scratch_key
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -42,26 +42,26 @@ _SCENARIOS = [
]
async def _seed_target(proxy_client, seeder, prefix, world, shape, caller) -> str:
async def _seed_target(gateway_client, seeder, prefix, actors, shape, caller) -> str:
if shape == "self":
return await create_scratch_key(
proxy_client, seeder, prefix, user_id=caller.user_id
gateway_client, seeder, prefix, user_id=caller.user_id
)
if shape == "team_alpha":
return await create_scratch_key(
proxy_client,
gateway_client,
seeder,
prefix,
user_id=world.keys[Actor.OWNER].user_id,
team_id=TEAM_ALPHA,
user_id=actors.keys[Actor.OWNER].user_id,
team_id=actors.team_alpha_id,
)
if shape == "team_beta":
return await create_scratch_key(
proxy_client,
gateway_client,
seeder,
prefix,
user_id=world.keys[Actor.CROSS_ORG_USER].user_id,
team_id=TEAM_BETA,
user_id=actors.keys[Actor.CROSS_ORG_USER].user_id,
team_id=actors.team_beta_id,
)
pytest.fail(f"unknown shape={shape}") # pragma: no cover
@ -75,22 +75,22 @@ async def test_key_reset_spend_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
caller = world.keys[actor]
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
caller = actors.keys[actor]
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
target = await _seed_target(
proxy_client, seeder, scratch.prefix, world, shape, caller
gateway_client, seeder, resources.prefix, actors, shape, caller
)
hashed = hash_token(target)
await prisma.db.litellm_verificationtoken.update(
await db.db.litellm_verificationtoken.update(
where={"token": hashed}, data={"spend": _SEED_SPEND}
)
resp = await proxy_client.post(
resp = await gateway_client.post(
f"/key/{target}/reset_spend",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"reset_to": _RESET_TO},
@ -99,7 +99,7 @@ async def test_key_reset_spend_authz_matrix(
resp.status_code == expected_status
), f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_verificationtoken.find_unique(where={"token": hashed})
row = await db.db.litellm_verificationtoken.find_unique(where={"token": hashed})
assert row is not None
if expected_status == 200:
assert row.spend == _RESET_TO
@ -110,25 +110,25 @@ async def test_key_reset_spend_authz_matrix(
@pytest.mark.parametrize(
"actor", [Actor.PROXY_ADMIN, Actor.TEAM_ADMIN], ids=["proxy_admin", "team_admin"]
)
async def test_key_reset_spend_missing_key_is_404(actor: Actor, proxy_client, world):
async def test_key_reset_spend_missing_key_is_404(actor: Actor, gateway_client, actors):
"""A well-formed but unseeded key is 404 before any spend validation."""
resp = await proxy_client.post(
resp = await gateway_client.post(
f"/key/sk-{uuid.uuid4().hex}/reset_spend",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
json={"reset_to": 0.0},
)
assert resp.status_code == 404, resp.text
async def test_key_reset_spend_above_current_spend_is_400(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""reset_to above the key's current spend is rejected 400."""
admin = world.keys[Actor.PROXY_ADMIN]
admin = actors.keys[Actor.PROXY_ADMIN]
target = await create_scratch_key(
proxy_client, admin.cleartext, scratch.prefix, user_id=admin.user_id
gateway_client, admin.cleartext, resources.prefix, user_id=admin.user_id
)
resp = await proxy_client.post(
resp = await gateway_client.post(
f"/key/{target}/reset_spend",
headers={"Authorization": f"Bearer {admin.cleartext}"},
json={"reset_to": 1.0},

36
tests/proxy_behavior/management/test_key_service_account_generate.py → tests/integration/management/test_key_service_account_generate.py
View file

@ -1,6 +1,6 @@
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -45,23 +45,23 @@ async def test_key_service_account_generate_authz_matrix(
actor: Actor,
team_target: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
caller = world.keys[actor]
caller = actors.keys[actor]
team_id = {
"own": world.team_alpha_id,
"cross_org": world.team_beta_id,
"own": actors.team_alpha_id,
"cross_org": actors.team_beta_id,
"none": None,
}[team_target]
body = {"key_alias": scratch.prefix}
body = {"key_alias": resources.prefix}
if team_id is not None:
body["team_id"] = team_id
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/service-account/generate",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json=body,
@ -70,8 +70,8 @@ async def test_key_service_account_generate_authz_matrix(
resp.status_code == expected_status
), f"{actor.value} {team_target}: {resp.status_code} {resp.text}"
rows = await prisma.db.litellm_verificationtoken.find_many(
where={"key_alias": scratch.prefix}
rows = await db.db.litellm_verificationtoken.find_many(
where={"key_alias": resources.prefix}
)
if expected_status == 200:
assert len(rows) == 1
@ -83,16 +83,16 @@ async def test_key_service_account_generate_authz_matrix(
async def test_key_service_account_generate_unknown_team_is_400(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""A team_id absent from the database is rejected 400."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/service-account/generate",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"key_alias": scratch.prefix, "team_id": scratch.tag("no-such-team")},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"key_alias": resources.prefix, "team_id": resources.tag("no-such-team")},
)
assert resp.status_code == 400, resp.text
rows = await prisma.db.litellm_verificationtoken.find_many(
where={"key_alias": scratch.prefix}
rows = await db.db.litellm_verificationtoken.find_many(
where={"key_alias": resources.prefix}
)
assert rows == []

84
tests/proxy_behavior/management/test_key_team_change.py → tests/integration/management/test_key_team_change.py
View file

@ -21,8 +21,8 @@ from prisma import Json
from litellm.proxy.utils import hash_token
from .actors import TEAM_ALPHA, Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import TEAM_ALPHA, Actor
from tests.integration._support.actors import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -30,7 +30,7 @@ KEY_MODEL = "phase4-f2-key-model"
async def _seed_key_with_limits(
prisma,
db,
scratch_prefix: str,
*,
user_id: str,
@ -39,7 +39,7 @@ async def _seed_key_with_limits(
tpm_limit: Optional[int] = None,
rpm_limit: Optional[int] = None,
) -> str:
"""Raw-seed a scratch key with explicit models / tpm / rpm — /key/generate
"""Raw-seed a resources key with explicit models / tpm / rpm — /key/generate
can't set these against a non-throughput team without firing F1's guards.
Returns the cleartext key for /key/update calls."""
cleartext = "sk-" + uuid.uuid4().hex
@ -55,45 +55,45 @@ async def _seed_key_with_limits(
data["tpm_limit"] = tpm_limit
if rpm_limit is not None:
data["rpm_limit"] = rpm_limit
await prisma.db.litellm_verificationtoken.create(data=data)
await db.db.litellm_verificationtoken.create(data=data)
return cleartext
# ---------------------------------------------------------------------------
# Accepted — proxy admin moves a key into a scratch team that has the model,
# Accepted — proxy admin moves a key into a resources team that has the model,
# accommodates the limits, and has the key owner as a member.
# ---------------------------------------------------------------------------
async def test_key_team_change_accepted(proxy_client, prisma, scratch, world):
owner_id = world.keys[Actor.OWNER].user_id
async def test_key_team_change_accepted(gateway_client, db, resources, actors):
owner_id = actors.keys[Actor.OWNER].user_id
target_team = await create_scratch_team(
prisma,
team_id=scratch.tag("target"),
db,
team_id=resources.tag("target"),
member_user_ids=[owner_id],
models=[KEY_MODEL],
tpm_limit=10_000,
rpm_limit=1_000,
)
key_cleartext = await _seed_key_with_limits(
prisma,
scratch.prefix,
db,
resources.prefix,
user_id=owner_id,
team_id=TEAM_ALPHA,
team_id=actors.team_alpha_id,
models=[KEY_MODEL],
tpm_limit=500,
rpm_limit=50,
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {seeder}"},
json={"key": key_cleartext, "team_id": target_team},
)
assert resp.status_code == 200, resp.text
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(key_cleartext)}
)
assert row is not None
@ -115,9 +115,9 @@ async def test_key_team_change_accepted(proxy_client, prisma, scratch, world):
],
)
async def test_key_team_change_rejected_guards(
scenario: str, proxy_client, prisma, scratch, world
scenario: str, gateway_client, db, resources, actors
):
owner_id = world.keys[Actor.OWNER].user_id
owner_id = actors.keys[Actor.OWNER].user_id
team_kwargs: Dict[str, Any] = {
"member_user_ids": [owner_id],
@ -141,20 +141,20 @@ async def test_key_team_change_rejected_guards(
team_kwargs["member_user_ids"] = []
target_team = await create_scratch_team(
prisma,
team_id=scratch.tag("target"),
db,
team_id=resources.tag("target"),
**team_kwargs,
)
key_cleartext = await _seed_key_with_limits(
prisma,
scratch.prefix,
db,
resources.prefix,
user_id=owner_id,
team_id=TEAM_ALPHA,
team_id=actors.team_alpha_id,
**key_kwargs,
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {seeder}"},
json={"key": key_cleartext, "team_id": target_team},
@ -168,11 +168,11 @@ async def test_key_team_change_rejected_guards(
403,
), f"{scenario}: expected 400/403, got {resp.status_code}: {resp.text}"
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(key_cleartext)}
)
assert row is not None
assert row.team_id == TEAM_ALPHA, (
assert row.team_id == actors.team_alpha_id, (
f"{scenario}: row team_id mutated despite rejection — " f"got {row.team_id!r}"
)
@ -186,34 +186,34 @@ async def test_key_team_change_rejected_guards(
async def test_key_team_change_rejected_initiator_not_admin(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
owner_id = world.keys[Actor.OWNER].user_id
owner_id = actors.keys[Actor.OWNER].user_id
target_team = await create_scratch_team(
prisma,
team_id=scratch.tag("target"),
db,
team_id=resources.tag("target"),
# Owner is the key's user_id; member of the team to clear membership
# guard. Internal-user role on world.keys[INTERNAL_USER] is also a
# member of TEAM_ALPHA but never an admin → role-gate path fires.
member_user_ids=[owner_id, world.keys[Actor.INTERNAL_USER].user_id],
# guard. Internal-user role on actors.keys[INTERNAL_USER] is also a
# member of actors.team_alpha_id but never an admin → role-gate path fires.
member_user_ids=[owner_id, actors.keys[Actor.INTERNAL_USER].user_id],
models=[KEY_MODEL],
tpm_limit=10_000,
rpm_limit=1_000,
)
key_cleartext = await _seed_key_with_limits(
prisma,
scratch.prefix,
db,
resources.prefix,
user_id=owner_id,
team_id=TEAM_ALPHA,
team_id=actors.team_alpha_id,
models=[KEY_MODEL],
tpm_limit=500,
rpm_limit=50,
)
# The internal_user actor is the initiator: a TEAM_ALPHA member, but not
# The internal_user actor is the initiator: a actors.team_alpha_id member, but not
# an admin anywhere, and has no team_member_permissions for /key/update.
initiator = world.keys[Actor.INTERNAL_USER].cleartext
initiator = actors.keys[Actor.INTERNAL_USER].cleartext
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {initiator}"},
json={"key": key_cleartext, "team_id": target_team},
@ -222,8 +222,8 @@ async def test_key_team_change_rejected_initiator_not_admin(
401,
403,
), f"expected 401/403, got {resp.status_code}: {resp.text}"
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(key_cleartext)}
)
assert row is not None
assert row.team_id == TEAM_ALPHA, "row team_id mutated despite rejection"
assert row.team_id == actors.team_alpha_id, "row team_id mutated despite rejection"

78
tests/proxy_behavior/management/test_key_update.py → tests/integration/management/test_key_update.py
View file

@ -4,8 +4,8 @@ import pytest
from litellm.proxy.utils import hash_token
from .actors import TEAM_ALPHA, TEAM_BETA, Actor
from .conftest import create_scratch_key
from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, Actor
from tests.integration._support.actors import create_scratch_key
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -50,40 +50,40 @@ async def test_key_update_authz_matrix(
actor: Actor,
target_shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
caller = world.keys[actor]
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
caller = actors.keys[actor]
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
if target_shape == "self":
target_cleartext = await create_scratch_key(
proxy_client, seeder, scratch.prefix, user_id=caller.user_id
gateway_client, seeder, resources.prefix, user_id=caller.user_id
)
elif target_shape == "owner":
target_cleartext = await create_scratch_key(
proxy_client,
gateway_client,
seeder,
scratch.prefix,
user_id=world.keys[Actor.OWNER].user_id,
team_id=TEAM_ALPHA,
resources.prefix,
user_id=actors.keys[Actor.OWNER].user_id,
team_id=actors.team_alpha_id,
)
elif target_shape == "cross_org":
target_cleartext = await create_scratch_key(
proxy_client,
gateway_client,
seeder,
scratch.prefix,
user_id=world.keys[Actor.CROSS_ORG_USER].user_id,
team_id=TEAM_BETA,
resources.prefix,
user_id=actors.keys[Actor.CROSS_ORG_USER].user_id,
team_id=actors.team_beta_id,
)
else:
pytest.fail(f"unknown target_shape={target_shape}")
target_hashed = hash_token(target_cleartext)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"key": target_cleartext, "models": [MARKER_MODEL]},
@ -92,7 +92,7 @@ async def test_key_update_authz_matrix(
resp.status_code == expected_status
), f"{actor.value} {target_shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": target_hashed}
)
assert row is not None
@ -102,35 +102,35 @@ async def test_key_update_authz_matrix(
assert row.models != [MARKER_MODEL], "denied but row mutated"
async def _seed_shape(proxy_client, seeder, prefix, world, shape, caller) -> str:
async def _seed_shape(gateway_client, seeder, prefix, actors, shape, caller) -> str:
if shape == "self":
return await create_scratch_key(
proxy_client, seeder, prefix, user_id=caller.user_id
gateway_client, seeder, prefix, user_id=caller.user_id
)
if shape == "owner":
return await create_scratch_key(
proxy_client,
gateway_client,
seeder,
prefix,
user_id=world.keys[Actor.OWNER].user_id,
team_id=TEAM_ALPHA,
user_id=actors.keys[Actor.OWNER].user_id,
team_id=actors.team_alpha_id,
)
if shape == "cross_org":
return await create_scratch_key(
proxy_client,
gateway_client,
seeder,
prefix,
user_id=world.keys[Actor.CROSS_ORG_USER].user_id,
team_id=TEAM_BETA,
user_id=actors.keys[Actor.CROSS_ORG_USER].user_id,
team_id=actors.team_beta_id,
)
pytest.fail(f"unknown shape={shape}") # pragma: no cover
async def test_key_update_missing_key_is_404(proxy_client, world):
async def test_key_update_missing_key_is_404(gateway_client, actors):
"""An update targeting a key absent from the DB is a 404 — not 401/403."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"key": "sk-" + uuid.uuid4().hex, "models": [MARKER_MODEL]},
)
assert resp.status_code == 404, resp.text
@ -154,19 +154,19 @@ async def test_key_update_denied_does_not_touch_budget_counters(
actor: Actor,
target_shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
caller = world.keys[actor]
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
caller = actors.keys[actor]
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
target = await _seed_shape(
proxy_client, seeder, scratch.prefix, world, target_shape, caller
gateway_client, seeder, resources.prefix, actors, target_shape, caller
)
target_hashed = hash_token(target)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/key/update",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"key": target, "max_budget": 999.0, "tpm_limit": 888, "rpm_limit": 777},
@ -175,7 +175,7 @@ async def test_key_update_denied_does_not_touch_budget_counters(
resp.status_code == expected_status
), f"{actor.value} {target_shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": target_hashed}
)
assert row is not None

22
tests/integration/management/test_smoke.py Normal file
View file

@ -0,0 +1,22 @@
import pytest
pytestmark = pytest.mark.asyncio(loop_scope="session")
async def test_key_generate_lands_in_db(gateway_client, db, resources):
from litellm.proxy.utils import hash_token
resp = await gateway_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {gateway_client.gateway.key}"},
json={"key_alias": resources.prefix},
)
assert resp.status_code == 200, resp.text
cleartext = resp.json()["key"]
assert cleartext.startswith("sk-")
hashed = hash_token(cleartext)
row = await db.db.litellm_verificationtoken.find_unique(where={"token": hashed})
assert row is not None
assert row.token == hashed != cleartext

10
tests/proxy_behavior/management/test_team_available.py → tests/integration/management/test_team_available.py
View file

@ -1,21 +1,21 @@
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
# GET /team/available lists teams from
# litellm.default_internal_user_params["available_teams"]. The behavior world
# litellm.default_internal_user_params["available_teams"]. The behavior actors
# configures no available_teams, so the handler returns [] for every actor
# before it even reads the caller — this is the route-coverage + default-path
# pin. /team/available is an info route, so every authenticated actor reaches
# the handler.
@pytest.mark.parametrize("actor", list(Actor), ids=[a.value for a in Actor])
async def test_team_available_default_is_empty(actor: Actor, proxy_client, world):
resp = await proxy_client.get(
async def test_team_available_default_is_empty(actor: Actor, gateway_client, actors):
resp = await gateway_client.get(
"/team/available",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
)
assert resp.status_code == 200, f"{actor.value}: {resp.status_code} {resp.text}"
assert resp.json() == []

64
tests/proxy_behavior/management/test_team_block_unblock.py → tests/integration/management/test_team_block_unblock.py
View file

@ -1,7 +1,7 @@
import pytest
from .actors import Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -28,10 +28,10 @@ _MATRIX = [
]
async def _seed_target(prisma, world, shape: str, team_id: str) -> str:
"""Raw-seed the scratch target team; returns its organization_id."""
org_id = world.org_a_id if shape == "alpha" else world.org_b_id
await create_scratch_team(prisma, team_id, organization_id=org_id)
async def _seed_target(db, actors, shape: str, team_id: str) -> str:
"""Raw-seed the resources target team; returns its organization_id."""
org_id = actors.org_a_id if shape == "alpha" else actors.org_b_id
await create_scratch_team(db, team_id, organization_id=org_id)
return org_id
@ -46,31 +46,31 @@ async def test_team_block_unblock_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
org_id = await _seed_target(prisma, world, shape, scratch.prefix)
caller = world.keys[actor]
org_id = await _seed_target(db, actors, shape, resources.prefix)
caller = actors.keys[actor]
# /unblock starts from a blocked row so a 200 is observable as True->False.
if route == "unblock":
await prisma.db.litellm_teamtable.update(
where={"team_id": scratch.prefix}, data={"blocked": True}
await db.db.litellm_teamtable.update(
where={"team_id": resources.prefix}, data={"blocked": True}
)
resp = await proxy_client.post(
resp = await gateway_client.post(
f"/team/{route}",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"team_id": scratch.prefix, "organization_id": org_id},
json={"team_id": resources.prefix, "organization_id": org_id},
)
assert (
resp.status_code == expected_status
), f"{route} {actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:
@ -79,36 +79,36 @@ async def test_team_block_unblock_authz_matrix(
assert bool(row.blocked) is (route == "unblock"), "denied but blocked mutated"
async def test_team_block_unblock_round_trip(proxy_client, prisma, scratch, world):
async def test_team_block_unblock_round_trip(gateway_client, db, resources, actors):
"""PROXY_ADMIN block then unblock flips the blocked column True then False."""
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
headers = {"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
headers = {"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"}
blocked = await proxy_client.post(
"/team/block", headers=headers, json={"team_id": scratch.prefix}
blocked = await gateway_client.post(
"/team/block", headers=headers, json={"team_id": resources.prefix}
)
assert blocked.status_code == 200, blocked.text
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None and row.blocked is True
unblocked = await proxy_client.post(
"/team/unblock", headers=headers, json={"team_id": scratch.prefix}
unblocked = await gateway_client.post(
"/team/unblock", headers=headers, json={"team_id": resources.prefix}
)
assert unblocked.status_code == 200, unblocked.text
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None and row.blocked is False
@pytest.mark.parametrize("route", ["block", "unblock"])
async def test_team_block_unblock_missing_team_is_404(route: str, proxy_client, world):
async def test_team_block_unblock_missing_team_is_404(route: str, gateway_client, actors):
"""A team_id absent from the DB is 404 — the existence check precedes authz."""
resp = await proxy_client.post(
resp = await gateway_client.post(
f"/team/{route}",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": "behavior-pin-no-such-team"},
)
assert resp.status_code == 404, resp.text

142
tests/proxy_behavior/management/test_team_budget_limits.py → tests/integration/management/test_team_budget_limits.py
View file

@ -26,14 +26,14 @@ import pytest
from litellm.proxy.utils import hash_token
from .actors import Actor
from .conftest import MASTER_KEY, create_scratch_org, create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_org, create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
async def _seed_scratch_actor_with_caps(
prisma,
db,
scratch_prefix: str,
*,
user_role: str = "internal_user",
@ -42,7 +42,7 @@ async def _seed_scratch_actor_with_caps(
tpm_limit: Optional[int] = None,
rpm_limit: Optional[int] = None,
) -> str:
"""Raw-seed a scratch actor + verification token, with the token carrying
"""Raw-seed a resources actor + verification token, with the token carrying
explicit caps that flow into `user_api_key_dict` at request time.
Returns cleartext key. Used by F3 user-team-limit scenarios where the
@ -54,7 +54,7 @@ async def _seed_scratch_actor_with_caps(
"""
user_id = f"{scratch_prefix}-team-creator"
cleartext = "sk-" + uuid.uuid4().hex
await prisma.db.litellm_usertable.create(
await db.db.litellm_usertable.create(
data={
"user_id": user_id,
"user_role": user_role,
@ -73,7 +73,7 @@ async def _seed_scratch_actor_with_caps(
token_data["tpm_limit"] = tpm_limit
if rpm_limit is not None:
token_data["rpm_limit"] = rpm_limit
await prisma.db.litellm_verificationtoken.create(data=token_data)
await db.db.litellm_verificationtoken.create(data=token_data)
return cleartext
@ -112,28 +112,28 @@ async def test_check_org_team_limits_models_subset(
org_models,
body_extras: Dict[str, Any],
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
org_id = await create_scratch_org(prisma, scratch.prefix, models=org_models)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
team_id = scratch.tag("team")
org_id = await create_scratch_org(db, resources.prefix, models=org_models)
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
team_id = resources.tag("team")
body: Dict[str, Any] = {
"team_id": team_id,
"team_alias": scratch.prefix,
"team_alias": resources.prefix,
"organization_id": org_id,
**body_extras,
}
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
)
assert resp.status_code == expected_status, f"{body!r} → {resp.status_code}: {resp.text}"
rows = await prisma.db.litellm_teamtable.find_many(where={"team_id": team_id})
rows = await db.db.litellm_teamtable.find_many(where={"team_id": team_id})
assert len(rows) == (1 if expected_status == 200 else 0)
@ -191,26 +191,26 @@ async def test_check_org_team_limits_budget_enforced(
org_budget,
body_extras: Dict[str, Any],
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
org_id = await create_scratch_org(prisma, scratch.prefix, **org_budget)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
team_id = scratch.tag("team")
resp = await proxy_client.post(
org_id = await create_scratch_org(db, resources.prefix, **org_budget)
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
team_id = resources.tag("team")
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {seeder}"},
json={
"team_id": team_id,
"team_alias": scratch.prefix,
"team_alias": resources.prefix,
"organization_id": org_id,
**body_extras,
},
)
assert resp.status_code == expected_status, f"{body_extras!r} → {resp.status_code}: {resp.text}"
rows = await prisma.db.litellm_teamtable.find_many(where={"team_id": team_id})
rows = await db.db.litellm_teamtable.find_many(where={"team_id": team_id})
assert len(rows) == (1 if expected_status == 200 else 0)
@ -223,21 +223,21 @@ async def test_check_org_team_limits_budget_enforced_on_update(
org_budget,
body_extras: Dict[str, Any],
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
org_id = await create_scratch_org(prisma, scratch.prefix, **org_budget)
team_id = await create_scratch_team(prisma, scratch.tag("team"), organization_id=org_id)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
org_id = await create_scratch_org(db, resources.prefix, **org_budget)
team_id = await create_scratch_team(db, resources.tag("team"), organization_id=org_id)
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {seeder}"},
json={"team_id": team_id, **body_extras},
)
assert resp.status_code == expected_status, f"{body_extras!r} → {resp.status_code}: {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
persisted = {field: getattr(row, field) for field in body_extras}
assert (persisted == body_extras) == (expected_status == 200)
@ -309,18 +309,18 @@ async def test_check_user_team_limits(
actor_caps,
body_extras: Dict[str, Any],
expected_status: int,
proxy_client,
prisma,
scratch,
gateway_client,
db,
resources,
):
caller = await _seed_scratch_actor_with_caps(prisma, scratch.prefix, **actor_caps)
team_id = scratch.tag("team")
resp = await proxy_client.post(
caller = await _seed_scratch_actor_with_caps(db, resources.prefix, **actor_caps)
team_id = resources.tag("team")
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {caller}"},
json={
"team_id": team_id,
"team_alias": scratch.prefix,
"team_alias": resources.prefix,
# Standalone team — no organization_id, so user-limit guard fires.
**body_extras,
},
@ -329,7 +329,7 @@ async def test_check_user_team_limits(
f"caps={actor_caps} body={body_extras} → {resp.status_code}: {resp.text}"
)
rows = await prisma.db.litellm_teamtable.find_many(where={"team_id": team_id})
rows = await db.db.litellm_teamtable.find_many(where={"team_id": team_id})
assert len(rows) == (1 if expected_status == 200 else 0)
@ -351,16 +351,16 @@ async def test_check_user_team_limits(
ids=["raise_with_generous_personal_budget", "lower_with_tiny_personal_budget"],
)
async def test_team_admin_cannot_change_budget_while_max_budget_is_not_editable(
proxy_client, prisma, scratch, personal_budget: float, requested_budget: float
gateway_client, db, resources, personal_budget: float, requested_budget: float
):
caller_cleartext = await _seed_scratch_actor_with_caps(prisma, scratch.prefix, max_budget=personal_budget)
caller_cleartext = await _seed_scratch_actor_with_caps(db, resources.prefix, max_budget=personal_budget)
team_id = await create_scratch_team(
prisma,
team_id=scratch.tag("team"),
admin_user_ids=[f"{scratch.prefix}-team-creator"],
db,
team_id=resources.tag("team"),
admin_user_ids=[f"{resources.prefix}-team-creator"],
max_budget=500.0,
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {caller_cleartext}"},
json={"team_id": team_id, "max_budget": requested_budget},
@ -368,69 +368,69 @@ async def test_team_admin_cannot_change_budget_while_max_budget_is_not_editable(
assert resp.status_code == 403, resp.text
assert "Team admin editable fields" in resp.text, resp.text
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
assert row.max_budget == 500.0, "a refused update must leave the team budget unchanged"
async def test_proxy_admin_raise_budget_allowed(proxy_client, prisma, scratch):
async def test_proxy_admin_raise_budget_allowed(gateway_client, db, resources):
"""A proxy admin may raise a team's budget."""
team_id = await create_scratch_team(
prisma,
team_id=scratch.tag("team"),
admin_user_ids=[f"{scratch.prefix}-team-creator"],
db,
team_id=resources.tag("team"),
admin_user_ids=[f"{resources.prefix}-team-creator"],
max_budget=50.0,
)
# MASTER_KEY acts as proxy admin.
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {MASTER_KEY}"},
headers={"Authorization": f"Bearer {gateway_client.gateway.key}"},
json={"team_id": team_id, "max_budget": 999.0},
)
assert resp.status_code == 200, resp.text
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
assert row.max_budget == 999.0, "proxy admin should be able to raise the budget"
async def test_team_admin_remove_budget_cap_blocked(proxy_client, prisma, scratch):
async def test_team_admin_remove_budget_cap_blocked(gateway_client, db, resources):
"""A team admin cannot strip the team's cap (max_budget=null); removing the
ceiling is the strongest possible raise -> proxy-admin only."""
caller_cleartext = await _seed_scratch_actor_with_caps(prisma, scratch.prefix, max_budget=100000.0)
caller_cleartext = await _seed_scratch_actor_with_caps(db, resources.prefix, max_budget=100000.0)
team_id = await create_scratch_team(
prisma,
team_id=scratch.tag("team"),
admin_user_ids=[f"{scratch.prefix}-team-creator"],
db,
team_id=resources.tag("team"),
admin_user_ids=[f"{resources.prefix}-team-creator"],
max_budget=50.0,
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {caller_cleartext}"},
json={"team_id": team_id, "max_budget": None},
)
assert resp.status_code == 403, resp.text
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
assert row.max_budget == 50.0, "team budget cap must not be removed by a team admin"
async def test_proxy_admin_remove_budget_cap_allowed(proxy_client, prisma, scratch):
async def test_proxy_admin_remove_budget_cap_allowed(gateway_client, db, resources):
"""A proxy admin may remove a team's cap (max_budget=null)."""
team_id = await create_scratch_team(
prisma,
team_id=scratch.tag("team"),
admin_user_ids=[f"{scratch.prefix}-team-creator"],
db,
team_id=resources.tag("team"),
admin_user_ids=[f"{resources.prefix}-team-creator"],
max_budget=50.0,
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {MASTER_KEY}"},
headers={"Authorization": f"Bearer {gateway_client.gateway.key}"},
json={"team_id": team_id, "max_budget": None},
)
assert resp.status_code == 200, resp.text
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
assert row.max_budget is None, "proxy admin should be able to remove the cap"

105
tests/integration/management/test_team_bulk_member_add.py Normal file
View file

@ -0,0 +1,105 @@
import pytest
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
def _member_ids(row) -> list:
return [m["user_id"] for m in (row.members_with_roles or [])]
async def test_team_bulk_member_add_proxy_admin_adds_explicit_members(
gateway_client, db, resources, actors
):
"""PROXY_ADMIN bulk-adds an explicit member list to a resources team."""
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
new_member = resources.tag("m1")
resp = await gateway_client.post(
"/team/bulk_member_add",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={
"team_id": resources.prefix,
"members": [{"user_id": new_member, "role": "user"}],
},
)
assert resp.status_code == 200, resp.text
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None and new_member in _member_ids(row)
async def test_team_bulk_member_add_empty_members_is_400(
gateway_client, db, resources, actors
):
"""An empty member list (with all_users unset) is rejected 400."""
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
resp = await gateway_client.post(
"/team/bulk_member_add",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "members": []},
)
assert resp.status_code == 400, resp.text
async def test_team_bulk_member_add_over_max_batch_is_400(
gateway_client, db, resources, actors
):
"""A member list larger than the 500-member cap is rejected 400."""
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
members = [
{"user_id": f"{resources.prefix}-u{i}", "role": "user"} for i in range(501)
]
resp = await gateway_client.post(
"/team/bulk_member_add",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "members": members},
)
assert resp.status_code == 400, resp.text
@pytest.mark.parametrize(
"actor",
[Actor.TEAM_ADMIN, Actor.INTERNAL_USER],
ids=["team_admin", "internal_user"],
)
async def test_team_bulk_member_add_non_admin_is_401(
actor: Actor, gateway_client, db, resources, actors
):
"""/team/bulk_member_add is neither an internal_user nor a self-managed
route — a non-proxy-admin with no org context is 401 at the route gate."""
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
resp = await gateway_client.post(
"/team/bulk_member_add",
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
json={
"team_id": resources.prefix,
"members": [{"user_id": resources.tag("m"), "role": "user"}],
},
)
assert resp.status_code == 401, f"{actor.value}: {resp.status_code} {resp.text}"
async def test_team_bulk_member_add_all_users_proxy_admin(
gateway_client, db, resources, actors
):
"""all_users=True pulls every user in the DB into the team. The route is
reachable only by PROXY_ADMIN (the route gate 401s every other actor — even
an org admin with organization_id in the body), so the handler's own
all_users PROXY_ADMIN gate is never the deciding check at the boundary."""
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
resp = await gateway_client.post(
"/team/bulk_member_add",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "all_users": True},
)
assert resp.status_code == 200, resp.text
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
member_ids = _member_ids(row)
# every actors actor is a user in the DB, so all are now team members
assert actors.keys[Actor.INTERNAL_USER].user_id in member_ids

140
tests/proxy_behavior/management/test_team_bulk_member_delete.py → tests/integration/management/test_team_bulk_member_delete.py
View file

@ -1,8 +1,8 @@
import pytest
from prisma import Json
from .actors import Actor
from .conftest import create_scratch_team, create_scratch_user
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team, create_scratch_user
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -28,20 +28,20 @@ _MATRIX = [
]
async def _seed_target(prisma, world, shape: str, team_id: str, victim_ids: list) -> None:
async def _seed_target(db, actors, shape: str, team_id: str, victim_ids: list) -> None:
if shape == "alpha":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_a_id,
admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id],
organization_id=actors.org_a_id,
admin_user_ids=[actors.keys[Actor.TEAM_ADMIN].user_id],
member_user_ids=victim_ids,
)
elif shape == "beta":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_b_id,
organization_id=actors.org_b_id,
member_user_ids=victim_ids,
)
else: # pragma: no cover - guard
@ -61,18 +61,18 @@ async def test_team_bulk_member_delete_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
victims = [scratch.tag("v1"), scratch.tag("v2")]
keep = scratch.tag("keep")
await _seed_target(prisma, world, shape, scratch.prefix, victims + [keep])
caller = world.keys[actor]
victims = [resources.tag("v1"), resources.tag("v2")]
keep = resources.tag("keep")
await _seed_target(db, actors, shape, resources.prefix, victims + [keep])
caller = actors.keys[actor]
resp = await proxy_client.post(
f"/management/v1/teams/{scratch.prefix}/members/bulk_delete",
resp = await gateway_client.post(
f"/management/v1/teams/{resources.prefix}/members/bulk_delete",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"members": [{"user_id": v} for v in victims]},
)
@ -81,7 +81,7 @@ async def test_team_bulk_member_delete_authz_matrix(
assert resp.headers["content-type"] == "application/problem+json"
assert resp.json()["type"] == "urn:litellm:error:forbidden"
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert row is not None
assert keep in _member_ids(row), "unrelated member removed"
if expected_status == 200:
@ -91,15 +91,15 @@ async def test_team_bulk_member_delete_authz_matrix(
assert set(victims) <= set(_member_ids(row)), "denied but members removed"
async def test_team_bulk_member_delete_reports_each_row_in_order(proxy_client, prisma, scratch, world):
victim = scratch.tag("victim")
keep = scratch.tag("keep")
stranger = scratch.tag("stranger")
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id, member_user_ids=[victim, keep])
async def test_team_bulk_member_delete_reports_each_row_in_order(gateway_client, db, resources, actors):
victim = resources.tag("victim")
keep = resources.tag("keep")
stranger = resources.tag("stranger")
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id, member_user_ids=[victim, keep])
resp = await proxy_client.post(
f"/management/v1/teams/{scratch.prefix}/members/bulk_delete",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
resp = await gateway_client.post(
f"/management/v1/teams/{resources.prefix}/members/bulk_delete",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"members": [{"user_id": stranger}, {"user_id": victim}]},
)
assert resp.status_code == 200, resp.text
@ -111,47 +111,47 @@ async def test_team_bulk_member_delete_reports_each_row_in_order(proxy_client, p
]
assert body["data"][0]["error"] == "User not found in team"
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
assert row is not None and _member_ids(row) == [keep]
row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert row is not None and set(_member_ids(row)) == {keep, "default_user_id"}
async def test_team_bulk_member_delete_by_id_removes_a_legacy_email_only_roster_entry(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
email = f"{scratch.prefix}@example.com"
victim = await create_scratch_user(prisma, scratch.prefix, suffix="victim", user_email=email)
keep = scratch.tag("keep")
await prisma.db.litellm_teamtable.create(
email = f"{resources.prefix}@example.com"
victim = await create_scratch_user(db, resources.prefix, suffix="victim", user_email=email)
keep = resources.tag("keep")
await db.db.litellm_teamtable.create(
data={
"team_id": scratch.prefix,
"team_alias": scratch.prefix,
"organization_id": world.org_a_id,
"team_id": resources.prefix,
"team_alias": resources.prefix,
"organization_id": actors.org_a_id,
"members_with_roles": Json([{"user_email": email, "role": "user"}, {"user_id": keep, "role": "user"}]),
}
)
await prisma.db.litellm_usertable.update(where={"user_id": victim}, data={"teams": [scratch.prefix]})
await db.db.litellm_usertable.update(where={"user_id": victim}, data={"teams": [resources.prefix]})
resp = await proxy_client.post(
f"/management/v1/teams/{scratch.prefix}/members/bulk_delete",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
resp = await gateway_client.post(
f"/management/v1/teams/{resources.prefix}/members/bulk_delete",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"members": [{"user_id": victim}]},
)
assert resp.status_code == 200, resp.text
assert [(r["user_id"], r["success"]) for r in resp.json()["data"]] == [(victim, True)]
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert row is not None and [(m["user_id"], m.get("user_email")) for m in row.members_with_roles] == [(keep, None)]
user = await prisma.db.litellm_usertable.find_unique(where={"user_id": victim})
user = await db.db.litellm_usertable.find_unique(where={"user_id": victim})
assert user is not None and user.teams == []
async def test_team_bulk_member_delete_row_naming_both_identifiers_is_422(proxy_client, prisma, scratch, world):
victim = scratch.tag("victim")
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id, member_user_ids=[victim])
async def test_team_bulk_member_delete_row_naming_both_identifiers_is_422(gateway_client, db, resources, actors):
victim = resources.tag("victim")
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id, member_user_ids=[victim])
resp = await proxy_client.post(
f"/management/v1/teams/{scratch.prefix}/members/bulk_delete",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
resp = await gateway_client.post(
f"/management/v1/teams/{resources.prefix}/members/bulk_delete",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"members": [{"user_id": victim, "user_email": f"{victim}@example.com"}]},
)
assert resp.status_code == 422, resp.text
@ -162,48 +162,48 @@ async def test_team_bulk_member_delete_row_naming_both_identifiers_is_422(proxy_
== "members.0: Value error, Each member must be identified by exactly one of user_id or user_email"
)
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert row is not None and victim in _member_ids(row)
async def test_team_bulk_member_delete_unknown_query_param_is_400(proxy_client, prisma, scratch, world):
victim = scratch.tag("victim")
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id, member_user_ids=[victim])
async def test_team_bulk_member_delete_unknown_query_param_is_400(gateway_client, db, resources, actors):
victim = resources.tag("victim")
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id, member_user_ids=[victim])
resp = await proxy_client.post(
f"/management/v1/teams/{scratch.prefix}/members/bulk_delete?dry_run=1",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
resp = await gateway_client.post(
f"/management/v1/teams/{resources.prefix}/members/bulk_delete?dry_run=1",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"members": [{"user_id": victim}]},
)
assert resp.status_code == 400, resp.text
assert resp.headers["content-type"] == "application/problem+json"
assert "dry_run" in resp.json()["detail"]
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert row is not None and victim in _member_ids(row)
async def test_team_bulk_member_delete_unknown_body_field_is_422(proxy_client, prisma, scratch, world):
victim = scratch.tag("victim")
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id, member_user_ids=[victim])
async def test_team_bulk_member_delete_unknown_body_field_is_422(gateway_client, db, resources, actors):
victim = resources.tag("victim")
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id, member_user_ids=[victim])
resp = await proxy_client.post(
f"/management/v1/teams/{scratch.prefix}/members/bulk_delete",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": scratch.prefix, "members": [{"user_id": victim}]},
resp = await gateway_client.post(
f"/management/v1/teams/{resources.prefix}/members/bulk_delete",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "members": [{"user_id": victim}]},
)
assert resp.status_code == 422, resp.text
assert "team_id" in resp.json()["detail"]
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert row is not None and victim in _member_ids(row)
async def test_team_bulk_member_delete_unknown_team_is_404_problem(proxy_client, scratch, world):
resp = await proxy_client.post(
f"/management/v1/teams/{scratch.tag('missing')}/members/bulk_delete",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"members": [{"user_id": scratch.tag("victim")}]},
async def test_team_bulk_member_delete_unknown_team_is_404_problem(gateway_client, resources, actors):
resp = await gateway_client.post(
f"/management/v1/teams/{resources.tag('missing')}/members/bulk_delete",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"members": [{"user_id": resources.tag("victim")}]},
)
assert resp.status_code == 404, resp.text
assert resp.headers["content-type"] == "application/problem+json"

32
tests/proxy_behavior/management/test_team_daily_activity.py → tests/integration/management/test_team_daily_activity.py
View file

@ -1,6 +1,9 @@
from typing import Final
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
from tests.integration._support.database import write_rows
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -41,6 +44,21 @@ _CASES = [
_DATES = "start_date=2024-01-01&end_date=2024-12-31"
_SEEDED_SPEND: Final = {"alpha": 1.25, "beta": 2.5}
@pytest.fixture(scope="module")
def seeded_team_spend(actors) -> None:
team_ids: Final = {"alpha": actors.team_alpha_id, "beta": actors.team_beta_id}
for team, spend in _SEEDED_SPEND.items():
write_rows(
'INSERT INTO "LiteLLM_DailyTeamSpend" (id, team_id, date, api_key, model, spend, updated_at) '
"VALUES (%s, %s, '2024-06-01', 'integration-daily-activity', 'gpt-4o', %s, now()) "
"ON CONFLICT (id) DO NOTHING",
(f"integration-daily-activity-{team}", team_ids[team], str(spend)),
)
@pytest.mark.parametrize(
"endpoint",
("/team/daily/activity", "/team/daily/activity/aggregated"),
@ -52,18 +70,20 @@ _DATES = "start_date=2024-01-01&end_date=2024-12-31"
ids=[c[0] for c in _CASES],
)
async def test_team_daily_activity_matrix(
actor: Actor, team: str, expected_status: int, endpoint: str, proxy_client, world
actor: Actor, team: str, expected_status: int, endpoint: str, gateway_client, actors, seeded_team_spend
):
query = _DATES
if team == "alpha":
query += f"&team_ids={world.team_alpha_id}"
query += f"&team_ids={actors.team_alpha_id}"
elif team == "beta":
query += f"&team_ids={world.team_beta_id}"
query += f"&team_ids={actors.team_beta_id}"
resp = await proxy_client.get(
resp = await gateway_client.get(
f"{endpoint}?{query}",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
)
assert (
resp.status_code == expected_status
), f"{actor.value} -> {team}: {resp.status_code} {resp.text}"
if actor == Actor.PROXY_ADMIN and team in _SEEDED_SPEND:
assert resp.json()["metadata"]["total_spend"] == _SEEDED_SPEND[team], resp.text

40
tests/proxy_behavior/management/test_team_delete.py → tests/integration/management/test_team_delete.py
View file

@ -1,7 +1,7 @@
import pytest
from .actors import Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -33,26 +33,26 @@ async def test_team_delete_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
org_id = world.org_a_id if shape == "alpha" else world.org_b_id
await create_scratch_team(prisma, scratch.prefix, organization_id=org_id)
caller = world.keys[actor]
org_id = actors.org_a_id if shape == "alpha" else actors.org_b_id
await create_scratch_team(db, resources.prefix, organization_id=org_id)
caller = actors.keys[actor]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/delete",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"team_ids": [scratch.prefix], "organization_id": org_id},
json={"team_ids": [resources.prefix], "organization_id": org_id},
)
assert (
resp.status_code == expected_status
), f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
if expected_status == 200:
assert row is None, "deleted but team row survives"
@ -61,18 +61,18 @@ async def test_team_delete_authz_matrix(
async def test_team_delete_batch_with_missing_id_deletes_nothing(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""A batch is validated whole before any deletion: one missing team_id
fails the request 404 and the accessible team in the batch survives."""
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
resp = await proxy_client.post(
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
resp = await gateway_client.post(
"/team/delete",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_ids": [scratch.prefix, "behavior-pin-no-such-team"]},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_ids": [resources.prefix, "behavior-pin-no-such-team"]},
)
assert resp.status_code == 404, resp.text
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None, "batch aborted but the accessible team was deleted"

22
tests/proxy_behavior/management/test_team_filter_ui.py → tests/integration/management/test_team_filter_ui.py
View file

@ -1,6 +1,6 @@
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -12,10 +12,10 @@ pytestmark = pytest.mark.asyncio(loop_scope="session")
# route) and 401s every non-proxy-admin before the handler runs. PROXY_ADMIN
# reaches the unscoped find_many and sees teams across every org.
@pytest.mark.parametrize("actor", list(Actor), ids=[a.value for a in Actor])
async def test_team_filter_ui_is_proxy_admin_only(actor: Actor, proxy_client, world):
resp = await proxy_client.get(
async def test_team_filter_ui_is_proxy_admin_only(actor: Actor, gateway_client, actors):
resp = await gateway_client.get(
"/team/filter/ui",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
)
expected = 200 if actor == Actor.PROXY_ADMIN else 401
assert (
@ -23,17 +23,17 @@ async def test_team_filter_ui_is_proxy_admin_only(actor: Actor, proxy_client, wo
), f"{actor.value}: {resp.status_code} {resp.text}"
async def test_team_filter_ui_proxy_admin_sees_cross_org_teams(proxy_client, world):
async def test_team_filter_ui_proxy_admin_sees_cross_org_teams(gateway_client, actors):
"""The handler runs an unscoped query — PROXY_ADMIN sees teams from every
org, including the three seeded world teams."""
resp = await proxy_client.get(
org, including the three seeded actors teams."""
resp = await gateway_client.get(
"/team/filter/ui",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 200, resp.text
team_ids = {t.get("team_id") for t in resp.json() if isinstance(t, dict)}
assert {
world.team_alpha_id,
world.team_beta_id,
world.team_gamma_id,
actors.team_alpha_id,
actors.team_beta_id,
actors.team_gamma_id,
} <= team_ids

22
tests/proxy_behavior/management/test_team_info.py → tests/integration/management/test_team_info.py
View file

@ -1,6 +1,6 @@
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -47,16 +47,16 @@ _SCENARIOS = [
ids=[s[0] for s in _SCENARIOS],
)
async def test_team_info_authz_matrix(
actor: Actor, target: str, expected_status: int, proxy_client, world
actor: Actor, target: str, expected_status: int, gateway_client, actors
):
caller = world.keys[actor]
caller = actors.keys[actor]
target_team_id = {
"alpha": world.team_alpha_id,
"gamma": world.team_gamma_id,
"beta": world.team_beta_id,
"alpha": actors.team_alpha_id,
"gamma": actors.team_gamma_id,
"beta": actors.team_beta_id,
}[target]
resp = await proxy_client.get(
resp = await gateway_client.get(
f"/team/info?team_id={target_team_id}",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)
@ -74,10 +74,10 @@ async def test_team_info_authz_matrix(
# alpha/org_b_admin already covers the branch in the matrix; this guard
# turns a silent rename of the exception detail into a CI red, which is the
# behavior tripwire that the matrix's status-only assertion cannot catch.
async def test_team_info_org_admin_cross_org_rejection_detail(proxy_client, world):
resp = await proxy_client.get(
f"/team/info?team_id={world.team_alpha_id}",
headers={"Authorization": f"Bearer {world.keys[Actor.ORG_B_ADMIN].cleartext}"},
async def test_team_info_org_admin_cross_org_rejection_detail(gateway_client, actors):
resp = await gateway_client.get(
f"/team/info?team_id={actors.team_alpha_id}",
headers={"Authorization": f"Bearer {actors.keys[Actor.ORG_B_ADMIN].cleartext}"},
)
assert resp.status_code == 403, resp.text
# validate_membership() at GET /team/info raises with this exact phrase.

96
tests/proxy_behavior/management/test_team_key_bulk_update.py → tests/integration/management/test_team_key_bulk_update.py
View file

@ -5,8 +5,8 @@ import pytest
from litellm.proxy._types import KeyManagementRoutes
from litellm.proxy.utils import hash_token
from .actors import Actor
from .conftest import create_scratch_key, create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_key, create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -31,20 +31,20 @@ _MATRIX = [
]
async def _seed_team_key(prisma, proxy_client, prefix: str, world, shape: str) -> str:
"""Raw-seed the scratch team for `shape`, return a team key's cleartext."""
internal = world.keys[Actor.INTERNAL_USER].user_id
owner = world.keys[Actor.OWNER].user_id
async def _seed_team_key(db, gateway_client, prefix: str, actors, shape: str) -> str:
"""Raw-seed the resources team for `shape`, return a team key's cleartext."""
internal = actors.keys[Actor.INTERNAL_USER].user_id
owner = actors.keys[Actor.OWNER].user_id
if shape == "admin":
await create_scratch_team(
prisma, prefix, organization_id=world.org_a_id, admin_user_ids=[internal]
db, prefix, organization_id=actors.org_a_id, admin_user_ids=[internal]
)
key_owner = internal
elif shape == "member_allowed":
await create_scratch_team(
prisma,
db,
prefix,
organization_id=world.org_a_id,
organization_id=actors.org_a_id,
admin_user_ids=[owner],
member_user_ids=[internal],
team_member_permissions=[_KEY_UPDATE],
@ -52,9 +52,9 @@ async def _seed_team_key(prisma, proxy_client, prefix: str, world, shape: str) -
key_owner = owner
elif shape == "member_denied":
await create_scratch_team(
prisma,
db,
prefix,
organization_id=world.org_a_id,
organization_id=actors.org_a_id,
admin_user_ids=[owner],
member_user_ids=[internal],
team_member_permissions=[],
@ -62,14 +62,14 @@ async def _seed_team_key(prisma, proxy_client, prefix: str, world, shape: str) -
key_owner = owner
elif shape == "nonmember":
await create_scratch_team(
prisma, prefix, organization_id=world.org_a_id, admin_user_ids=[owner]
db, prefix, organization_id=actors.org_a_id, admin_user_ids=[owner]
)
key_owner = owner
else:
pytest.fail(f"unknown shape={shape}") # pragma: no cover
return await create_scratch_key(
proxy_client,
world.keys[Actor.PROXY_ADMIN].cleartext,
gateway_client,
actors.keys[Actor.PROXY_ADMIN].cleartext,
prefix,
user_id=key_owner,
team_id=prefix,
@ -85,20 +85,20 @@ async def test_team_key_bulk_update_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
key = await _seed_team_key(prisma, proxy_client, scratch.prefix, world, shape)
key = await _seed_team_key(db, gateway_client, resources.prefix, actors, shape)
hashed = hash_token(key)
caller = world.keys[actor]
caller = actors.keys[actor]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/key/bulk_update",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={
"team_id": scratch.prefix,
"team_id": resources.prefix,
"key_ids": [key],
"update_fields": {"max_budget": _MARKER_BUDGET},
},
@ -107,7 +107,7 @@ async def test_team_key_bulk_update_authz_matrix(
resp.status_code == expected_status
), f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_verificationtoken.find_unique(where={"token": hashed})
row = await db.db.litellm_verificationtoken.find_unique(where={"token": hashed})
assert row is not None
if expected_status == 200:
assert len(resp.json()["successful_updates"]) == 1
@ -117,12 +117,12 @@ async def test_team_key_bulk_update_authz_matrix(
async def test_team_key_bulk_update_requires_team_id(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""An empty team_id is rejected 400."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/key/bulk_update",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={
"team_id": "",
"key_ids": ["sk-" + uuid.uuid4().hex],
@ -133,28 +133,28 @@ async def test_team_key_bulk_update_requires_team_id(
async def test_team_key_bulk_update_all_keys_in_team(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""all_keys_in_team=True broadcasts the update to every key in the team."""
admin = world.keys[Actor.PROXY_ADMIN].cleartext
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
admin = actors.keys[Actor.PROXY_ADMIN].cleartext
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
keys = [
await create_scratch_key(
proxy_client,
gateway_client,
admin,
scratch.prefix,
user_id=world.keys[Actor.OWNER].user_id,
team_id=scratch.prefix,
key_alias=f"{scratch.prefix}-k{i}",
resources.prefix,
user_id=actors.keys[Actor.OWNER].user_id,
team_id=resources.prefix,
key_alias=f"{resources.prefix}-k{i}",
)
for i in range(2)
]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/key/bulk_update",
headers={"Authorization": f"Bearer {admin}"},
json={
"team_id": scratch.prefix,
"team_id": resources.prefix,
"all_keys_in_team": True,
"update_fields": {"max_budget": _MARKER_BUDGET},
},
@ -162,22 +162,22 @@ async def test_team_key_bulk_update_all_keys_in_team(
assert resp.status_code == 200, resp.text
assert len(resp.json()["successful_updates"]) == 2
for key in keys:
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(key)}
)
assert row is not None and row.max_budget == _MARKER_BUDGET
async def test_team_key_bulk_update_no_keys_found_is_404(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""all_keys_in_team=True on a team with no keys is a top-level 404."""
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
resp = await proxy_client.post(
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
resp = await gateway_client.post(
"/team/key/bulk_update",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={
"team_id": scratch.prefix,
"team_id": resources.prefix,
"all_keys_in_team": True,
"update_fields": {"max_budget": _MARKER_BUDGET},
},
@ -186,21 +186,21 @@ async def test_team_key_bulk_update_no_keys_found_is_404(
async def test_team_key_bulk_update_missing_key_is_isolated(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""A key_id absent from the team lands in failed_updates; the batch still
returns 200 and the real key is updated."""
admin = world.keys[Actor.PROXY_ADMIN].cleartext
admin = actors.keys[Actor.PROXY_ADMIN].cleartext
real = await _seed_team_key(
prisma, proxy_client, scratch.prefix, world, "nonmember"
db, gateway_client, resources.prefix, actors, "nonmember"
)
missing = "sk-" + uuid.uuid4().hex
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/key/bulk_update",
headers={"Authorization": f"Bearer {admin}"},
json={
"team_id": scratch.prefix,
"team_id": resources.prefix,
"key_ids": [real, missing],
"update_fields": {"max_budget": _MARKER_BUDGET},
},
@ -211,7 +211,7 @@ async def test_team_key_bulk_update_missing_key_is_isolated(
assert len(body["successful_updates"]) == 1
assert len(body["failed_updates"]) == 1
row = await prisma.db.litellm_verificationtoken.find_unique(
row = await db.db.litellm_verificationtoken.find_unique(
where={"token": hash_token(real)}
)
assert row is not None and row.max_budget == _MARKER_BUDGET

28
tests/proxy_behavior/management/test_team_list.py → tests/integration/management/test_team_list.py
View file

@ -2,18 +2,18 @@ from typing import FrozenSet, Optional
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
# The behavior DB may hold teams beyond the three seeded ones, so every
# assertion intersects the returned team_ids with the known seeded set.
def _seeded_visible(resp_json, world) -> set:
def _seeded_visible(resp_json, actors) -> set:
known = {
world.team_alpha_id: "alpha",
world.team_beta_id: "beta",
world.team_gamma_id: "gamma",
actors.team_alpha_id: "alpha",
actors.team_beta_id: "beta",
actors.team_gamma_id: "gamma",
}
return {
known[entry["team_id"]]
@ -47,11 +47,11 @@ async def test_team_list_bare_authz(
actor: Actor,
expected_status: int,
expected_visible: Optional[set],
proxy_client,
world,
gateway_client,
actors,
):
caller = world.keys[actor]
resp = await proxy_client.get(
caller = actors.keys[actor]
resp = await gateway_client.get(
"/team/list",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)
@ -60,7 +60,7 @@ async def test_team_list_bare_authz(
), f"{actor.value}: {resp.status_code} {resp.text}"
if expected_status == 200:
visible = _seeded_visible(resp.json(), world)
visible = _seeded_visible(resp.json(), actors)
assert visible == expected_visible, (
f"{actor.value}: expected {sorted(expected_visible)}, "
f"got {sorted(visible)}"
@ -90,16 +90,16 @@ _OWN = {
ids=[a.value for a in _OWN],
)
async def test_team_list_own_query(
actor: Actor, expected_visible: FrozenSet[str], proxy_client, world
actor: Actor, expected_visible: FrozenSet[str], gateway_client, actors
):
caller = world.keys[actor]
resp = await proxy_client.get(
caller = actors.keys[actor]
resp = await gateway_client.get(
f"/team/list?user_id={caller.user_id}",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)
assert resp.status_code == 200, f"{actor.value}: {resp.status_code} {resp.text}"
visible = _seeded_visible(resp.json(), world)
visible = _seeded_visible(resp.json(), actors)
assert visible == set(expected_visible), (
f"{actor.value}: expected {sorted(expected_visible)}, " f"got {sorted(visible)}"
)

54
tests/proxy_behavior/management/test_team_list_v2.py → tests/integration/management/test_team_list_v2.py
View file

@ -2,26 +2,26 @@ from typing import FrozenSet, Optional
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
def _seeded(team_ids: set, world) -> set:
def _seeded(team_ids: set, actors) -> set:
known = {
world.team_alpha_id: "alpha",
world.team_beta_id: "beta",
world.team_gamma_id: "gamma",
actors.team_alpha_id: "alpha",
actors.team_beta_id: "beta",
actors.team_gamma_id: "gamma",
}
return {known[t] for t in team_ids if t in known}
async def _v2_team_ids(proxy_client, caller_cleartext: str, extra: str = "") -> set:
async def _v2_team_ids(gateway_client, caller_cleartext: str, extra: str = "") -> set:
"""Walk every /v2/team/list page and collect the returned team_ids."""
ids: set = set()
page = 1
while True:
resp = await proxy_client.get(
resp = await gateway_client.get(
f"/v2/team/list?page={page}&page_size=100{extra}",
headers={"Authorization": f"Bearer {caller_cleartext}"},
)
@ -64,19 +64,19 @@ async def test_team_list_v2_bare(
actor: Actor,
expected_status: int,
expected_visible: Optional[FrozenSet[str]],
proxy_client,
world,
gateway_client,
actors,
):
caller = world.keys[actor]
caller = actors.keys[actor]
if expected_status != 200:
resp = await proxy_client.get(
resp = await gateway_client.get(
"/v2/team/list",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)
assert resp.status_code == expected_status, resp.text
return
visible = _seeded(await _v2_team_ids(proxy_client, caller.cleartext), world)
visible = _seeded(await _v2_team_ids(gateway_client, caller.cleartext), actors)
assert visible == set(
expected_visible
), f"{actor.value}: expected {sorted(expected_visible)}, got {sorted(visible)}"
@ -98,44 +98,44 @@ _OWN = {
"actor,expected_visible", list(_OWN.items()), ids=[a.value for a in _OWN]
)
async def test_team_list_v2_own_user_id_query(
actor: Actor, expected_visible: FrozenSet[str], proxy_client, world
actor: Actor, expected_visible: FrozenSet[str], gateway_client, actors
):
caller = world.keys[actor]
caller = actors.keys[actor]
visible = _seeded(
await _v2_team_ids(
proxy_client, caller.cleartext, f"&user_id={caller.user_id}"
gateway_client, caller.cleartext, f"&user_id={caller.user_id}"
),
world,
actors,
)
assert visible == set(
expected_visible
), f"{actor.value}: expected {sorted(expected_visible)}, got {sorted(visible)}"
async def test_team_list_v2_user_id_filter_other_user_is_401(proxy_client, world):
async def test_team_list_v2_user_id_filter_other_user_is_401(gateway_client, actors):
"""A regular user filtering by another user's user_id is rejected 401."""
resp = await proxy_client.get(
f"/v2/team/list?user_id={world.keys[Actor.OWNER].user_id}",
resp = await gateway_client.get(
f"/v2/team/list?user_id={actors.keys[Actor.OWNER].user_id}",
headers={
"Authorization": f"Bearer {world.keys[Actor.INTERNAL_USER].cleartext}"
"Authorization": f"Bearer {actors.keys[Actor.INTERNAL_USER].cleartext}"
},
)
assert resp.status_code == 401, resp.text
async def test_team_list_v2_org_filter_foreign_org_is_403(proxy_client, world):
async def test_team_list_v2_org_filter_foreign_org_is_403(gateway_client, actors):
"""An org admin filtering by an organization it does not administer is 403."""
resp = await proxy_client.get(
f"/v2/team/list?organization_id={world.org_b_id}",
headers={"Authorization": f"Bearer {world.keys[Actor.ORG_ADMIN].cleartext}"},
resp = await gateway_client.get(
f"/v2/team/list?organization_id={actors.org_b_id}",
headers={"Authorization": f"Bearer {actors.keys[Actor.ORG_ADMIN].cleartext}"},
)
assert resp.status_code == 403, resp.text
async def test_team_list_v2_invalid_status_is_400(proxy_client, world):
async def test_team_list_v2_invalid_status_is_400(gateway_client, actors):
"""status accepts only 'deleted' — any other value is 400."""
resp = await proxy_client.get(
resp = await gateway_client.get(
"/v2/team/list?status=bogus",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 400, resp.text

135
tests/proxy_behavior/management/test_team_member_add.py → tests/integration/management/test_team_member_add.py
View file

@ -1,8 +1,7 @@
import litellm
import pytest
from .actors import Actor
from .conftest import create_scratch_team, create_scratch_user
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team, create_scratch_user, owned_settings_client
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -40,16 +39,16 @@ _MATRIX = [
]
async def _seed_target(prisma, world, shape: str, team_id: str) -> None:
async def _seed_target(db, actors, shape: str, team_id: str) -> None:
if shape == "alpha":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_a_id,
admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id],
organization_id=actors.org_a_id,
admin_user_ids=[actors.keys[Actor.TEAM_ADMIN].user_id],
)
elif shape == "beta":
await create_scratch_team(prisma, team_id, organization_id=world.org_b_id)
await create_scratch_team(db, team_id, organization_id=actors.org_b_id)
else: # pragma: no cover - guard
pytest.fail(f"unknown shape={shape}")
@ -67,22 +66,22 @@ async def test_team_member_add_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
await _seed_target(prisma, world, shape, scratch.prefix)
caller = world.keys[actor]
await _seed_target(db, actors, shape, resources.prefix)
caller = actors.keys[actor]
new_member_id = await create_scratch_user(
prisma, scratch.prefix, suffix="newmember"
db, resources.prefix, suffix="newmember"
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={
"team_id": scratch.prefix,
"team_id": resources.prefix,
"member": {"user_id": new_member_id, "role": "user"},
},
)
@ -90,8 +89,8 @@ async def test_team_member_add_authz_matrix(
resp.status_code == expected_status
), f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:
@ -121,28 +120,28 @@ _UNPROVISIONED = [
async def test_team_member_add_unprovisioned_user_id_is_proxy_admin_only(
actor: Actor,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
await _seed_target(prisma, world, "alpha", scratch.prefix)
caller = world.keys[actor]
await _seed_target(db, actors, "alpha", resources.prefix)
caller = actors.keys[actor]
# Deliberately NOT seeded — create_scratch_user is what the matrix above
# calls, and its absence here is the whole scenario.
unprovisioned_id = scratch.tag("unprovisioned")
unprovisioned_id = resources.tag("unprovisioned")
assert (
await prisma.db.litellm_usertable.find_unique(
await db.db.litellm_usertable.find_unique(
where={"user_id": unprovisioned_id}
)
is None
), "setup: user must not exist yet"
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={
"team_id": scratch.prefix,
"team_id": resources.prefix,
"member": {"user_id": unprovisioned_id, "role": "user"},
},
)
@ -150,11 +149,11 @@ async def test_team_member_add_unprovisioned_user_id_is_proxy_admin_only(
resp.status_code == expected_status
), f"{actor.value}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
created = await prisma.db.litellm_usertable.find_unique(
created = await db.db.litellm_usertable.find_unique(
where={"user_id": unprovisioned_id}
)
if expected_status == 200:
@ -167,10 +166,10 @@ async def test_team_member_add_unprovisioned_user_id_is_proxy_admin_only(
async def test_team_member_add_email_invite_open_to_team_admin(
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
"""The escape hatch the provisioning 403 names must actually work.
@ -179,34 +178,31 @@ async def test_team_member_add_email_invite_open_to_team_admin(
it stays open to team admins; if it ever closed, the 403 would be sending
callers down a dead end.
"""
await _seed_target(prisma, world, "alpha", scratch.prefix)
caller = world.keys[Actor.TEAM_ADMIN]
email = f"{scratch.prefix}-invitee@example.com"
await _seed_target(db, actors, "alpha", resources.prefix)
caller = actors.keys[Actor.TEAM_ADMIN]
email = f"{resources.prefix}-invitee@example.com"
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={
"team_id": scratch.prefix,
"team_id": resources.prefix,
"member": {"user_email": email, "role": "user"},
},
)
assert resp.status_code == 200, resp.text
invited = await prisma.db.litellm_usertable.find_first(
invited = await db.db.litellm_usertable.find_first(
where={"user_email": email}
)
assert invited is not None, "invite did not create the user"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
assert invited.user_id in _member_ids(row)
# Available-team self-join: a non-admin caller may add ITSELF to a team listed
# in litellm.default_internal_user_params["available_teams"], but the bypass
# must not escalate to role=admin or inject another user.
_SELF_JOIN = [
("self_as_user", "self", "user", 200),
("self_as_admin", "self", "admin", 403),
@ -223,36 +219,33 @@ async def test_team_member_add_available_team_self_join(
who: str,
role: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
monkeypatch,
gateway_client,
db,
resources,
actors,
tmp_path,
):
# Org-less team with no admins: the INTERNAL_USER caller is neither team
# nor org admin, so it lands on the available-team branch.
await create_scratch_team(prisma, scratch.prefix)
monkeypatch.setattr(
litellm, "default_internal_user_params", {"available_teams": [scratch.prefix]}
)
await create_scratch_team(db, resources.prefix)
caller = actors.keys[Actor.INTERNAL_USER]
member_id = caller.user_id if who == "self" else actors.keys[Actor.OWNER].user_id
caller = world.keys[Actor.INTERNAL_USER]
member_id = caller.user_id if who == "self" else world.keys[Actor.OWNER].user_id
resp = await proxy_client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={
"team_id": scratch.prefix,
"member": {"user_id": member_id, "role": role},
},
)
with owned_settings_client(
gateway_client, tmp_path, {"default_internal_user_params": {"available_teams": [resources.prefix]}}
) as client:
resp = await client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={
"team_id": resources.prefix,
"member": {"user_id": member_id, "role": role},
},
)
assert (
resp.status_code == expected_status
), f"{who}/{role}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:

40
tests/proxy_behavior/management/test_team_member_delete.py → tests/integration/management/test_team_member_delete.py
View file

@ -1,12 +1,12 @@
import pytest
from .actors import Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
# POST /team/member_delete — actor x team-shape matrix. The scratch team is
# POST /team/member_delete — actor x team-shape matrix. The resources team is
# raw-seeded with a victim member already in it; PROXY_ADMIN, the team's team
# admin, or an org admin of the team's org may remove members; else 403.
_MATRIX = [
@ -31,20 +31,20 @@ _MATRIX = [
]
async def _seed_target(prisma, world, shape: str, team_id: str, victim_id: str) -> None:
async def _seed_target(db, actors, shape: str, team_id: str, victim_id: str) -> None:
if shape == "alpha":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_a_id,
admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id],
organization_id=actors.org_a_id,
admin_user_ids=[actors.keys[Actor.TEAM_ADMIN].user_id],
member_user_ids=[victim_id],
)
elif shape == "beta":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_b_id,
organization_id=actors.org_b_id,
member_user_ids=[victim_id],
)
else: # pragma: no cover - guard
@ -64,26 +64,26 @@ async def test_team_member_delete_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
victim_id = scratch.tag("victim")
await _seed_target(prisma, world, shape, scratch.prefix, victim_id)
caller = world.keys[actor]
victim_id = resources.tag("victim")
await _seed_target(db, actors, shape, resources.prefix, victim_id)
caller = actors.keys[actor]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/member_delete",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"team_id": scratch.prefix, "user_id": victim_id},
json={"team_id": resources.prefix, "user_id": victim_id},
)
assert (
resp.status_code == expected_status
), f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:

87
tests/proxy_behavior/management/test_team_member_info_validation.py → tests/integration/management/test_team_member_info_validation.py
View file

@ -13,8 +13,9 @@ user_id is invisible from response-body alone).
import pytest
from .actors import Actor
from .conftest import create_scratch_team, create_scratch_user
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team, create_scratch_user
from tests.integration._support.database import write_rows
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -24,10 +25,10 @@ pytestmark = pytest.mark.asyncio(loop_scope="session")
# ---------------------------------------------------------------------------
async def test_member_add_both_none_rejected(proxy_client, prisma, scratch, world):
team_id = await create_scratch_team(prisma, team_id=scratch.tag("team"))
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
async def test_member_add_both_none_rejected(gateway_client, db, resources, actors):
team_id = await create_scratch_team(db, team_id=resources.tag("team"))
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {seeder}"},
json={"team_id": team_id, "member": {"role": "user"}},
@ -36,8 +37,8 @@ async def test_member_add_both_none_rejected(proxy_client, prisma, scratch, worl
# (422). Either shape proves the empty-Member payload is rejected before
# any membership row is written — pin both.
assert resp.status_code in (400, 422), resp.text
rows = await prisma.db.litellm_teammembership.find_many(where={"team_id": team_id})
assert rows == [], "empty-Member payload leaked a membership row"
rows = await db.db.litellm_teammembership.find_many(where={"team_id": team_id})
assert all(row.user_id == "default_user_id" for row in rows), "empty-Member payload leaked a membership row"
# ---------------------------------------------------------------------------
@ -46,17 +47,17 @@ async def test_member_add_both_none_rejected(proxy_client, prisma, scratch, worl
async def test_member_add_email_id_mismatch_rejected(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
email = f"{scratch.prefix}-mismatch@example.com"
email = f"{resources.prefix}-mismatch@example.com"
real_user_id = await create_scratch_user(
prisma, scratch.prefix, suffix="real", user_email=email
db, resources.prefix, suffix="real", user_email=email
)
other_user_id = await create_scratch_user(prisma, scratch.prefix, suffix="other")
other_user_id = await create_scratch_user(db, resources.prefix, suffix="other")
assert real_user_id != other_user_id # sanity
team_id = await create_scratch_team(prisma, team_id=scratch.tag("team"))
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
team_id = await create_scratch_team(db, team_id=resources.tag("team"))
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {seeder}"},
json={
@ -70,8 +71,8 @@ async def test_member_add_email_id_mismatch_rejected(
)
assert resp.status_code == 400, resp.text
assert "do not belong to the same user" in resp.text, resp.text
rows = await prisma.db.litellm_teammembership.find_many(where={"team_id": team_id})
assert rows == [], "mismatch payload leaked a membership row"
rows = await db.db.litellm_teammembership.find_many(where={"team_id": team_id})
assert not any(row.user_id in {real_user_id, other_user_id} for row in rows), "mismatch payload leaked a membership row"
# ---------------------------------------------------------------------------
@ -80,15 +81,15 @@ async def test_member_add_email_id_mismatch_rejected(
async def test_member_add_email_only_resolves_user_id(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
email = f"{scratch.prefix}-resolve@example.com"
email = f"{resources.prefix}-resolve@example.com"
user_id = await create_scratch_user(
prisma, scratch.prefix, suffix="lookup", user_email=email
db, resources.prefix, suffix="lookup", user_email=email
)
team_id = await create_scratch_team(prisma, team_id=scratch.tag("team"))
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
team_id = await create_scratch_team(db, team_id=resources.tag("team"))
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {seeder}"},
json={
@ -102,7 +103,7 @@ async def test_member_add_email_only_resolves_user_id(
# team's members_with_roles JSON. Re-read that and assert the resolved
# user_id landed — the regression shape is "email resolved to the WRONG
# user_id and was silently written to members_with_roles".
team_row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
team_row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert team_row is not None
member_user_ids = [m.get("user_id") for m in team_row.members_with_roles]
assert (
@ -116,15 +117,15 @@ async def test_member_add_email_only_resolves_user_id(
async def test_member_add_unknown_user_id_upserted(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
team_id = await create_scratch_team(prisma, team_id=scratch.tag("team"))
new_user_id = f"{scratch.prefix}-fresh"
team_id = await create_scratch_team(db, team_id=resources.tag("team"))
new_user_id = f"{resources.prefix}-fresh"
# Sanity — user does not exist yet.
pre = await prisma.db.litellm_usertable.find_unique(where={"user_id": new_user_id})
pre = await db.db.litellm_usertable.find_unique(where={"user_id": new_user_id})
assert pre is None
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {seeder}"},
json={
@ -133,7 +134,7 @@ async def test_member_add_unknown_user_id_upserted(
},
)
assert resp.status_code == 200, resp.text
post = await prisma.db.litellm_usertable.find_unique(where={"user_id": new_user_id})
post = await db.db.litellm_usertable.find_unique(where={"user_id": new_user_id})
assert post is not None, "user_id was not upserted"
# The user row was created with NULL email (the helper returned the
# member as-is, no email lookup happened because the user didn't exist).
@ -143,20 +144,24 @@ async def test_member_add_unknown_user_id_upserted(
# ---------------------------------------------------------------------------
# Duplicate-email rejection — two scratch users share an email; email-only
# Duplicate-email rejection — two resources users share an email; email-only
# add → 400 with "Multiple users found" detail.
# ---------------------------------------------------------------------------
async def test_member_add_duplicate_email_rejected(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
email = f"{scratch.prefix}-dup@example.com"
await create_scratch_user(prisma, scratch.prefix, suffix="dup1", user_email=email)
await create_scratch_user(prisma, scratch.prefix, suffix="dup2", user_email=email)
team_id = await create_scratch_team(prisma, team_id=scratch.tag("team"))
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
email = f"{resources.prefix}-dup@example.com"
first = await create_scratch_user(db, resources.prefix, suffix="dup1", user_email=email)
write_rows(
'INSERT INTO "LiteLLM_UserTable" (user_id, user_email, user_role, models, metadata) '
"VALUES (%s, %s, 'internal_user', '{}', '{}')",
(resources.tag("dup2"), email),
)
team_id = await create_scratch_team(db, team_id=resources.tag("team"))
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/member_add",
headers={"Authorization": f"Bearer {seeder}"},
json={
@ -166,5 +171,7 @@ async def test_member_add_duplicate_email_rejected(
)
assert resp.status_code == 400, resp.text
assert "Multiple users found" in resp.text, resp.text
rows = await prisma.db.litellm_teammembership.find_many(where={"team_id": team_id})
rows = await db.db.litellm_teammembership.find_many(
where={"team_id": team_id, "user_id": {"in": [first, resources.tag("dup2")]}}
)
assert rows == [], "duplicate-email payload leaked a membership row"

28
tests/proxy_behavior/management/test_team_member_me.py → tests/integration/management/test_team_member_me.py
View file

@ -4,7 +4,7 @@ import pytest
from litellm.proxy.utils import hash_token
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -38,16 +38,16 @@ _CASES = [
ids=[c[0] for c in _CASES],
)
async def test_team_member_me_matrix(
actor: Actor, team: str, expected_status: int, proxy_client, world
actor: Actor, team: str, expected_status: int, gateway_client, actors
):
team_id = {
"alpha": world.team_alpha_id,
"beta": world.team_beta_id,
"gamma": world.team_gamma_id,
"alpha": actors.team_alpha_id,
"beta": actors.team_beta_id,
"gamma": actors.team_gamma_id,
}[team]
caller = world.keys[actor]
caller = actors.keys[actor]
resp = await proxy_client.get(
resp = await gateway_client.get(
f"/team/{team_id}/members/me",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)
@ -62,22 +62,22 @@ async def test_team_member_me_matrix(
async def test_team_member_me_team_key_without_user_id_is_400(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""A key with no associated user_id (a team / service-account key) cannot
resolve 'me' — the caller has no identity to look up — so it is 400."""
cleartext = "sk-" + uuid.uuid4().hex
await prisma.db.litellm_verificationtoken.create(
await db.db.litellm_verificationtoken.create(
data={
"token": hash_token(cleartext),
"key_name": f"{scratch.prefix}-teamkey",
"key_alias": f"{scratch.prefix}-teamkey",
"team_id": world.team_alpha_id,
"key_name": f"{resources.prefix}-teamkey",
"key_alias": f"{resources.prefix}-teamkey",
"team_id": actors.team_alpha_id,
"models": [],
}
)
resp = await proxy_client.get(
f"/team/{world.team_alpha_id}/members/me",
resp = await gateway_client.get(
f"/team/{actors.team_alpha_id}/members/me",
headers={"Authorization": f"Bearer {cleartext}"},
)
assert resp.status_code == 400, resp.text

200
tests/integration/management/test_team_member_reset_budget.py Normal file
View file

@ -0,0 +1,200 @@
import uuid
import pytest
from prisma import Json
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
_SEED_SPEND = 5.0
_TEAM_DEFAULT_MAX_BUDGET = 100.0
_CUSTOM_MAX_BUDGET = 50.0
_MATRIX = [
("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200),
("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200),
("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 200),
("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403),
("alpha/owner", Actor.OWNER, "alpha", 403),
("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403),
("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403),
("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403),
("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403),
("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200),
("beta/org_admin", Actor.ORG_ADMIN, "beta", 403),
("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403),
("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200),
]
async def _seed_budget(db, budget_id: str, max_budget: float) -> str:
await db.db.litellm_budgettable.create(
data={
"budget_id": budget_id,
"max_budget": max_budget,
"created_by": "phase4-resources",
"updated_by": "phase4-resources",
}
)
return budget_id
async def _seed_team_with_default_budget(db, actors, shape: str, team_id: str, resources) -> str:
default_budget_id = await _seed_budget(db, resources.tag("team-default-budget"), _TEAM_DEFAULT_MAX_BUDGET)
if shape == "alpha":
await create_scratch_team(
db,
team_id,
organization_id=actors.org_a_id,
admin_user_ids=[actors.keys[Actor.TEAM_ADMIN].user_id],
)
elif shape == "beta":
await create_scratch_team(db, team_id, organization_id=actors.org_b_id)
else:
pytest.fail(f"unknown shape={shape}")
await db.db.litellm_teamtable.update(
where={"team_id": team_id},
data={"metadata": Json({"team_member_budget_id": default_budget_id})},
)
return default_budget_id
async def _seed_custom_member(db, team_id: str, member_id: str, resources) -> str:
custom_budget_id = await _seed_budget(db, resources.tag("custom-budget"), _CUSTOM_MAX_BUDGET)
await db.db.litellm_teammembership.create(
data={
"user_id": member_id,
"team_id": team_id,
"spend": _SEED_SPEND,
"litellm_budget_table": {"connect": {"budget_id": custom_budget_id}},
}
)
return custom_budget_id
async def _membership(db, team_id: str, member_id: str):
row = await db.db.litellm_teammembership.find_unique(
where={"user_id_team_id": {"user_id": member_id, "team_id": team_id}}
)
assert row is not None
return row
@pytest.mark.parametrize(
"actor,shape,expected_status",
[(a, sh, s) for (_id, a, sh, s) in _MATRIX],
ids=[s[0] for s in _MATRIX],
)
async def test_team_member_reset_budget_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
gateway_client,
db,
resources,
actors,
):
member_id = resources.tag("member")
default_budget_id = await _seed_team_with_default_budget(db, actors, shape, resources.prefix, resources)
custom_budget_id = await _seed_custom_member(db, resources.prefix, member_id, resources)
caller = actors.keys[actor]
resp = await gateway_client.post(
f"/team/{resources.prefix}/member/{member_id}/reset_budget",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)
assert resp.status_code == expected_status, f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await _membership(db, resources.prefix, member_id)
assert row.spend == _SEED_SPEND, "reset_budget must never touch spend"
if expected_status == 200:
assert row.budget_id == default_budget_id
body = resp.json()
assert body["budget_id"] == default_budget_id
assert body["previous_budget_id"] == custom_budget_id
assert body["budget_source"] == "team_default"
else:
assert row.budget_id == custom_budget_id, "denied but budget relinked"
async def test_team_member_reset_budget_leaves_shared_default_row_untouched(gateway_client, db, resources, actors):
member_id = resources.tag("member")
default_budget_id = await _seed_team_with_default_budget(db, actors, "alpha", resources.prefix, resources)
await _seed_custom_member(db, resources.prefix, member_id, resources)
resp = await gateway_client.post(
f"/team/{resources.prefix}/member/{member_id}/reset_budget",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 200, resp.text
default_row = await db.db.litellm_budgettable.find_unique(where={"budget_id": default_budget_id})
assert default_row is not None and default_row.max_budget == _TEAM_DEFAULT_MAX_BUDGET
info = await gateway_client.get(
f"/team/info?team_id={resources.prefix}",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert info.status_code == 200, info.text
memberships = {tm["user_id"]: tm for tm in info.json()["team_memberships"]}
assert memberships[member_id]["budget_source"] == "team_default"
assert memberships[member_id]["litellm_budget_table"]["max_budget"] == _TEAM_DEFAULT_MAX_BUDGET
async def test_team_member_reset_budget_without_team_default_detaches_member(gateway_client, db, resources, actors):
member_id = resources.tag("member")
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
await _seed_custom_member(db, resources.prefix, member_id, resources)
resp = await gateway_client.post(
f"/team/{resources.prefix}/member/{member_id}/reset_budget",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 200, resp.text
assert resp.json()["budget_id"] is None
assert resp.json()["budget_source"] == "none"
row = await _membership(db, resources.prefix, member_id)
assert row.budget_id is None
assert row.spend == _SEED_SPEND
async def test_team_member_reset_budget_with_deleted_team_default_detaches_member(gateway_client, db, resources, actors):
member_id = resources.tag("member")
await create_scratch_team(
db,
resources.prefix,
organization_id=actors.org_a_id,
metadata={"team_member_budget_id": resources.tag("deleted-budget")},
)
await _seed_custom_member(db, resources.prefix, member_id, resources)
resp = await gateway_client.post(
f"/team/{resources.prefix}/member/{member_id}/reset_budget",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 200, resp.text
assert resp.json()["budget_id"] is None
assert resp.json()["budget_source"] == "none"
row = await _membership(db, resources.prefix, member_id)
assert row.budget_id is None
async def test_team_member_reset_budget_missing_team_is_404(gateway_client, actors):
resp = await gateway_client.post(
f"/team/behavior-pin-no-such-team/member/{uuid.uuid4().hex}/reset_budget",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 404, resp.text
async def test_team_member_reset_budget_missing_membership_is_404(gateway_client, db, resources, actors):
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
resp = await gateway_client.post(
f"/team/{resources.prefix}/member/{uuid.uuid4().hex}/reset_budget",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 404, resp.text

94
tests/proxy_behavior/management/test_team_member_reset_spend.py → tests/integration/management/test_team_member_reset_spend.py
View file

@ -2,8 +2,9 @@ import uuid
import pytest
from .actors import Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team
from tests.integration._support.database import write_rows
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -32,19 +33,19 @@ _MATRIX = [
]
async def _seed_target(prisma, world, shape: str, team_id: str, member_id: str) -> None:
async def _seed_target(db, actors, shape: str, team_id: str, member_id: str) -> None:
if shape == "alpha":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_a_id,
admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id],
organization_id=actors.org_a_id,
admin_user_ids=[actors.keys[Actor.TEAM_ADMIN].user_id],
)
elif shape == "beta":
await create_scratch_team(prisma, team_id, organization_id=world.org_b_id)
await create_scratch_team(db, team_id, organization_id=actors.org_b_id)
else: # pragma: no cover - guard
pytest.fail(f"unknown shape={shape}")
await prisma.db.litellm_teammembership.create(
await db.db.litellm_teammembership.create(
data={"user_id": member_id, "team_id": team_id, "spend": _SEED_SPEND}
)
@ -58,17 +59,17 @@ async def test_team_member_reset_spend_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
member_id = scratch.tag("member")
await _seed_target(prisma, world, shape, scratch.prefix, member_id)
caller = world.keys[actor]
member_id = resources.tag("member")
await _seed_target(db, actors, shape, resources.prefix, member_id)
caller = actors.keys[actor]
resp = await proxy_client.post(
f"/team/{scratch.prefix}/member/{member_id}/reset_spend",
resp = await gateway_client.post(
f"/team/{resources.prefix}/member/{member_id}/reset_spend",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"reset_to": _RESET_TO},
)
@ -76,8 +77,8 @@ async def test_team_member_reset_spend_authz_matrix(
resp.status_code == expected_status
), f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teammembership.find_unique(
where={"user_id_team_id": {"user_id": member_id, "team_id": scratch.prefix}}
row = await db.db.litellm_teammembership.find_unique(
where={"user_id_team_id": {"user_id": member_id, "team_id": resources.prefix}}
)
assert row is not None
if expected_status == 200:
@ -86,67 +87,68 @@ async def test_team_member_reset_spend_authz_matrix(
assert row.spend == _SEED_SPEND, "denied but spend reset"
async def test_team_member_reset_spend_missing_team_is_404(proxy_client, world):
resp = await proxy_client.post(
async def test_team_member_reset_spend_missing_team_is_404(gateway_client, actors):
resp = await gateway_client.post(
f"/team/behavior-pin-no-such-team/member/{uuid.uuid4().hex}/reset_spend",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"reset_to": 0.0},
)
assert resp.status_code == 404, resp.text
async def test_team_member_reset_spend_missing_membership_is_404(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""A well-formed team but a user_id with no LiteLLM_TeamMembership row is 404."""
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
resp = await proxy_client.post(
f"/team/{scratch.prefix}/member/{uuid.uuid4().hex}/reset_spend",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
resp = await gateway_client.post(
f"/team/{resources.prefix}/member/{uuid.uuid4().hex}/reset_spend",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"reset_to": 0.0},
)
assert resp.status_code == 404, resp.text
async def test_team_member_reset_spend_above_current_spend_is_400(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
member_id = scratch.tag("member")
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
await prisma.db.litellm_teammembership.create(
data={"user_id": member_id, "team_id": scratch.prefix, "spend": 1.0}
member_id = resources.tag("member")
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
await db.db.litellm_teammembership.create(
data={"user_id": member_id, "team_id": resources.prefix, "spend": 1.0}
)
resp = await proxy_client.post(
f"/team/{scratch.prefix}/member/{member_id}/reset_spend",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
resp = await gateway_client.post(
f"/team/{resources.prefix}/member/{member_id}/reset_spend",
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"reset_to": 5.0},
)
assert resp.status_code == 400, resp.text
async def test_team_member_reset_spend_team_admin_cannot_reset_own_spend(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""A team admin targeting their own LiteLLM_TeamMembership row is 403: unchecked, an
admin could repeatedly zero their own spend right before it crosses their per-member
cap, consuming the shared team budget without the configured limit ever binding."""
team_admin = world.keys[Actor.TEAM_ADMIN]
team_admin = actors.keys[Actor.TEAM_ADMIN]
await create_scratch_team(
prisma,
scratch.prefix,
organization_id=world.org_a_id,
db,
resources.prefix,
organization_id=actors.org_a_id,
admin_user_ids=[team_admin.user_id],
)
await prisma.db.litellm_teammembership.create(
data={"user_id": team_admin.user_id, "team_id": scratch.prefix, "spend": _SEED_SPEND}
write_rows(
'UPDATE "LiteLLM_TeamMembership" SET spend = %s WHERE user_id = %s AND team_id = %s',
(str(_SEED_SPEND), team_admin.user_id, resources.prefix),
)
resp = await proxy_client.post(
f"/team/{scratch.prefix}/member/{team_admin.user_id}/reset_spend",
resp = await gateway_client.post(
f"/team/{resources.prefix}/member/{team_admin.user_id}/reset_spend",
headers={"Authorization": f"Bearer {team_admin.cleartext}"},
json={"reset_to": 0.0},
)
assert resp.status_code == 403, resp.text
row = await prisma.db.litellm_teammembership.find_unique(
where={"user_id_team_id": {"user_id": team_admin.user_id, "team_id": scratch.prefix}}
row = await db.db.litellm_teammembership.find_unique(
where={"user_id_team_id": {"user_id": team_admin.user_id, "team_id": resources.prefix}}
)
assert row is not None and row.spend == _SEED_SPEND, "denied but spend reset"

40
tests/proxy_behavior/management/test_team_member_update.py → tests/integration/management/test_team_member_update.py
View file

@ -1,12 +1,12 @@
import pytest
from .actors import Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
# POST /team/member_update — actor x team-shape matrix. The scratch team is
# POST /team/member_update — actor x team-shape matrix. The resources team is
# raw-seeded with a "user"-role member; each scenario tries to promote it to
# "admin". PROXY_ADMIN, the team's team admin, or an org admin of the team's
# org may update members; else 403. (The harness forces premium_user, so the
@ -33,20 +33,20 @@ _MATRIX = [
]
async def _seed_target(prisma, world, shape: str, team_id: str, member_id: str) -> None:
async def _seed_target(db, actors, shape: str, team_id: str, member_id: str) -> None:
if shape == "alpha":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_a_id,
admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id],
organization_id=actors.org_a_id,
admin_user_ids=[actors.keys[Actor.TEAM_ADMIN].user_id],
member_user_ids=[member_id],
)
elif shape == "beta":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_b_id,
organization_id=actors.org_b_id,
member_user_ids=[member_id],
)
else: # pragma: no cover - guard
@ -69,26 +69,26 @@ async def test_team_member_update_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
member_id = scratch.tag("member")
await _seed_target(prisma, world, shape, scratch.prefix, member_id)
caller = world.keys[actor]
member_id = resources.tag("member")
await _seed_target(db, actors, shape, resources.prefix, member_id)
caller = actors.keys[actor]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/member_update",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"team_id": scratch.prefix, "user_id": member_id, "role": "admin"},
json={"team_id": resources.prefix, "user_id": member_id, "role": "admin"},
)
assert (
resp.status_code == expected_status
), f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:

14
tests/proxy_behavior/management/test_team_metadata_schema.py → tests/integration/management/test_team_metadata_schema.py
View file

@ -1,25 +1,25 @@
"""GET /team/metadata_schema — the behavior world declares no
"""GET /team/metadata_schema — the behavior actors declares no
``general_settings.team_metadata_schema``, so the route is an info route that
returns an empty field list to every authenticated actor and 401s without a key.
"""
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@pytest.mark.parametrize("actor", list(Actor), ids=[a.value for a in Actor])
async def test_team_metadata_schema_default_is_empty(actor: Actor, proxy_client, world):
resp = await proxy_client.get(
async def test_team_metadata_schema_default_is_empty(actor: Actor, gateway_client, actors):
resp = await gateway_client.get(
"/team/metadata_schema",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
)
assert resp.status_code == 200, f"{actor.value}: {resp.status_code} {resp.text}"
assert resp.json() == {"fields": []}
async def test_team_metadata_schema_requires_auth(proxy_client, world):
resp = await proxy_client.get("/team/metadata_schema")
async def test_team_metadata_schema_requires_auth(gateway_client, actors):
resp = await gateway_client.get("/team/metadata_schema", headers={"Authorization": "Bearer invalid"})
assert resp.status_code == 401, resp.text

30
tests/proxy_behavior/management/test_team_model.py → tests/integration/management/test_team_model.py
View file

@ -1,7 +1,7 @@
import pytest
from .actors import Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -37,28 +37,28 @@ async def test_team_model_authz_matrix(
route: str,
actor: Actor,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
initial = [] if route == "add" else [_MARKER_MODEL]
await create_scratch_team(
prisma, scratch.prefix, organization_id=world.org_a_id, models=initial
db, resources.prefix, organization_id=actors.org_a_id, models=initial
)
caller = world.keys[actor]
caller = actors.keys[actor]
resp = await proxy_client.post(
resp = await gateway_client.post(
_ROUTE_URL[route],
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"team_id": scratch.prefix, "models": [_MARKER_MODEL]},
json={"team_id": resources.prefix, "models": [_MARKER_MODEL]},
)
assert (
resp.status_code == expected_status
), f"{route} {actor.value}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:
@ -68,11 +68,11 @@ async def test_team_model_authz_matrix(
@pytest.mark.parametrize("route", ["add", "delete"])
async def test_team_model_missing_team_is_404(route: str, proxy_client, world):
async def test_team_model_missing_team_is_404(route: str, gateway_client, actors):
"""A team_id absent from the DB is 404 — the existence check precedes authz."""
resp = await proxy_client.post(
resp = await gateway_client.post(
_ROUTE_URL[route],
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": "behavior-pin-no-such-team", "models": [_MARKER_MODEL]},
)
assert resp.status_code == 404, resp.text

56
tests/proxy_behavior/management/test_team_new.py → tests/integration/management/test_team_new.py
View file

@ -2,7 +2,7 @@ from typing import Any, Dict
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -51,30 +51,30 @@ async def test_team_new_authz_matrix(
actor: Actor,
org_target: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
caller = world.keys[actor]
caller = actors.keys[actor]
org_id = {
"none": None,
"org_a": world.org_a_id,
"org_b": world.org_b_id,
"org_a": actors.org_a_id,
"org_b": actors.org_b_id,
}[org_target]
body: Dict[str, Any] = {"team_id": scratch.prefix, "team_alias": scratch.prefix}
body: Dict[str, Any] = {"team_id": resources.prefix, "team_alias": resources.prefix}
if org_id is not None:
body["organization_id"] = org_id
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json=body,
)
assert resp.status_code == expected_status, f"{actor.value} org={org_target}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
if expected_status == 200:
assert row is not None
assert row.organization_id == org_id
@ -82,49 +82,49 @@ async def test_team_new_authz_matrix(
assert row is None, f"{actor.value}: denied but team row leaked"
async def test_team_new_rejects_negative_budget(proxy_client, prisma, scratch, world):
async def test_team_new_rejects_negative_budget(gateway_client, db, resources, actors):
"""Input-validation pin: max_budget < 0 is a 400, no row created."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": scratch.prefix, "max_budget": -1},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "max_budget": -1},
)
assert resp.status_code == 400, resp.text
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert row is None
async def test_team_new_rejects_duplicate_team_id(proxy_client, prisma, scratch, world):
async def test_team_new_rejects_duplicate_team_id(gateway_client, db, resources, actors):
"""Input-validation pin: a colliding team_id is a 400 on the second call."""
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
first = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
first = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {seeder}"},
json={"team_id": scratch.prefix, "team_alias": scratch.prefix},
json={"team_id": resources.prefix, "team_alias": resources.prefix},
)
assert first.status_code == 200, first.text
second = await proxy_client.post(
second = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {seeder}"},
json={"team_id": scratch.prefix, "team_alias": scratch.prefix},
json={"team_id": resources.prefix, "team_alias": resources.prefix},
)
assert second.status_code == 400, second.text
async def test_team_new_unknown_organization_is_400(proxy_client, prisma, scratch, world):
async def test_team_new_unknown_organization_is_400(gateway_client, db, resources, actors):
"""A /team/new with an organization_id that does not exist fails 400:
OrganizationNotFoundError is routed into the handler's own
'Organization not found' guard instead of escaping as a 500."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/new",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={
"team_id": scratch.prefix,
"organization_id": scratch.tag("no-such-org"),
"team_id": resources.prefix,
"organization_id": resources.tag("no-such-org"),
},
)
assert resp.status_code == 400, resp.text
assert "Organization not found" in resp.text
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert row is None

104
tests/proxy_behavior/management/test_team_permissions.py → tests/integration/management/test_team_permissions.py
View file

@ -1,10 +1,9 @@
import litellm
import pytest
from litellm.proxy._types import KeyManagementRoutes
from .actors import Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team, owned_settings_client
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -13,7 +12,7 @@ _PERM = KeyManagementRoutes.KEY_INFO.value
# GET /team/permissions_list and POST /team/permissions_update are self-managed
# routes, so every actor reaches the handler. Both grant access to PROXY_ADMIN,
# the team admin, or an org admin of the team's org. The scratch team is in
# the team admin, or an org admin of the team's org. The resources team is in
# ORG_A with TEAM_ADMIN as its team admin.
_MATRIX = [
("proxy_admin", Actor.PROXY_ADMIN, 200),
@ -28,17 +27,17 @@ _MATRIX = [
]
async def _seed_team(prisma, scratch_prefix, world) -> None:
async def _seed_team(db, scratch_prefix, actors) -> None:
await create_scratch_team(
prisma,
db,
scratch_prefix,
organization_id=world.org_a_id,
admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id],
organization_id=actors.org_a_id,
admin_user_ids=[actors.keys[Actor.TEAM_ADMIN].user_id],
member_user_ids=[
world.keys[Actor.INTERNAL_USER].user_id,
world.keys[Actor.OWNER].user_id,
world.keys[Actor.UNRELATED_SAME_ORG].user_id,
world.keys[Actor.SERVICE_ACCOUNT].user_id,
actors.keys[Actor.INTERNAL_USER].user_id,
actors.keys[Actor.OWNER].user_id,
actors.keys[Actor.UNRELATED_SAME_ORG].user_id,
actors.keys[Actor.SERVICE_ACCOUNT].user_id,
],
)
@ -49,18 +48,18 @@ async def _seed_team(prisma, scratch_prefix, world) -> None:
ids=[s[0] for s in _MATRIX],
)
async def test_team_permissions_list_authz_matrix(
actor: Actor, expected_status: int, proxy_client, prisma, scratch, world
actor: Actor, expected_status: int, gateway_client, db, resources, actors
):
await _seed_team(prisma, scratch.prefix, world)
resp = await proxy_client.get(
f"/team/permissions_list?team_id={scratch.prefix}",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
await _seed_team(db, resources.prefix, actors)
resp = await gateway_client.get(
f"/team/permissions_list?team_id={resources.prefix}",
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
)
assert (
resp.status_code == expected_status
), f"{actor.value}: {resp.status_code} {resp.text}"
if expected_status == 200:
assert resp.json()["team_id"] == scratch.prefix
assert resp.json()["team_id"] == resources.prefix
@pytest.mark.parametrize(
@ -69,20 +68,20 @@ async def test_team_permissions_list_authz_matrix(
ids=[s[0] for s in _MATRIX],
)
async def test_team_permissions_update_authz_matrix(
actor: Actor, expected_status: int, proxy_client, prisma, scratch, world
actor: Actor, expected_status: int, gateway_client, db, resources, actors
):
await _seed_team(prisma, scratch.prefix, world)
resp = await proxy_client.post(
await _seed_team(db, resources.prefix, actors)
resp = await gateway_client.post(
"/team/permissions_update",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
json={"team_id": scratch.prefix, "team_member_permissions": [_PERM]},
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
json={"team_id": resources.prefix, "team_member_permissions": [_PERM]},
)
assert (
resp.status_code == expected_status
), f"{actor.value}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:
@ -92,29 +91,24 @@ async def test_team_permissions_update_authz_matrix(
async def test_team_permissions_available_team_self_join_divergence(
proxy_client, prisma, scratch, world, monkeypatch
gateway_client, db, resources, actors, tmp_path
):
"""permissions_list honours the available-team self-join — a non-admin can
READ an available team's permissions — but permissions_update deliberately
does not: the same caller is 403 on update. default_internal_user_params is
module-level litellm.* state, so monkeypatch save/restores it."""
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
monkeypatch.setattr(
litellm, "default_internal_user_params", {"available_teams": [scratch.prefix]}
)
caller = world.keys[Actor.CROSS_ORG_USER] # non-admin, unrelated to the team
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
caller = actors.keys[Actor.CROSS_ORG_USER]
listed = await proxy_client.get(
f"/team/permissions_list?team_id={scratch.prefix}",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)
with owned_settings_client(
gateway_client, tmp_path, {"default_internal_user_params": {"available_teams": [resources.prefix]}}
) as client:
listed = await client.get(
f"/team/permissions_list?team_id={resources.prefix}",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)
updated = await client.post(
"/team/permissions_update",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"team_id": resources.prefix, "team_member_permissions": [_PERM]},
)
assert listed.status_code == 200, listed.text
updated = await proxy_client.post(
"/team/permissions_update",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"team_id": scratch.prefix, "team_member_permissions": [_PERM]},
)
assert updated.status_code == 403, updated.text
@ -138,20 +132,20 @@ _BULK_MATRIX = [
ids=[s[0] for s in _BULK_MATRIX],
)
async def test_team_permissions_bulk_update_authz_matrix(
actor: Actor, expected_status: int, proxy_client, prisma, scratch, world
actor: Actor, expected_status: int, gateway_client, db, resources, actors
):
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
resp = await proxy_client.post(
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id)
resp = await gateway_client.post(
"/team/permissions_bulk_update",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
json={"team_ids": [scratch.prefix], "permissions": [_PERM]},
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
json={"team_ids": [resources.prefix], "permissions": [_PERM]},
)
assert (
resp.status_code == expected_status
), f"{actor.value}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:
@ -160,11 +154,11 @@ async def test_team_permissions_bulk_update_authz_matrix(
assert _PERM not in (row.team_member_permissions or []), "denied but mutated"
async def test_team_permissions_bulk_update_no_selector_is_400(proxy_client, world):
async def test_team_permissions_bulk_update_no_selector_is_400(gateway_client, actors):
"""Neither team_ids nor apply_to_all_teams is a 400."""
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/permissions_bulk_update",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"permissions": [_PERM]},
)
assert resp.status_code == 400, resp.text

136
tests/proxy_behavior/management/test_team_permissions_bulk_update.py → tests/integration/management/test_team_permissions_bulk_update.py
View file

@ -8,22 +8,20 @@ Pins the three helpers
The route is admin-only — PROXY_ADMIN is the only legal caller. The
contracts under test:
* specific-team list → ONLY listed teams mutate; every other team
(including world teams) is byte-identical on re-read.
* apply_to_all_teams → every team gains the permission, idempotently
(including actors teams) is byte-identical on re-read.
* apply_to_all_teams → every team has the permission, idempotently
merged (re-running with the same permission is a no-op).
* unknown team_id → 404 (the missing_ids guard); no mutation anywhere.
* malformed payload (neither / both selector flags) → 400; no mutation.
The all-teams scenario mutates world teams as a deliberate side effect of
the path under test. The test snapshots every team's
`team_member_permissions` up front and restores them on exit so the
read-world stays immutable for downstream tests.
"""
from typing import Final
import pytest
from .actors import TEAM_ALPHA, TEAM_BETA, TEAM_GAMMA, Actor
from .conftest import create_scratch_team
from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, TEAM_GAMMA, Actor
from tests.integration._support.actors import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -34,28 +32,28 @@ pytestmark = pytest.mark.asyncio(loop_scope="session")
async def test_bulk_update_specific_team_only_mutates_listed(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
target = await create_scratch_team(
prisma,
team_id=scratch.tag("target"),
db,
team_id=resources.tag("target"),
team_member_permissions=[],
)
bystander = await create_scratch_team(
prisma,
team_id=scratch.tag("bystander"),
db,
team_id=resources.tag("bystander"),
team_member_permissions=["pre-existing-perm"],
)
# Snapshot world teams so we can assert byte-equality on re-read.
world_team_ids = [TEAM_ALPHA, TEAM_BETA, TEAM_GAMMA]
# Snapshot actors teams so we can assert byte-equality on re-read.
world_team_ids = [actors.team_alpha_id, actors.team_beta_id, actors.team_gamma_id]
before = {}
for tid in world_team_ids + [bystander]:
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": tid})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": tid})
before[tid] = list(row.team_member_permissions or [])
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
perm = "/key/info"
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/permissions_bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={"team_ids": [target], "permissions": [perm]},
@ -63,7 +61,7 @@ async def test_bulk_update_specific_team_only_mutates_listed(
assert resp.status_code == 200, resp.text
assert resp.json()["teams_updated"] == 1
target_row = await prisma.db.litellm_teamtable.find_unique(
target_row = await db.db.litellm_teamtable.find_unique(
where={"team_id": target}
)
assert perm in (
@ -71,7 +69,7 @@ async def test_bulk_update_specific_team_only_mutates_listed(
), f"target team did not gain perm; got={target_row.team_member_permissions}"
for tid, before_perms in before.items():
after_row = await prisma.db.litellm_teamtable.find_unique(
after_row = await db.db.litellm_teamtable.find_unique(
where={"team_id": tid}
)
assert list(after_row.team_member_permissions or []) == before_perms, (
@ -87,15 +85,15 @@ async def test_bulk_update_specific_team_only_mutates_listed(
async def test_bulk_update_specific_team_idempotent(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
target = await create_scratch_team(
prisma,
team_id=scratch.tag("target"),
db,
team_id=resources.tag("target"),
team_member_permissions=["/key/info"],
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/permissions_bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={
@ -113,25 +111,25 @@ async def test_bulk_update_specific_team_idempotent(
async def test_bulk_update_unknown_team_id_rejected(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
existing = await create_scratch_team(
prisma,
team_id=scratch.tag("real"),
db,
team_id=resources.tag("real"),
team_member_permissions=[],
)
before = list(
(
await prisma.db.litellm_teamtable.find_unique(where={"team_id": existing})
await db.db.litellm_teamtable.find_unique(where={"team_id": existing})
).team_member_permissions
or []
)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/permissions_bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={
"team_ids": [existing, f"{scratch.prefix}-ghost"],
"team_ids": [existing, f"{resources.prefix}-ghost"],
"permissions": ["/key/info"],
},
)
@ -141,7 +139,7 @@ async def test_bulk_update_unknown_team_id_rejected(
assert "not found" in resp.text.lower() or "ghost" in resp.text, resp.text
# Critical: the partial-success regression shape is "real team got
# mutated before the ghost-id check ran". Re-read and assert it didn't.
after_real = await prisma.db.litellm_teamtable.find_unique(
after_real = await db.db.litellm_teamtable.find_unique(
where={"team_id": existing}
)
assert (
@ -173,10 +171,10 @@ async def test_bulk_update_unknown_team_id_rejected(
ids=["neither_selector", "both_selectors"],
)
async def test_bulk_update_selector_validation(
body, expected_substring: str, proxy_client, world
body, expected_substring: str, gateway_client, actors
):
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/permissions_bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json=body,
@ -186,60 +184,40 @@ async def test_bulk_update_selector_validation(
# ---------------------------------------------------------------------------
# apply_to_all_teams — mutates every team. Snapshot + restore world teams
# so the read-world contract holds for downstream tests.
# apply_to_all_teams
# ---------------------------------------------------------------------------
async def test_bulk_update_apply_to_all_mutates_every_team(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
# Two scratch teams so we can assert "every" includes our own targets.
# Two resources teams so we can assert "every" includes our own targets.
a = await create_scratch_team(
prisma, team_id=scratch.tag("a"), team_member_permissions=[]
db, team_id=resources.tag("a"), team_member_permissions=[]
)
b = await create_scratch_team(
prisma, team_id=scratch.tag("b"), team_member_permissions=["other"]
db, team_id=resources.tag("b"), team_member_permissions=["other"]
)
perm = "/key/health" # distinct from the specific-team scenarios above
# Snapshot every team's permission list so we can restore world teams.
all_teams = await prisma.db.litellm_teamtable.find_many()
all_teams = await db.db.litellm_teamtable.find_many()
snapshot = {t.team_id: list(t.team_member_permissions or []) for t in all_teams}
expected_updates: Final = sum(
1 for permissions in snapshot.values() if perm not in permissions
)
try:
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
"/team/permissions_bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={"apply_to_all_teams": True, "permissions": [perm]},
)
assert resp.status_code == 200, resp.text
# `teams_updated` counts teams that didn't already have the perm —
# i.e. every team in the DB at call time.
assert resp.json()["teams_updated"] == len(all_teams)
seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
resp = await gateway_client.post(
"/team/permissions_bulk_update",
headers={"Authorization": f"Bearer {seeder}"},
json={"apply_to_all_teams": True, "permissions": [perm]},
)
assert resp.status_code == 200, resp.text
assert resp.json()["teams_updated"] == expected_updates
post = await prisma.db.litellm_teamtable.find_many()
for team in post:
perms = list(team.team_member_permissions or [])
assert perm in perms, f"team {team.team_id} missing the all-perm: {perms}"
# Existing perms preserved (merge, not replace).
for prior in snapshot.get(team.team_id, []):
assert (
prior in perms
), f"team {team.team_id} lost prior perm {prior!r}: {perms}"
finally:
# Restore every team — scratch teardown handles {a, b}; we must
# explicitly restore world teams (and any other non-scratch teams
# that snuck in) so downstream tests see the immutable world.
for team_id, prior_perms in snapshot.items():
current = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
if current is None:
continue
if list(current.team_member_permissions or []) != prior_perms:
await prisma.db.litellm_teamtable.update(
where={"team_id": team_id},
data={"team_member_permissions": prior_perms},
)
post = await db.db.litellm_teamtable.find_many()
for team in post:
perms = list(team.team_member_permissions or [])
assert perm in perms, f"team {team.team_id} missing the all-perm: {perms}"
for prior in snapshot.get(team.team_id, []):
assert prior in perms, f"team {team.team_id} lost prior perm {prior!r}: {perms}"

10
tests/proxy_behavior/management/test_team_spend_by_user.py → tests/integration/management/test_team_spend_by_user.py
View file

@ -1,6 +1,6 @@
import pytest
from .actors import Actor
from tests.integration._support.actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -43,13 +43,13 @@ _DATES = "start_date=2024-01-01&end_date=2024-12-31"
[(a, t, s) for (_id, a, t, s) in _CASES],
ids=[c[0] for c in _CASES],
)
async def test_team_spend_by_user_matrix(actor: Actor, team: str, expected_status: int, proxy_client, world):
team_id = {"alpha": world.team_alpha_id, "beta": world.team_beta_id}.get(team)
async def test_team_spend_by_user_matrix(actor: Actor, team: str, expected_status: int, gateway_client, actors):
team_id = {"alpha": actors.team_alpha_id, "beta": actors.team_beta_id}.get(team)
query = _DATES if team_id is None else f"{_DATES}&team_ids={team_id}"
resp = await proxy_client.get(
resp = await gateway_client.get(
f"/team/spend/by_user?{query}",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
)
assert resp.status_code == expected_status, f"{actor.value} -> {team}: {resp.status_code} {resp.text}"
if expected_status == 200:

130
tests/proxy_behavior/management/test_team_update.py → tests/integration/management/test_team_update.py
View file

@ -2,8 +2,8 @@ import pytest
from litellm.proxy._types import LitellmUserRoles
from .actors import Actor
from .conftest import create_scratch_actor, create_scratch_team
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_actor, create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -38,30 +38,30 @@ _MATRIX = [
]
async def _seed_target(prisma, world, shape: str, team_id: str) -> str:
"""Raw-seed the scratch target team; returns its organization_id."""
async def _seed_target(db, actors, shape: str, team_id: str) -> str:
"""Raw-seed the resources target team; returns its organization_id."""
if shape == "alpha":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_a_id,
admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id],
organization_id=actors.org_a_id,
admin_user_ids=[actors.keys[Actor.TEAM_ADMIN].user_id],
member_user_ids=[
world.keys[Actor.INTERNAL_USER].user_id,
world.keys[Actor.OWNER].user_id,
world.keys[Actor.UNRELATED_SAME_ORG].user_id,
world.keys[Actor.SERVICE_ACCOUNT].user_id,
actors.keys[Actor.INTERNAL_USER].user_id,
actors.keys[Actor.OWNER].user_id,
actors.keys[Actor.UNRELATED_SAME_ORG].user_id,
actors.keys[Actor.SERVICE_ACCOUNT].user_id,
],
)
return world.org_a_id
return actors.org_a_id
if shape == "beta":
await create_scratch_team(
prisma,
db,
team_id,
organization_id=world.org_b_id,
member_user_ids=[world.keys[Actor.CROSS_ORG_USER].user_id],
organization_id=actors.org_b_id,
member_user_ids=[actors.keys[Actor.CROSS_ORG_USER].user_id],
)
return world.org_b_id
return actors.org_b_id
pytest.fail(f"unknown shape={shape}") # pragma: no cover
@ -74,19 +74,19 @@ async def test_team_update_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
org_id = await _seed_target(prisma, world, shape, scratch.prefix)
caller = world.keys[actor]
org_id = await _seed_target(db, actors, shape, resources.prefix)
caller = actors.keys[actor]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={
"team_id": scratch.prefix,
"team_id": resources.prefix,
"team_alias": MARKER_ALIAS,
"organization_id": org_id,
},
@ -95,8 +95,8 @@ async def test_team_update_authz_matrix(
resp.status_code == expected_status
), f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:
@ -106,41 +106,41 @@ async def test_team_update_authz_matrix(
async def test_team_update_org_admin_resolved_from_team_without_org_context(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""With no organization_id in the body the route gate resolves the target
team's org from team_id, so an org admin of the team's own org is allowed
(200), same as PROXY_ADMIN. A team admin of that same team reaches the
handler but is refused (403) until a proxy admin enables fields for team
admins, and the response says so."""
await _seed_target(prisma, world, "alpha", scratch.prefix)
await _seed_target(db, actors, "alpha", resources.prefix)
allowed_org_admin = await proxy_client.post(
allowed_org_admin = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {world.keys[Actor.ORG_ADMIN].cleartext}"},
json={"team_id": scratch.prefix, "team_alias": MARKER_ALIAS},
headers={"Authorization": f"Bearer {actors.keys[Actor.ORG_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "team_alias": MARKER_ALIAS},
)
assert allowed_org_admin.status_code == 200, allowed_org_admin.text
allowed_proxy_admin = await proxy_client.post(
allowed_proxy_admin = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": scratch.prefix, "team_alias": MARKER_ALIAS},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "team_alias": MARKER_ALIAS},
)
assert allowed_proxy_admin.status_code == 200, allowed_proxy_admin.text
denied_team_admin = await proxy_client.post(
denied_team_admin = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {world.keys[Actor.TEAM_ADMIN].cleartext}"},
json={"team_id": scratch.prefix, "team_alias": MARKER_ALIAS},
headers={"Authorization": f"Bearer {actors.keys[Actor.TEAM_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "team_alias": MARKER_ALIAS},
)
assert denied_team_admin.status_code == 403, denied_team_admin.text
assert "cannot edit team settings" in denied_team_admin.text, denied_team_admin.text
assert "Team admin editable fields" in denied_team_admin.text, denied_team_admin.text
# Relocation gate — moving a team to a different org. The scratch team starts
# in ORG_A; each scenario relocates it to ORG_B. PROXY_ADMIN bypasses;
# Relocation gate — moving a team to a different org. The resources team starts
# in actors.org_a_id; each scenario relocates it to actors.org_b_id. PROXY_ADMIN bypasses;
# ORG_B_ADMIN reaches the handler but holds no role on the source team (403);
# ORG_ADMIN holds the source team but not the destination org (403 from the
# relocation gate); the team admin is refused by the empty field allow-list and
@ -164,31 +164,31 @@ _RELOCATION = [
async def test_team_update_org_relocation_gate(
actor: Actor,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
await _seed_target(prisma, world, "alpha", scratch.prefix)
caller = world.keys[actor]
await _seed_target(db, actors, "alpha", resources.prefix)
caller = actors.keys[actor]
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"team_id": scratch.prefix, "organization_id": world.org_b_id},
json={"team_id": resources.prefix, "organization_id": actors.org_b_id},
)
assert (
resp.status_code == expected_status
), f"{actor.value}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
row = await db.db.litellm_teamtable.find_unique(
where={"team_id": resources.prefix}
)
assert row is not None
if expected_status == 200:
assert row.organization_id == world.org_b_id
assert row.organization_id == actors.org_b_id
else:
assert row.organization_id == world.org_a_id, "denied but team relocated"
assert row.organization_id == actors.org_a_id, "denied but team relocated"
# Phase 4 F6 — explicit pin on the `team_access_denied` 403 detail string
@ -197,44 +197,44 @@ async def test_team_update_org_relocation_gate(
# status; this guard turns a silent rename of the helper's exception detail
# into a CI red.
async def test_team_update_org_b_admin_relocation_rejection_detail(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
await _seed_target(prisma, world, "alpha", scratch.prefix)
resp = await proxy_client.post(
await _seed_target(db, actors, "alpha", resources.prefix)
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {world.keys[Actor.ORG_B_ADMIN].cleartext}"},
json={"team_id": scratch.prefix, "organization_id": world.org_b_id},
headers={"Authorization": f"Bearer {actors.keys[Actor.ORG_B_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "organization_id": actors.org_b_id},
)
assert resp.status_code == 403, resp.text
assert "do not have access to this team" in resp.text, resp.text
async def test_team_update_org_relocation_allowed_for_dual_org_admin(
proxy_client, prisma, scratch, world
gateway_client, db, resources, actors
):
"""Relocation-allowed branch: a caller who is org admin of BOTH the source
and destination org may relocate a team between them. Completes the
_RELOCATION matrix, whose allowed branch PR2 left open — no seeded actor is
a dual-org admin, so one is minted with create_scratch_actor."""
actor = await create_scratch_actor(
prisma,
scratch.prefix,
db,
resources.prefix,
user_role=LitellmUserRoles.ORG_ADMIN.value,
org_admin_of=(world.org_a_id, world.org_b_id),
org_admin_of=(actors.org_a_id, actors.org_b_id),
)
team_id = await create_scratch_team(
prisma, scratch.tag("team"), organization_id=world.org_a_id
db, resources.tag("team"), organization_id=actors.org_a_id
)
resp = await proxy_client.post(
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {actor.cleartext}"},
json={"team_id": team_id, "organization_id": world.org_b_id},
json={"team_id": team_id, "organization_id": actors.org_b_id},
)
assert resp.status_code == 200, resp.text
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
row = await db.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
assert (
row.organization_id == world.org_b_id
row.organization_id == actors.org_b_id
), "dual-org admin relocation not applied"

70
tests/proxy_behavior/management/test_users_bulk_delete.py → tests/integration/management/test_users_bulk_delete.py
View file

@ -1,7 +1,7 @@
import pytest
from .actors import Actor
from .conftest import create_scratch_team, create_scratch_user
from tests.integration._support.actors import Actor
from tests.integration._support.actors import create_scratch_team, create_scratch_user
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -25,16 +25,16 @@ def _member_ids(row) -> list:
return [m["user_id"] for m in (row.members_with_roles or [])]
async def _seed_team_members(prisma, scratch, world, member_ids: list, org_id) -> None:
async def _seed_team_members(db, resources, actors, member_ids: list, org_id) -> None:
"""Leave behind what /team/member_add would: roster entry, `teams` array, and org membership."""
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id, member_user_ids=member_ids)
await prisma.db.litellm_usertable.update_many(
where={"user_id": {"in": member_ids}}, data={"teams": {"set": [scratch.prefix]}}
await create_scratch_team(db, resources.prefix, organization_id=actors.org_a_id, member_user_ids=member_ids)
await db.db.litellm_usertable.update_many(
where={"user_id": {"in": member_ids}}, data={"teams": {"set": [resources.prefix]}}
)
if org_id is None:
return
for uid in member_ids:
await prisma.db.litellm_organizationmembership.create(
await db.db.litellm_organizationmembership.create(
data={"user_id": uid, "organization_id": org_id, "user_role": "internal_user"}
)
@ -49,25 +49,25 @@ async def test_users_bulk_delete_authz_matrix(
org,
expected_status: int,
expect_deleted: bool,
proxy_client,
prisma,
scratch,
world,
gateway_client,
db,
resources,
actors,
):
victims = [await create_scratch_user(prisma, scratch.prefix, suffix=s) for s in ("v1", "v2")]
keep = await create_scratch_user(prisma, scratch.prefix, suffix="keep")
await _seed_team_members(prisma, scratch, world, victims + [keep], world.org_a_id if org == "a" else None)
victims = [await create_scratch_user(db, resources.prefix, suffix=s) for s in ("v1", "v2")]
keep = await create_scratch_user(db, resources.prefix, suffix="keep")
await _seed_team_members(db, resources, actors, victims + [keep], actors.org_a_id if org == "a" else None)
resp = await proxy_client.post(
resp = await gateway_client.post(
_URL,
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
json={"user_ids": victims},
)
assert resp.status_code == expected_status, f"{actor.value}: {resp.status_code} {resp.text}"
team = await prisma.db.litellm_teamtable.find_unique(where={"team_id": scratch.prefix})
team = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert team is not None and keep in _member_ids(team), "unrelated member removed"
remaining = {u.user_id for u in await prisma.db.litellm_usertable.find_many(where={"user_id": {"in": victims}})}
remaining = {u.user_id for u in await db.db.litellm_usertable.find_many(where={"user_id": {"in": victims}})}
if expected_status == 403:
assert resp.headers["content-type"] == "application/problem+json"
assert resp.json()["type"] == "urn:litellm:error:forbidden"
@ -79,7 +79,7 @@ async def test_users_bulk_delete_authz_matrix(
assert set(body) == {"data"}
rows = [(r["user_id"], r["success"], r["teams_removed"]) for r in body["data"]]
if expect_deleted:
assert rows == [(v, True, [scratch.prefix]) for v in victims]
assert rows == [(v, True, [resources.prefix]) for v in victims]
assert remaining == set()
assert not set(victims) & set(_member_ids(team))
return
@ -89,13 +89,13 @@ async def test_users_bulk_delete_authz_matrix(
assert set(victims) <= set(_member_ids(team))
async def test_users_bulk_delete_reports_each_row_in_order(proxy_client, prisma, scratch, world):
victim = await create_scratch_user(prisma, scratch.prefix, suffix="victim")
ghost = scratch.tag("ghost")
async def test_users_bulk_delete_reports_each_row_in_order(gateway_client, db, resources, actors):
victim = await create_scratch_user(db, resources.prefix, suffix="victim")
ghost = resources.tag("ghost")
resp = await proxy_client.post(
resp = await gateway_client.post(
_URL,
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"user_ids": [ghost, victim, victim]},
)
assert resp.status_code == 200, resp.text
@ -104,34 +104,34 @@ async def test_users_bulk_delete_reports_each_row_in_order(proxy_client, prisma,
(victim, True),
(victim, False),
]
assert await prisma.db.litellm_usertable.find_unique(where={"user_id": victim}) is None
assert await db.db.litellm_usertable.find_unique(where={"user_id": victim}) is None
async def test_users_bulk_delete_unknown_query_param_is_400_problem(proxy_client, prisma, scratch, world):
victim = await create_scratch_user(prisma, scratch.prefix, suffix="victim")
async def test_users_bulk_delete_unknown_query_param_is_400_problem(gateway_client, db, resources, actors):
victim = await create_scratch_user(db, resources.prefix, suffix="victim")
resp = await proxy_client.post(
resp = await gateway_client.post(
f"{_URL}?dry_run=1",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"user_ids": [victim]},
)
assert resp.status_code == 400, resp.text
assert resp.headers["content-type"] == "application/problem+json"
assert resp.json()["type"] == "urn:litellm:error:unknown-query-parameter"
assert "dry_run" in resp.json()["detail"]
assert await prisma.db.litellm_usertable.find_unique(where={"user_id": victim}) is not None
assert await db.db.litellm_usertable.find_unique(where={"user_id": victim}) is not None
async def test_users_bulk_delete_unknown_body_field_is_422_problem(proxy_client, prisma, scratch, world):
victim = await create_scratch_user(prisma, scratch.prefix, suffix="victim")
async def test_users_bulk_delete_unknown_body_field_is_422_problem(gateway_client, db, resources, actors):
victim = await create_scratch_user(db, resources.prefix, suffix="victim")
resp = await proxy_client.post(
resp = await gateway_client.post(
_URL,
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
headers={"Authorization": f"Bearer {actors.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"user_ids": [victim], "dry_run": True},
)
assert resp.status_code == 422, resp.text
assert resp.headers["content-type"] == "application/problem+json"
assert resp.json()["type"] == "urn:litellm:error:invalid-request-body"
assert "dry_run" in resp.json()["detail"]
assert await prisma.db.litellm_usertable.find_unique(where={"user_id": victim}) is not None
assert await db.db.litellm_usertable.find_unique(where={"user_id": victim}) is not None

0
tests/proxy_behavior/management/__init__.py → tests/integration/spend/__init__.py
View file

179
tests/proxy_behavior/spend/test_autorouter_session_rollup.py → tests/integration/spend/test_autorouter_session_rollup.py
View file

@ -6,12 +6,18 @@ tests/unit/proxy/db/test_autorouter_session_rollup.py.
"""
import asyncio
import hashlib
import time
import uuid
from collections.abc import AsyncIterator, Iterator
from dataclasses import dataclass
from contextvars import ContextVar, Token
from datetime import datetime, timedelta, timezone
from itertools import count
from types import SimpleNamespace
from typing import Final, TypedDict, cast
import pytest_asyncio
import pytest
from prisma import Prisma
from prisma.errors import RawQueryError
@ -26,12 +32,49 @@ from litellm.proxy.db.autorouter_session_rollup import (
write_autorouter_turn,
)
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup
from tests.integration._support.prisma_clients import connected_prisma
pytestmark = pytest.mark.asyncio(loop_scope="session")
T0 = datetime(2026, 8, 1, 12, 0, 0)
@dataclass(frozen=True, slots=True)
class ScopedIdentity:
namespace: str
sequence: Iterator[int]
def new(self) -> str:
return f"{self.namespace}-{next(self.sequence)}"
IDENTITY: Final = ContextVar[ScopedIdentity | None]("autorouter_test_identity", default=None)
def _new_identity() -> str:
identity: Final = IDENTITY.get()
assert identity is not None
return identity.new()
@pytest_asyncio.fixture(scope="module", loop_scope="session")
async def db() -> AsyncIterator[Prisma]:
async with connected_prisma() as client:
yield client
@pytest_asyncio.fixture(autouse=True, loop_scope="session")
async def cleanup_rows(db: Prisma) -> AsyncIterator[None]:
identity: Final = ScopedIdentity(uuid.uuid4().hex, count())
token: Final[Token[ScopedIdentity | None]] = IDENTITY.set(identity)
yield
key_pattern: Final = f"k-{identity.namespace}-%"
await db.execute_raw('DELETE FROM "LiteLLM_AutoRouterDailySpend" WHERE api_key LIKE $1', key_pattern)
await db.execute_raw('DELETE FROM "LiteLLM_AutoRouterUserSession" WHERE api_key LIKE $1', key_pattern)
await db.execute_raw('DELETE FROM "LiteLLM_AutoRouterSession" WHERE api_key LIKE $1', key_pattern)
IDENTITY.reset(token)
def _utc_epoch(moment: datetime) -> float:
return moment.replace(tzinfo=timezone.utc).timestamp()
@ -115,7 +158,7 @@ async def _row(db, key: str, session_id: str = "s1", router: str = "auto-1") ->
async def test_every_turn_lands_in_exactly_one_bucket(db):
key = f"k-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
await _turn(db, key, "A", T0, ttl=300)
await _turn(db, key, "A", T0 + timedelta(seconds=10), hit=1)
await _turn(db, key, "B", T0 + timedelta(seconds=20), ttl=3600)
@ -145,7 +188,7 @@ async def test_every_turn_lands_in_exactly_one_bucket(db):
async def test_return_misses_attribute_against_the_recorded_ttl(db):
key = f"k-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
await _turn(db, key, "A", T0, ttl=300)
await _turn(db, key, "B", T0 + timedelta(seconds=10), ttl=3600)
await _turn(db, key, "A", T0 + timedelta(seconds=400))
@ -157,7 +200,7 @@ async def test_return_misses_attribute_against_the_recorded_ttl(db):
async def test_a_return_miss_with_no_recorded_ttl_stays_unattributed(db):
key = f"k-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
await _turn(db, key, "A", T0)
await _turn(db, key, "B", T0 + timedelta(seconds=10))
await _turn(db, key, "A", T0 + timedelta(seconds=20))
@ -169,7 +212,7 @@ async def test_a_return_miss_with_no_recorded_ttl_stays_unattributed(db):
async def test_a_hit_refreshes_the_models_cache_clock(db):
key = f"k-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
await _turn(db, key, "A", T0, ttl=300)
await _turn(db, key, "B", T0 + timedelta(seconds=250), ttl=3600)
await _turn(db, key, "A", T0 + timedelta(seconds=290), hit=1)
@ -183,7 +226,7 @@ async def test_a_hit_refreshes_the_models_cache_clock(db):
async def test_out_of_order_turns_do_not_rewind_the_session(db):
key = f"k-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
await _turn(db, key, "A", T0 + timedelta(seconds=100), ttl=300)
await _turn(db, key, "B", T0 + timedelta(seconds=200))
await _turn(db, key, "A", T0)
@ -197,7 +240,7 @@ async def test_out_of_order_turns_do_not_rewind_the_session(db):
async def test_concurrent_writers_compose_without_losing_turns(db):
key = f"k-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
await _turn(db, key, "A", T0, classifier_cost=0.001)
await asyncio.gather(
*(_turn(db, key, "A", T0 + timedelta(seconds=1 + offset), hit=1, classifier_cost=0.002) for offset in range(30))
@ -231,7 +274,7 @@ async def _legacy_turn(db, key: str, at: datetime, session_id: str = "s1", route
@pytest.mark.parametrize("writers", [(False,), (True,), (False, True), (True, False)])
async def test_subtotal_coverage_survives_legacy_and_rolling_writers(db, writers: tuple[bool, ...]):
key: Final = f"k-{uuid.uuid4()}"
key: Final = f"k-{_new_identity()}"
for offset, records_cost in enumerate(writers):
at: Final = T0 + timedelta(seconds=offset)
if records_cost:
@ -260,7 +303,7 @@ async def test_subtotal_coverage_survives_legacy_and_rolling_writers(db, writers
async def test_unknown_and_legacy_turns_preserve_actual_spend_without_entering_the_estimated_cohort(db: Prisma) -> None:
key: Final = f"k-{uuid.uuid4()}"
key: Final = f"k-{_new_identity()}"
await _turn(db, key, "A", T0, spend=0.25, saved=-0.05, baseline="opus")
await _turn(
db, key, "B", T0 + timedelta(seconds=1), spend=0.7, saved=0, baseline="sonnet", estimated=False
@ -280,10 +323,10 @@ async def test_unknown_and_legacy_turns_preserve_actual_spend_without_entering_t
async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db):
key = f"k-{uuid.uuid4()}"
router = f"r-{uuid.uuid4()}"
in_window = f"s-{uuid.uuid4()}"
out_of_window = f"s-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
router = f"r-{_new_identity()}"
in_window = f"s-{_new_identity()}"
out_of_window = f"s-{_new_identity()}"
await _turn(
db,
key,
@ -315,9 +358,9 @@ async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db):
async def test_the_benchmarks_aggregate_can_filter_to_one_key(db):
router = f"r-{uuid.uuid4()}"
first_key = f"k-{uuid.uuid4()}"
second_key = f"k-{uuid.uuid4()}"
router = f"r-{_new_identity()}"
first_key = f"k-{_new_identity()}"
second_key = f"k-{_new_identity()}"
await _turn(db, first_key, "A", T0, router=router, saved=0.5, classifier_cost=0.01)
await _turn(db, second_key, "A", T0, router=router, saved=9.0, classifier_cost=0.09)
@ -330,7 +373,7 @@ async def test_the_benchmarks_aggregate_can_filter_to_one_key(db):
assert day["classifier_cost"] == pytest.approx(0.01)
assert day["classifier_cost_recorded_turns"] == 1
unknown_key_rows = await _benchmark_rows(db, (T0 - timedelta(days=1)), (T0 + timedelta(days=1)), f"k-{uuid.uuid4()}", None)
unknown_key_rows = await _benchmark_rows(db, (T0 - timedelta(days=1)), (T0 + timedelta(days=1)), f"k-{_new_identity()}", None)
assert [row for row in unknown_key_rows if row["router_name"] == router] == []
@ -360,11 +403,11 @@ async def _scoped_benchmarks(
async def test_users_keep_written_identity_across_shared_keys_and_keyless_sessions(db: Prisma) -> None:
router: Final = f"r-{uuid.uuid4()}"
alice: Final = f"u-{uuid.uuid4()}"
bob: Final = f"u-{uuid.uuid4()}"
first_key: Final = f"k-{uuid.uuid4()}"
second_key: Final = f"k-{uuid.uuid4()}"
router: Final = f"r-{_new_identity()}"
alice: Final = f"u-{_new_identity()}"
bob: Final = f"u-{_new_identity()}"
first_key: Final = f"k-{_new_identity()}"
second_key: Final = f"k-{_new_identity()}"
await _legacy_turn(db, first_key, T0, router=router)
await _turn(db, first_key, "A", T0 + timedelta(seconds=10), router=router, user_id=alice, tier="simple")
await _turn(
@ -404,7 +447,7 @@ async def test_users_keep_written_identity_across_shared_keys_and_keyless_sessio
assert (intersection[0]["sessions"], intersection[0]["turns"]) == (1, 1)
assert intersection[0]["spend"] == pytest.approx(0.01)
assert await _scoped_benchmarks(db, router, user_id=bob, key=second_key) == ()
assert await _scoped_benchmarks(db, router, user_id=f"u-{uuid.uuid4()}") == ()
assert await _scoped_benchmarks(db, router, user_id=f"u-{_new_identity()}") == ()
assert [
row
for row in await _benchmark_rows(db, (T0 - timedelta(days=1)), (T0 + timedelta(days=1)), None, "")
@ -413,8 +456,8 @@ async def test_users_keep_written_identity_across_shared_keys_and_keyless_sessio
async def test_a_failed_user_projection_rolls_back_the_keys_increment(db: Prisma) -> None:
key: Final = f"k-{uuid.uuid4()}"
user_id: Final = "".join(str(uuid.uuid4()) for _ in range(200))
key: Final = f"k-{_new_identity()}"
user_id: Final = hashlib.shake_256(_new_identity().encode()).hexdigest(4096)
await _turn(db, key, "A", T0)
before: Final = await _row(db, key)
@ -425,8 +468,8 @@ async def test_a_failed_user_projection_rolls_back_the_keys_increment(db: Prisma
assert await db.query_raw('SELECT user_id FROM "LiteLLM_AutoRouterUserSession" WHERE user_id = $1', user_id) == []
assert [day["turns"] for day in await _days(db, key)] == [1]
first_user: Final = f"u-{uuid.uuid4()}"
second_user: Final = f"u-{uuid.uuid4()}"
first_user: Final = f"u-{_new_identity()}"
second_user: Final = f"u-{_new_identity()}"
turns: Final = tuple(
AutoRouterTurnTransaction(
api_key=key,
@ -479,9 +522,9 @@ async def test_a_failed_user_projection_rolls_back_the_keys_increment(db: Prisma
async def test_user_session_cleanup_keeps_another_users_recent_keyless_session(db: Prisma) -> None:
router: Final = f"r-{uuid.uuid4()}"
expired_user: Final = f"u-{uuid.uuid4()}"
recent_user: Final = f"u-{uuid.uuid4()}"
router: Final = f"r-{_new_identity()}"
expired_user: Final = f"u-{_new_identity()}"
recent_user: Final = f"u-{_new_identity()}"
await _turn(db, "", "A", T0 - timedelta(days=1), router=router, user_id=expired_user)
await _turn(db, "", "A", T0 + timedelta(days=1), router=router, user_id=recent_user)
cleaner: Final = SpendLogCleanup(general_settings={})
@ -496,11 +539,11 @@ async def test_user_session_cleanup_keeps_another_users_recent_keyless_session(d
async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db):
key = f"k-{uuid.uuid4()}"
router = f"r-{uuid.uuid4()}"
await _turn(db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, router_type="complexity")
key = f"k-{_new_identity()}"
router = f"r-{_new_identity()}"
await _turn(db, key, "A", T0, session_id=f"s-{_new_identity()}", router=router, router_type="complexity")
await _turn(
db, key, "A", T0 + timedelta(seconds=10), session_id=f"s-{uuid.uuid4()}", router=router, router_type="quality"
db, key, "A", T0 + timedelta(seconds=10), session_id=f"s-{_new_identity()}", router=router, router_type="quality"
)
rows = await _benchmark_rows(db, (T0 - timedelta(days=1)), (T0 + timedelta(days=1)), None, None)
@ -512,7 +555,7 @@ async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db
async def test_tier_turns_count_each_tier_that_served_a_turn(db):
key = f"k-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
await _turn(db, key, "A", T0, tier="simple")
await _turn(db, key, "B", T0 + timedelta(seconds=10), tier="complex")
await _turn(db, key, "A", T0 + timedelta(seconds=20), tier="simple")
@ -521,7 +564,7 @@ async def test_tier_turns_count_each_tier_that_served_a_turn(db):
async def test_an_untiered_turn_increments_no_tier_counter(db):
key = f"k-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
await _turn(db, key, "A", T0, tier=None)
assert (await _row(db, key))["tier_turns"] == {}
@ -533,7 +576,7 @@ async def test_an_untiered_turn_increments_no_tier_counter(db):
async def test_a_mid_session_router_type_change_keeps_foreign_tier_names_out_of_the_map(db):
key = f"k-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
await _turn(db, key, "A", T0, router_type="complexity", tier="medium")
await _turn(db, key, "A", T0 + timedelta(seconds=10), router_type="quality", tier="2")
await _turn(db, key, "A", T0 + timedelta(seconds=20), router_type="complexity", tier="medium")
@ -544,7 +587,7 @@ async def test_a_mid_session_router_type_change_keeps_foreign_tier_names_out_of_
async def test_baseline_models_count_the_turns_priced_against_each_baseline(db):
key = f"k-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
await _turn(db, key, "A", T0, baseline="opus")
await _turn(db, key, "B", T0 + timedelta(seconds=10), baseline="opus")
await _turn(db, key, "A", T0 + timedelta(seconds=20), baseline="sonnet")
@ -553,7 +596,7 @@ async def test_baseline_models_count_the_turns_priced_against_each_baseline(db):
async def test_a_turn_priced_against_no_baseline_leaves_the_map_alone(db):
key = f"k-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
await _turn(db, key, "A", T0, baseline=None)
assert (await _row(db, key))["baseline_models"] == {}
@ -565,7 +608,7 @@ async def test_a_turn_priced_against_no_baseline_leaves_the_map_alone(db):
async def test_an_out_of_order_turn_still_counts_toward_its_tier(db):
key = f"k-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
await _turn(db, key, "A", T0 + timedelta(seconds=60), tier="simple")
await _turn(db, key, "A", T0, tier="simple")
@ -575,12 +618,12 @@ async def test_an_out_of_order_turn_still_counts_toward_its_tier(db):
async def test_the_benchmarks_aggregate_sums_tier_turns_across_sessions(db):
key = f"k-{uuid.uuid4()}"
router = f"r-{uuid.uuid4()}"
await _turn(db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, tier="simple")
await _turn(db, key, "A", T0 + timedelta(seconds=10), session_id=f"s-{uuid.uuid4()}", router=router, tier="simple")
await _turn(db, key, "B", T0 + timedelta(seconds=20), session_id=f"s-{uuid.uuid4()}", router=router, tier="complex")
await _turn(db, key, "C", T0 + timedelta(seconds=30), session_id=f"s-{uuid.uuid4()}", router=router, tier=None)
key = f"k-{_new_identity()}"
router = f"r-{_new_identity()}"
await _turn(db, key, "A", T0, session_id=f"s-{_new_identity()}", router=router, tier="simple")
await _turn(db, key, "A", T0 + timedelta(seconds=10), session_id=f"s-{_new_identity()}", router=router, tier="simple")
await _turn(db, key, "B", T0 + timedelta(seconds=20), session_id=f"s-{_new_identity()}", router=router, tier="complex")
await _turn(db, key, "C", T0 + timedelta(seconds=30), session_id=f"s-{_new_identity()}", router=router, tier=None)
rows = await _benchmark_rows(db, (T0 - timedelta(days=1)), (T0 + timedelta(days=1)), None, None)
grouped = next(row for row in rows if row["router_name"] == router)
@ -589,17 +632,17 @@ async def test_the_benchmarks_aggregate_sums_tier_turns_across_sessions(db):
async def test_tier_maps_stay_separate_per_router_type_on_a_reconfigured_alias(db):
key = f"k-{uuid.uuid4()}"
router = f"r-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
router = f"r-{_new_identity()}"
await _turn(
db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, router_type="complexity", tier="medium"
db, key, "A", T0, session_id=f"s-{_new_identity()}", router=router, router_type="complexity", tier="medium"
)
await _turn(
db,
key,
"A",
T0 + timedelta(seconds=10),
session_id=f"s-{uuid.uuid4()}",
session_id=f"s-{_new_identity()}",
router=router,
router_type="quality",
tier="2",
@ -611,9 +654,9 @@ async def test_tier_maps_stay_separate_per_router_type_on_a_reconfigured_alias(d
async def test_a_window_with_no_tiered_turns_aggregates_to_an_empty_map(db):
key = f"k-{uuid.uuid4()}"
router = f"r-{uuid.uuid4()}"
await _turn(db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, tier=None)
key = f"k-{_new_identity()}"
router = f"r-{_new_identity()}"
await _turn(db, key, "A", T0, session_id=f"s-{_new_identity()}", router=router, tier=None)
rows = await _benchmark_rows(db, (T0 - timedelta(days=1)), (T0 + timedelta(days=1)), None, None)
grouped = next(row for row in rows if row["router_name"] == router)
@ -621,7 +664,7 @@ async def test_a_window_with_no_tiered_turns_aggregates_to_an_empty_map(db):
async def test_a_miss_that_touched_no_cache_does_not_advance_the_ttl_clock(db):
key = f"k-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
await _turn(db, key, "A", T0, ttl=300)
await _turn(db, key, "B", T0 + timedelta(seconds=10), ttl=3600)
await _turn(db, key, "A", T0 + timedelta(seconds=400))
@ -634,7 +677,7 @@ async def test_a_miss_that_touched_no_cache_does_not_advance_the_ttl_clock(db):
async def test_an_out_of_order_hit_still_counts_toward_the_overall_hit_rate(db):
key = f"k-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
await _turn(db, key, "A", T0 + timedelta(seconds=100))
await _turn(db, key, "A", T0 + timedelta(seconds=50), hit=1)
@ -645,8 +688,8 @@ async def test_an_out_of_order_hit_still_counts_toward_the_overall_hit_rate(db):
async def test_a_cross_midnight_session_splits_its_money_by_request_day(db):
key = f"k-{uuid.uuid4()}"
router = f"auto-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
router = f"auto-{_new_identity()}"
midnight = datetime(2026, 9, 2)
await _turn(db, key, "A", midnight - timedelta(minutes=10), router=router, tokens=100, spend=1.0, saved=7.0, user_id="u1")
await _turn(db, key, "A", midnight + timedelta(minutes=10), router=router, tokens=200, spend=1.0, saved=3.0, user_id="u1")
@ -669,8 +712,8 @@ async def test_a_cross_midnight_session_splits_its_money_by_request_day(db):
async def test_a_router_type_change_within_a_day_keeps_each_types_money_apart(db):
key = f"k-{uuid.uuid4()}"
router = f"auto-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
router = f"auto-{_new_identity()}"
await _turn(db, key, "A", T0, router=router, router_type="complexity", spend=1.0, saved=4.0)
await _turn(db, key, "A", T0 + timedelta(hours=1), router=router, router_type="quality", spend=2.0, saved=0.0)
@ -679,8 +722,8 @@ async def test_a_router_type_change_within_a_day_keeps_each_types_money_apart(db
async def test_a_router_type_change_mid_session_keeps_session_shape_with_the_sessions_type(db):
key = f"k-{uuid.uuid4()}"
router = f"auto-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
router = f"auto-{_new_identity()}"
await _turn(db, key, "A", T0, router=router, router_type="complexity", spend=1.0, saved=4.0)
await _turn(db, key, "B", T0 + timedelta(hours=1), router=router, router_type="quality", spend=2.0, saved=0.0)
@ -693,8 +736,8 @@ async def test_a_router_type_change_mid_session_keeps_session_shape_with_the_ses
@pytest.mark.parametrize("statement", [UPSERT_AUTOROUTER_SESSION_SQL, UPSERT_AUTOROUTER_USER_SESSION_SQL])
async def test_a_sessionless_turn_writes_its_router_day_row_and_no_session_row(db, statement: str):
key = f"k-{uuid.uuid4()}"
router = f"auto-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
router = f"auto-{_new_identity()}"
for offset in range(2):
await write_autorouter_turn(
db,
@ -732,8 +775,8 @@ async def test_a_sessionless_turn_writes_its_router_day_row_and_no_session_row(d
@pytest.mark.parametrize("historical", [True, False])
async def test_daily_token_coverage_stays_unknown_with_old_writers(db: Prisma, historical: bool) -> None:
key: Final = f"k-{uuid.uuid4()}"
router: Final = f"auto-{uuid.uuid4()}"
key: Final = f"k-{_new_identity()}"
router: Final = f"auto-{_new_identity()}"
if historical:
await db.execute_raw(
'INSERT INTO "LiteLLM_AutoRouterDailySpend" '
@ -754,8 +797,8 @@ async def test_daily_token_coverage_stays_unknown_with_old_writers(db: Prisma, h
@pytest.mark.parametrize("missing_first", [True, False])
async def test_missing_usage_never_completes_daily_token_coverage(db: Prisma, missing_first: bool) -> None:
key: Final = f"k-{uuid.uuid4()}"
router: Final = f"auto-{uuid.uuid4()}"
key: Final = f"k-{_new_identity()}"
router: Final = f"auto-{_new_identity()}"
for offset, recorded in enumerate((not missing_first, missing_first)):
await _turn(
db, key, "A", T0 + timedelta(seconds=offset), router=router,
@ -769,8 +812,8 @@ async def test_missing_usage_never_completes_daily_token_coverage(db: Prisma, mi
async def test_router_day_money_reconciles_with_the_overall_daily_total_including_sessionless_requests(db):
from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key
key = f"k-{uuid.uuid4()}"
router = f"auto-{uuid.uuid4()}"
key = f"k-{_new_identity()}"
router = f"auto-{_new_identity()}"
requests = (("session-1", 0.25, 1.5), ("session-1", 0.5, 2.0), ("", 0.1, 0.25))
for offset, (session_id, spend, saved) in enumerate(requests):
await write_autorouter_turn(

28
tests/proxy_behavior/spend/test_baseline_accounting.py → tests/integration/spend/test_baseline_accounting.py
View file

@ -8,6 +8,7 @@ from datetime import datetime, timedelta, timezone
from typing import Final
import pytest
import pytest_asyncio
from prisma import Prisma
import litellm
@ -23,10 +24,17 @@ from litellm.proxy.db.create_views import SupportsRawQueries
from litellm.proxy.spend_tracking.baseline_accounting import BaselineObservation
from litellm.proxy.spend_tracking.savings import BaselineCostSnapshot
from litellm.types.utils import Usage
from tests.integration._support.prisma_clients import connected_prisma
pytestmark = pytest.mark.asyncio(loop_scope="session")
@pytest_asyncio.fixture(scope="module", loop_scope="session")
async def db() -> AsyncIterator[Prisma]:
async with connected_prisma() as client:
yield client
@asynccontextmanager
async def _transaction(db: Prisma, *, before_commit: bool = False, after_commit: bool = False) -> AsyncIterator[SupportsRawQueries]:
async with db.tx() as tx:
@ -44,8 +52,8 @@ def _store(db: Prisma, **faults: bool) -> BaselineAccountingStore:
return BaselineAccountingStore(transaction)
@pytest.fixture
def record() -> Callable[..., BaselineAccountingRecord]:
@pytest_asyncio.fixture(loop_scope="session")
async def record(db: Prisma) -> AsyncIterator[Callable[..., BaselineAccountingRecord]]:
run: Final = uuid.uuid4().hex
marker: Final = CountedBreakpoint("prefix", 3600, 6000, ("prefix",), "content", ("content",))
usage: Final = Usage(
@ -87,7 +95,21 @@ def record() -> Callable[..., BaselineAccountingRecord]:
),
)
return create
yield create
key: Final = run
scope: Final = "autorouter-baseline:v3:" + run * 2
await db.execute_raw('DELETE FROM "LiteLLM_DailyTagSpend" WHERE api_key=$1', key)
await db.execute_raw('DELETE FROM "LiteLLM_DailyTeamSpend" WHERE api_key=$1', key)
await db.execute_raw('DELETE FROM "LiteLLM_DailyOrganizationSpend" WHERE api_key=$1', key)
await db.execute_raw('DELETE FROM "LiteLLM_DailyEndUserSpend" WHERE api_key=$1', key)
await db.execute_raw('DELETE FROM "LiteLLM_DailyAgentSpend" WHERE api_key=$1', key)
await db.execute_raw('DELETE FROM "LiteLLM_DailyUserSpend" WHERE api_key=$1', key)
await db.execute_raw('DELETE FROM "LiteLLM_AutoRouterDailySpend" WHERE api_key=$1', key)
await db.execute_raw('DELETE FROM "LiteLLM_AutoRouterUserSession" WHERE api_key=$1', key)
await db.execute_raw('DELETE FROM "LiteLLM_AutoRouterSession" WHERE api_key=$1', key)
await db.execute_raw('DELETE FROM "LiteLLM_AutoRouterBaselineObservation" WHERE scope=$1', scope)
await db.execute_raw('DELETE FROM "LiteLLM_AutoRouterBaselineComparison" WHERE scope=$1', scope)
await db.execute_raw('DELETE FROM "LiteLLM_SpendLogs" WHERE api_key=$1', key)
async def _log(db: Prisma, record: BaselineAccountingRecord) -> None:

10
tests/proxy_behavior/spend/test_cache_activity.py → tests/integration/spend/test_cache_activity.py
View file

@ -7,10 +7,13 @@ tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py.
import json
import uuid
from collections.abc import AsyncIterator
from datetime import datetime
from typing import Final
import pytest
import pytest_asyncio
from prisma import Prisma
from litellm.proxy.analytics_endpoints.cache_activity import (
ERROR_BREAKDOWN_SQL,
@ -19,6 +22,7 @@ from litellm.proxy.analytics_endpoints.cache_activity import (
KEY_ALIAS_OPTIONS_SQL,
MODEL_OPTIONS_SQL,
)
from tests.integration._support.prisma_clients import connected_prisma
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -34,6 +38,12 @@ INFO_ONLY_MODEL: Final = f"ghost-model-{RUN}"
NO_FILTER: Final = "[]"
@pytest_asyncio.fixture(scope="module", loop_scope="session")
async def db() -> AsyncIterator[Prisma]:
async with connected_prisma() as client:
yield client
async def _spend_log(db, api_key: str, call_type: str, status: str, model: str = "", error_code: str = "") -> None:
metadata: Final = {"error_information": {"error_code": error_code, "error_class": "ProxyException"}}
await db.execute_raw(

20
tests/proxy_behavior/auth/conftest.py
View file

@ -1,20 +0,0 @@
"""Session-scoped PrismaClient for auth behavior tests that run raw SQL against a real Postgres."""
import os
from unittest.mock import MagicMock
import pytest
import pytest_asyncio
from litellm.proxy.utils import PrismaClient
@pytest_asyncio.fixture(scope="session", loop_scope="session")
async def prisma():
database_url = os.environ.get("DATABASE_URL")
if not database_url:
pytest.skip("DATABASE_URL not set") # test-quality-ok: this suite exists to run SQL on a real Postgres
client = PrismaClient(database_url=database_url, proxy_logging_obj=MagicMock())
await client.connect()
yield client
await client.disconnect()

279
tests/proxy_behavior/management/actors.py
View file

@ -1,279 +0,0 @@
"""Read-world seed for the authz matrix tests: 2 orgs, 3 teams, 9 actors."""
import enum
import uuid
from dataclasses import dataclass
from typing import Any, Dict
from prisma import Json
from litellm.proxy._types import LitellmUserRoles
from litellm.proxy.utils import PrismaClient, hash_token
class Actor(str, enum.Enum):
PROXY_ADMIN = "proxy_admin"
ORG_ADMIN = "org_admin"
TEAM_ADMIN = "team_admin"
INTERNAL_USER = "internal_user"
OWNER = "owner"
UNRELATED_SAME_ORG = "unrelated_same_org"
CROSS_ORG_USER = "cross_org_user"
SERVICE_ACCOUNT = "service_account"
ORG_B_ADMIN = "org_b_admin"
PREFIX = "behavior-pin-"
ORG_A = PREFIX + "org-a"
ORG_B = PREFIX + "org-b"
TEAM_ALPHA = PREFIX + "team-alpha"
TEAM_BETA = PREFIX + "team-beta"
TEAM_GAMMA = PREFIX + "team-gamma"
BUDGET_ID = PREFIX + "budget"
@dataclass(frozen=True)
class SeededKey:
user_id: str
cleartext: str
hashed: str
@dataclass(frozen=True)
class World:
org_a_id: str
org_b_id: str
team_alpha_id: str
team_beta_id: str
team_gamma_id: str
keys: Dict[Actor, SeededKey]
def _new_clear_key() -> str:
return "sk-" + uuid.uuid4().hex
def _actor_profile() -> Dict[Actor, Dict[str, Any]]:
return {
Actor.PROXY_ADMIN: {
"user_role": LitellmUserRoles.PROXY_ADMIN.value,
"team_id": None,
"organization_id": None,
},
Actor.ORG_ADMIN: {
"user_role": LitellmUserRoles.ORG_ADMIN.value,
"team_id": None,
"organization_id": ORG_A,
},
Actor.TEAM_ADMIN: {
"user_role": LitellmUserRoles.INTERNAL_USER.value,
"team_id": TEAM_ALPHA,
"organization_id": ORG_A,
},
Actor.INTERNAL_USER: {
"user_role": LitellmUserRoles.INTERNAL_USER.value,
"team_id": TEAM_ALPHA,
"organization_id": ORG_A,
},
Actor.OWNER: {
"user_role": LitellmUserRoles.INTERNAL_USER.value,
"team_id": TEAM_ALPHA,
"organization_id": ORG_A,
},
Actor.UNRELATED_SAME_ORG: {
"user_role": LitellmUserRoles.INTERNAL_USER.value,
"team_id": TEAM_ALPHA,
"organization_id": ORG_A,
},
Actor.CROSS_ORG_USER: {
"user_role": LitellmUserRoles.INTERNAL_USER.value,
"team_id": TEAM_BETA,
"organization_id": ORG_B,
},
Actor.SERVICE_ACCOUNT: {
"user_role": LitellmUserRoles.INTERNAL_USER.value,
"team_id": TEAM_ALPHA,
"organization_id": ORG_A,
},
Actor.ORG_B_ADMIN: {
"user_role": LitellmUserRoles.ORG_ADMIN.value,
"team_id": None,
"organization_id": ORG_B,
},
}
async def _wipe_world(prisma: PrismaClient) -> None:
await prisma.db.litellm_verificationtoken.delete_many(
where={"user_id": {"startswith": PREFIX}}
)
await prisma.db.litellm_organizationmembership.delete_many(
where={"user_id": {"startswith": PREFIX}}
)
await prisma.db.litellm_teammembership.delete_many(
where={"user_id": {"startswith": PREFIX}}
)
await prisma.db.litellm_usertable.delete_many(
where={"user_id": {"startswith": PREFIX}}
)
await prisma.db.litellm_teamtable.delete_many(
where={"team_id": {"startswith": PREFIX}}
)
await prisma.db.litellm_organizationtable.delete_many(
where={"organization_id": {"startswith": PREFIX}}
)
await prisma.db.litellm_budgettable.delete_many(where={"budget_id": BUDGET_ID})
async def seed_world(prisma: PrismaClient) -> World:
await _wipe_world(prisma)
await prisma.db.litellm_budgettable.create(
data={
"budget_id": BUDGET_ID,
"created_by": "behavior-pin-seeder",
"updated_by": "behavior-pin-seeder",
}
)
for org_id, alias in [(ORG_A, "alpha"), (ORG_B, "beta")]:
await prisma.db.litellm_organizationtable.create(
data={
"organization_id": org_id,
"organization_alias": alias,
"budget_id": BUDGET_ID,
"created_by": "behavior-pin-seeder",
"updated_by": "behavior-pin-seeder",
}
)
profiles = _actor_profile()
user_ids: Dict[Actor, str] = {actor: PREFIX + actor.value for actor in Actor}
for actor, profile in profiles.items():
teams_list = [profile["team_id"]] if profile["team_id"] else []
await prisma.db.litellm_usertable.create(
data={
"user_id": user_ids[actor],
"user_role": profile["user_role"],
"team_id": profile["team_id"],
"organization_id": profile["organization_id"],
"teams": teams_list,
}
)
# _get_user_in_team in key_management_endpoints.py walks members_with_roles
# (a JSON list of {user_id, role}), not the String[] members column —
# populate both to match what /team/new produces.
await prisma.db.litellm_teamtable.create(
data={
"team_id": TEAM_ALPHA,
"team_alias": "alpha-1",
"organization_id": ORG_A,
"admins": [user_ids[Actor.TEAM_ADMIN]],
"members": [
user_ids[Actor.TEAM_ADMIN],
user_ids[Actor.INTERNAL_USER],
user_ids[Actor.OWNER],
user_ids[Actor.UNRELATED_SAME_ORG],
user_ids[Actor.SERVICE_ACCOUNT],
],
"members_with_roles": Json(
[
{"user_id": user_ids[Actor.TEAM_ADMIN], "role": "admin"},
{"user_id": user_ids[Actor.INTERNAL_USER], "role": "user"},
{"user_id": user_ids[Actor.OWNER], "role": "user"},
{"user_id": user_ids[Actor.UNRELATED_SAME_ORG], "role": "user"},
{"user_id": user_ids[Actor.SERVICE_ACCOUNT], "role": "user"},
]
),
}
)
await prisma.db.litellm_teamtable.create(
data={
"team_id": TEAM_BETA,
"team_alias": "beta-1",
"organization_id": ORG_B,
"admins": [],
"members": [user_ids[Actor.CROSS_ORG_USER]],
"members_with_roles": Json(
[
{"user_id": user_ids[Actor.CROSS_ORG_USER], "role": "user"},
]
),
}
)
# TEAM_GAMMA: ORG_A team with no actor members — the "same-org,
# not-my-team" read target.
await prisma.db.litellm_teamtable.create(
data={
"team_id": TEAM_GAMMA,
"team_alias": "gamma-1",
"organization_id": ORG_A,
"admins": [],
"members": [],
"members_with_roles": Json([]),
}
)
for actor, org_id, role in [
(Actor.ORG_ADMIN, ORG_A, "org_admin"),
(Actor.TEAM_ADMIN, ORG_A, "internal_user"),
(Actor.INTERNAL_USER, ORG_A, "internal_user"),
(Actor.OWNER, ORG_A, "internal_user"),
(Actor.UNRELATED_SAME_ORG, ORG_A, "internal_user"),
(Actor.SERVICE_ACCOUNT, ORG_A, "internal_user"),
(Actor.CROSS_ORG_USER, ORG_B, "internal_user"),
(Actor.ORG_B_ADMIN, ORG_B, "org_admin"),
]:
await prisma.db.litellm_organizationmembership.create(
data={
"user_id": user_ids[actor],
"organization_id": org_id,
"user_role": role,
}
)
for actor, team_id in [
(Actor.TEAM_ADMIN, TEAM_ALPHA),
(Actor.INTERNAL_USER, TEAM_ALPHA),
(Actor.OWNER, TEAM_ALPHA),
(Actor.UNRELATED_SAME_ORG, TEAM_ALPHA),
(Actor.SERVICE_ACCOUNT, TEAM_ALPHA),
(Actor.CROSS_ORG_USER, TEAM_BETA),
]:
await prisma.db.litellm_teammembership.create(
data={"user_id": user_ids[actor], "team_id": team_id}
)
keys: Dict[Actor, SeededKey] = {}
for actor, profile in profiles.items():
cleartext = _new_clear_key()
hashed = hash_token(cleartext)
token_data: Dict[str, Any] = {
"token": hashed,
"key_name": PREFIX + actor.value + "-key",
"user_id": user_ids[actor],
# LiteLLM_VerificationTokenView's models field rejects NULL even
# though the column is nullable in Postgres.
"models": [],
}
if profile["team_id"]:
token_data["team_id"] = profile["team_id"]
if profile["organization_id"]:
token_data["organization_id"] = profile["organization_id"]
if actor == Actor.SERVICE_ACCOUNT:
token_data["metadata"] = Json({"service_account_id": user_ids[actor]})
await prisma.db.litellm_verificationtoken.create(data=token_data)
keys[actor] = SeededKey(
user_id=user_ids[actor], cleartext=cleartext, hashed=hashed
)
return World(
org_a_id=ORG_A,
org_b_id=ORG_B,
team_alpha_id=TEAM_ALPHA,
team_beta_id=TEAM_BETA,
team_gamma_id=TEAM_GAMMA,
keys=keys,
)

396
tests/proxy_behavior/management/conftest.py
View file

@ -1,396 +0,0 @@
"""Session-scoped async ASGI client for HTTP-boundary behavior tests."""
import os
import tempfile
import uuid
from dataclasses import dataclass
from typing import Any, AsyncIterator, Dict, Optional
import httpx
import pytest_asyncio
import yaml
from prisma import Json
from litellm.proxy.utils import hash_token
MASTER_KEY = "sk-proxy-behavior-master-key"
SCRATCH_PREFIX = "scratch-"
def _write_minimal_proxy_config() -> str:
config = {
"general_settings": {"master_key": MASTER_KEY},
"litellm_settings": {},
}
database_url = os.environ.get("DATABASE_URL")
if database_url:
config["general_settings"]["database_url"] = database_url
f = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
yaml.dump(config, f)
f.close()
return f.name
@pytest_asyncio.fixture(scope="package")
async def proxy_app():
from litellm.proxy import proxy_server
from litellm.proxy.proxy_server import (
app,
cleanup_router_config_variables,
initialize,
proxy_startup_event,
)
cleanup_router_config_variables()
config_path = _write_minimal_proxy_config()
# proxy_startup_event re-reads master_key from LITELLM_MASTER_KEY and
# unconditionally overwrites the global, even when initialize() already
# set it from the config YAML. Force (not setdefault) both vars: an
# ambient LITELLM_MASTER_KEY with a different value would make the proxy
# authenticate on that key while the tests still send MASTER_KEY.
os.environ["LITELLM_MASTER_KEY"] = MASTER_KEY
os.environ["CONFIG_FILE_PATH"] = config_path
await initialize(config=config_path)
# /key/regenerate is gated behind premium_user; flipping it lets the matrix
# pin authz behavior instead of the licensing gate.
proxy_server.premium_user = True
async with proxy_startup_event(app):
proxy_server.premium_user = True # lifespan re-runs _license_check
# The lifespan fires check_view_exists() as a background task; on a
# fresh DB the first auth call races it and resolves user_id=None.
if proxy_server.prisma_client is not None:
await proxy_server.prisma_client.check_view_exists()
yield app
@pytest_asyncio.fixture(scope="package")
async def proxy_client(proxy_app) -> AsyncIterator[httpx.AsyncClient]:
transport = httpx.ASGITransport(app=proxy_app)
async with httpx.AsyncClient(
transport=transport, base_url="http://testserver"
) as client:
yield client
@pytest_asyncio.fixture(scope="package")
async def prisma(proxy_app):
from litellm.proxy import proxy_server
assert proxy_server.prisma_client is not None
return proxy_server.prisma_client
@pytest_asyncio.fixture(scope="package")
async def world(prisma):
from .actors import seed_world
return await seed_world(prisma)
@dataclass(frozen=True)
class Scratch:
prefix: str
def tag(self, suffix: str = "") -> str:
return f"{self.prefix}-{suffix}" if suffix else self.prefix
async def create_scratch_key(
proxy_client,
seeder_cleartext: str,
scratch_prefix: str,
*,
user_id: str,
team_id: Optional[str] = None,
organization_id: Optional[str] = None,
key_alias: Optional[str] = None,
) -> str:
"""Seed a scratch-tagged key via /key/generate; returns its cleartext.
Shared by the write-scenario matrices (key update/regenerate/delete).
key_alias defaults to scratch_prefix; pass a distinct scratch-prefixed
alias when a single scenario needs more than one key (/key/generate
enforces unique aliases).
"""
body: Dict[str, Any] = {
"key_alias": key_alias or scratch_prefix,
"user_id": user_id,
}
if team_id is not None:
body["team_id"] = team_id
if organization_id is not None:
body["organization_id"] = organization_id
resp = await proxy_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder_cleartext}"},
json=body,
)
assert resp.status_code == 200, f"setup failed: {resp.text}"
return resp.json()["key"]
async def create_scratch_team(
prisma,
team_id: str,
*,
organization_id: Optional[str] = None,
admin_user_ids: Optional[list] = None,
member_user_ids: Optional[list] = None,
team_member_permissions: Optional[list] = None,
models: Optional[list] = None,
max_budget: Optional[float] = None,
tpm_limit: Optional[int] = None,
rpm_limit: Optional[int] = None,
metadata: Optional[dict] = None,
) -> str:
"""Raw-seed a scratch-tagged team row; returns its team_id.
The target team for the team write matrices (update / member_*). Raw
prisma (not POST /team/new) avoids creation side effects — no creator
auto-add, no membership rows written onto the world's users — so seeding
never mutates the immutable read-world. The authz gates read the team's
members_with_roles JSON, so a raw-seeded team exercises them exactly as
a /team/new-created team would. team_id must start with the scratch
prefix so the `scratch` fixture reclaims the row.
team_member_permissions / models seed the matching raw columns — needed
by the team-key-permission and team-model matrices.
max_budget / tpm_limit / rpm_limit / metadata seed the team's own limit
columns (Phase 4 F1+F3) — they live directly on LiteLLM_TeamTable, no
budget-table relation needed.
"""
admin_user_ids = list(admin_user_ids or [])
member_user_ids = list(member_user_ids or [])
members_with_roles = [
{"user_id": uid, "role": "admin"} for uid in admin_user_ids
] + [{"user_id": uid, "role": "user"} for uid in member_user_ids]
data: Dict[str, Any] = {
"team_id": team_id,
"team_alias": team_id,
"admins": admin_user_ids,
"members": admin_user_ids + member_user_ids,
"members_with_roles": Json(members_with_roles),
}
if organization_id is not None:
data["organization_id"] = organization_id
if team_member_permissions is not None:
data["team_member_permissions"] = team_member_permissions
if models is not None:
data["models"] = models
if max_budget is not None:
data["max_budget"] = max_budget
if tpm_limit is not None:
data["tpm_limit"] = tpm_limit
if rpm_limit is not None:
data["rpm_limit"] = rpm_limit
if metadata is not None:
data["metadata"] = Json(metadata)
await prisma.db.litellm_teamtable.create(data=data)
return team_id
async def create_scratch_user(
prisma,
scratch_prefix: str,
*,
suffix: str,
user_email: Optional[str] = None,
) -> str:
"""Raw-seed a scratch-tagged user row with no key; returns its user_id.
The passive counterpart to create_scratch_actor: a user that requests are
made *about* rather than *by*, so it needs no verification token. The
user_id matches Scratch.tag(suffix), so the teardown reclaims it.
A member named by user_id must already exist for a non-proxy-admin caller
— `_validate_member_user_id_provisioning` 403s a user_id with no user row
— so an authz matrix targeting a member has to seed one, or every
non-proxy-admin row 403s on provisioning and the permission gate under
test goes unexercised.
"""
user_id = f"{scratch_prefix}-{suffix}"
data: Dict[str, Any] = {"user_id": user_id, "user_role": "internal_user"}
if user_email is not None:
data["user_email"] = user_email
await prisma.db.litellm_usertable.create(data=data)
return user_id
async def create_scratch_org(
prisma,
scratch_prefix: str,
*,
max_budget: Optional[float] = None,
tpm_limit: Optional[int] = None,
rpm_limit: Optional[int] = None,
models: Optional[list] = None,
metadata: Optional[dict] = None,
suffix: str = "org",
) -> str:
"""Seed a scratch-tagged org + its own budget row; returns organization_id.
The org's `budget_id` points at a fresh `litellm_budgettable` row that
carries the per-org limits (`_check_org_key_limits` and the team budget
helpers read `org_table.litellm_budget_table.<limit>`, not columns on the
org row itself). Both rows share the scratch prefix so the teardown
reclaims them — budget by `budget_id` prefix (already swept), org by
`organization_id` prefix (added in this PR to the `scratch` fixture).
models / metadata seed the matching org columns; `_check_org_team_limits`
(F3) reads `org_table.models`, and the org metadata mirror of
model_rpm_limit / model_tpm_limit is what F1's model-specific org guard
consults.
"""
org_id = f"{scratch_prefix}-{suffix}"
budget_id = f"{scratch_prefix}-{suffix}-budget"
budget_data: Dict[str, Any] = {
"budget_id": budget_id,
"created_by": "phase4-scratch",
"updated_by": "phase4-scratch",
}
if max_budget is not None:
budget_data["max_budget"] = max_budget
if tpm_limit is not None:
budget_data["tpm_limit"] = tpm_limit
if rpm_limit is not None:
budget_data["rpm_limit"] = rpm_limit
await prisma.db.litellm_budgettable.create(data=budget_data)
org_data: Dict[str, Any] = {
"organization_id": org_id,
"organization_alias": org_id,
"budget_id": budget_id,
"created_by": "phase4-scratch",
"updated_by": "phase4-scratch",
}
if models is not None:
org_data["models"] = models
if metadata is not None:
org_data["metadata"] = Json(metadata)
await prisma.db.litellm_organizationtable.create(data=org_data)
return org_id
@dataclass(frozen=True)
class SeededActor:
user_id: str
cleartext: str
hashed: str
async def create_scratch_actor(
prisma,
scratch_prefix: str,
*,
user_role: str,
org_admin_of: tuple = (),
organization_id: Optional[str] = None,
suffix: str = "actor",
) -> SeededActor:
"""Mint a scratch-prefixed user + verification token (+ org memberships).
Reclaimed by the existing `scratch` teardown, which sweeps
litellm_usertable, litellm_verificationtoken, and
litellm_organizationmembership by scratch prefix — no bespoke cleanup
needed. Does NOT write litellm_teammembership against world teams: the
teardown reclaims that table only by team_id prefix, so a scratch actor
needing team membership must join a scratch team instead. The cleartext
is hashed with the real hash_token so the key authenticates end-to-end;
models=[] satisfies LiteLLM_VerificationTokenView.
"""
user_id = f"{scratch_prefix}-{suffix}"
cleartext = "sk-" + uuid.uuid4().hex
hashed = hash_token(cleartext)
await prisma.db.litellm_usertable.create(
data={
"user_id": user_id,
"user_role": user_role,
"organization_id": organization_id,
}
)
token_data: Dict[str, Any] = {
"token": hashed,
"key_name": f"{scratch_prefix}-{suffix}-key",
"key_alias": f"{scratch_prefix}-{suffix}-alias",
"user_id": user_id,
"models": [],
}
if organization_id is not None:
token_data["organization_id"] = organization_id
await prisma.db.litellm_verificationtoken.create(data=token_data)
for org_id in org_admin_of:
await prisma.db.litellm_organizationmembership.create(
data={
"user_id": user_id,
"organization_id": org_id,
"user_role": "org_admin",
}
)
return SeededActor(user_id=user_id, cleartext=cleartext, hashed=hashed)
@pytest_asyncio.fixture
async def scratch(prisma):
handle = Scratch(prefix=f"{SCRATCH_PREFIX}{uuid.uuid4().hex[:12]}")
try:
yield handle
finally:
# Children before parents to avoid FK violations.
await prisma.db.litellm_verificationtoken.delete_many(
where={
"OR": [
{"key_alias": {"startswith": handle.prefix}},
{"key_name": {"startswith": handle.prefix}},
]
}
)
await prisma.db.litellm_teammembership.delete_many(
where={"team_id": {"startswith": handle.prefix}}
)
await prisma.db.litellm_organizationmembership.delete_many(
where={"user_id": {"startswith": handle.prefix}}
)
await prisma.db.litellm_teamtable.delete_many(
where={"team_id": {"startswith": handle.prefix}}
)
# Inviting a member by user_email allocates the user_id server-side
# (a uuid), so a scratch-prefixed email is the only handle on that
# row — sweep both, matching the token sweep above.
await prisma.db.litellm_usertable.delete_many(
where={
"OR": [
{"user_id": {"startswith": handle.prefix}},
{"user_email": {"startswith": handle.prefix}},
]
}
)
# F1+F3 seed scratch orgs via create_scratch_org; the world seeder is
# the only other writer of LiteLLM_OrganizationTable and uses the
# behavior-pin- prefix, so a scratch-prefixed sweep here cannot
# collide with the read-world. Org must be reclaimed BEFORE its
# budget — org.budget_id → budget.budget_id, so deleting the parent
# first would FK-violate on any still-attached scratch org.
await prisma.db.litellm_organizationtable.delete_many(
where={"organization_id": {"startswith": handle.prefix}}
)
await prisma.db.litellm_budgettable.delete_many(
where={"budget_id": {"startswith": handle.prefix}}
)
# /team/member_add writes LiteLLM_UserTable.teams; the available-team
# self-join writes it on a world actor whose row must survive. Strip
# dangling scratch-team refs so the read-world stays immutable.
polluted = await prisma.db.litellm_usertable.find_many(
where={"teams": {"isEmpty": False}}
)
for user in polluted:
cleaned = [t for t in user.teams if not t.startswith(handle.prefix)]
if cleaned != list(user.teams):
await prisma.db.litellm_usertable.update(
where={"user_id": user.user_id},
data={"teams": {"set": cleaned}},
)

46
tests/proxy_behavior/management/test_no_management_imports.py
View file

@ -1,46 +0,0 @@
import pathlib
import re
REPO_ROOT = pathlib.Path(__file__).resolve().parents[3]
BEHAVIOR_DIR = REPO_ROOT / "tests" / "proxy_behavior"
FORBIDDEN_IMPORT = re.compile(r"^\s*from\s+litellm\.proxy\.management_endpoints\b")
FORBIDDEN_AUTH_MOCK = re.compile(
r"(?:mock\.[A-Za-z_]+|patch[a-z_]*)\([^)]*user_api_key_auth"
)
# This file is the only place the forbidden patterns appear as regex source;
# exclude it so it can describe what it forbids.
SELF = pathlib.Path(__file__).resolve()
def _iter_py_files():
for path in BEHAVIOR_DIR.rglob("*.py"):
if path.resolve() != SELF:
yield path
def _scan(pattern):
violations = []
for path in _iter_py_files():
for lineno, line in enumerate(path.read_text().splitlines(), start=1):
if pattern.search(line):
violations.append(
f"{path.relative_to(REPO_ROOT)}:{lineno}: {line.strip()}"
)
return violations
def test_no_management_endpoint_imports():
violations = _scan(FORBIDDEN_IMPORT)
assert not violations, (
"tests/proxy_behavior/ must not import from litellm.proxy.management_endpoints. "
"Violations:\n " + "\n ".join(violations)
)
def test_no_user_api_key_auth_mocking():
violations = _scan(FORBIDDEN_AUTH_MOCK)
assert not violations, (
"tests/proxy_behavior/ must not mock user_api_key_auth. "
"Violations:\n " + "\n ".join(violations)
)

91
tests/proxy_behavior/management/test_route_coverage.py
View file

@ -1,91 +0,0 @@
"""PR3.M1 — codified route coverage.
Every route declared in the two management-endpoint source files must be
exercised by at least one behavior-suite scenario. This is a permanent
regression guard: a future route added without a behavior test fails CI here,
the same way test_no_management_imports.py codifies the G3 import grep.
"""
import ast
import pathlib
import re
REPO_ROOT = pathlib.Path(__file__).resolve().parents[3]
SOURCE_FILES = [
REPO_ROOT / "litellm/proxy/management_endpoints/key_management_endpoints.py",
REPO_ROOT / "litellm/proxy/management_endpoints/team_endpoints.py",
]
TEST_DIR = pathlib.Path(__file__).resolve().parent
SELF = pathlib.Path(__file__).resolve()
# Captures the route literal from `@router.<method>("<literal>"` — `\s*` spans
# newlines so multi-line decorators are matched too.
_ROUTE_DECORATOR = re.compile(
r"@router\.(?:get|post|put|delete|patch)\(\s*[\"']([^\"']+)[\"']"
)
def _source_routes() -> set:
routes: set = set()
for path in SOURCE_FILES:
routes.update(_ROUTE_DECORATOR.findall(path.read_text()))
return routes
def _route_to_regex(route: str) -> re.Pattern:
# A plain path param ({team_id}) matches a single path segment; a Starlette
# ':path' param ({key:path}) matches across '/'. Keeping plain params
# slash-bounded stops a loose regex from falsely reporting a future
# multi-segment route as already covered.
pattern = ["^"]
pos = 0
for match in re.finditer(r"\{([^}]+)\}", route):
pattern.append(re.escape(route[pos : match.start()]))
pattern.append("[^?]+" if match.group(1).endswith(":path") else "[^/?]+")
pos = match.end()
pattern.append(re.escape(route[pos:]) + "$")
return re.compile("".join(pattern))
def _test_urls() -> set:
"""Every request-URL string literal across the behavior test suite.
f-strings are reconstructed with each interpolation collapsed to a single
placeholder char, so f"/key/{target}/regenerate" becomes /key/X/regenerate.
Query strings are dropped — coverage is a path-level property.
"""
urls: set = set()
for path in sorted(TEST_DIR.glob("test_*.py")):
if path.resolve() == SELF:
continue
tree = ast.parse(path.read_text())
for node in ast.walk(tree):
literal = None
if isinstance(node, ast.Constant) and isinstance(node.value, str):
literal = node.value
elif isinstance(node, ast.JoinedStr):
chunks = []
for value in node.values:
if isinstance(value, ast.Constant) and isinstance(value.value, str):
chunks.append(value.value)
else:
chunks.append("X") # interpolated path / query segment
literal = "".join(chunks)
if literal and literal.startswith("/"):
urls.add(literal.split("?", 1)[0])
return urls
def test_every_management_route_has_a_behavior_scenario():
routes = _source_routes()
assert routes, "no @router routes parsed — the decorator regex is stale"
urls = _test_urls()
uncovered = sorted(
route
for route in routes
if not any(_route_to_regex(route).match(url) for url in urls)
)
assert (
not uncovered
), "management routes with no behavior-suite scenario:\n " + "\n ".join(uncovered)

61
tests/proxy_behavior/management/test_scratch_teardown.py
View file

@ -1,61 +0,0 @@
import pytest
from litellm.proxy._types import LitellmUserRoles
from .actors import ORG_A, ORG_B
from .conftest import MASTER_KEY, SCRATCH_PREFIX, create_scratch_actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
# The minting tests run in file order, then _b runs after their fixture
# teardown and asserts no scratch row survived in any reclaimed table. A leak
# in either direction fails _b on the next collection.
async def test_a_scratch_key_lands_in_db(proxy_client, prisma, scratch):
resp = await proxy_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {MASTER_KEY}"},
json={"key_alias": scratch.prefix},
)
assert resp.status_code == 200, resp.text
rows = await prisma.db.litellm_verificationtoken.find_many(
where={"key_alias": scratch.prefix}
)
assert len(rows) == 1
async def test_a2_scratch_actor_lands_in_db(proxy_client, prisma, scratch):
actor = await create_scratch_actor(
prisma,
scratch.prefix,
user_role=LitellmUserRoles.ORG_ADMIN.value,
org_admin_of=(ORG_A, ORG_B),
)
user_row = await prisma.db.litellm_usertable.find_unique(
where={"user_id": actor.user_id}
)
assert user_row is not None
info = await proxy_client.get(
"/key/info", headers={"Authorization": f"Bearer {actor.cleartext}"}
)
assert info.status_code == 200, info.text
memberships = await prisma.db.litellm_organizationmembership.find_many(
where={"user_id": actor.user_id}
)
assert {m.organization_id for m in memberships} == {ORG_A, ORG_B}
async def test_b_scratch_namespace_is_clean(prisma):
tokens = await prisma.db.litellm_verificationtoken.find_many(
where={"key_alias": {"startswith": SCRATCH_PREFIX}}
)
users = await prisma.db.litellm_usertable.find_many(
where={"user_id": {"startswith": SCRATCH_PREFIX}}
)
memberships = await prisma.db.litellm_organizationmembership.find_many(
where={"user_id": {"startswith": SCRATCH_PREFIX}}
)
assert tokens == [] and users == [] and memberships == []

28
tests/proxy_behavior/management/test_smoke.py
View file

@ -1,28 +0,0 @@
import pytest
from .conftest import MASTER_KEY
pytestmark = pytest.mark.asyncio(loop_scope="session")
async def test_liveliness(proxy_client):
resp = await proxy_client.get("/health/liveliness")
assert resp.status_code == 200
async def test_key_generate_lands_in_db(proxy_client, prisma, scratch):
from litellm.proxy.utils import hash_token
resp = await proxy_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {MASTER_KEY}"},
json={"key_alias": scratch.prefix},
)
assert resp.status_code == 200, resp.text
cleartext = resp.json()["key"]
assert cleartext.startswith("sk-")
hashed = hash_token(cleartext)
row = await prisma.db.litellm_verificationtoken.find_unique(where={"token": hashed})
assert row is not None
assert row.token == hashed != cleartext

105
tests/proxy_behavior/management/test_team_bulk_member_add.py
View file

@ -1,105 +0,0 @@
import pytest
from .actors import Actor
from .conftest import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
def _member_ids(row) -> list:
return [m["user_id"] for m in (row.members_with_roles or [])]
async def test_team_bulk_member_add_proxy_admin_adds_explicit_members(
proxy_client, prisma, scratch, world
):
"""PROXY_ADMIN bulk-adds an explicit member list to a scratch team."""
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
new_member = scratch.tag("m1")
resp = await proxy_client.post(
"/team/bulk_member_add",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
json={
"team_id": scratch.prefix,
"members": [{"user_id": new_member, "role": "user"}],
},
)
assert resp.status_code == 200, resp.text
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
)
assert row is not None and new_member in _member_ids(row)
async def test_team_bulk_member_add_empty_members_is_400(
proxy_client, prisma, scratch, world
):
"""An empty member list (with all_users unset) is rejected 400."""
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
resp = await proxy_client.post(
"/team/bulk_member_add",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": scratch.prefix, "members": []},
)
assert resp.status_code == 400, resp.text
async def test_team_bulk_member_add_over_max_batch_is_400(
proxy_client, prisma, scratch, world
):
"""A member list larger than the 500-member cap is rejected 400."""
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
members = [
{"user_id": f"{scratch.prefix}-u{i}", "role": "user"} for i in range(501)
]
resp = await proxy_client.post(
"/team/bulk_member_add",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": scratch.prefix, "members": members},
)
assert resp.status_code == 400, resp.text
@pytest.mark.parametrize(
"actor",
[Actor.TEAM_ADMIN, Actor.INTERNAL_USER],
ids=["team_admin", "internal_user"],
)
async def test_team_bulk_member_add_non_admin_is_401(
actor: Actor, proxy_client, prisma, scratch, world
):
"""/team/bulk_member_add is neither an internal_user nor a self-managed
route — a non-proxy-admin with no org context is 401 at the route gate."""
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
resp = await proxy_client.post(
"/team/bulk_member_add",
headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"},
json={
"team_id": scratch.prefix,
"members": [{"user_id": scratch.tag("m"), "role": "user"}],
},
)
assert resp.status_code == 401, f"{actor.value}: {resp.status_code} {resp.text}"
async def test_team_bulk_member_add_all_users_proxy_admin(
proxy_client, prisma, scratch, world
):
"""all_users=True pulls every user in the DB into the team. The route is
reachable only by PROXY_ADMIN (the route gate 401s every other actor — even
an org admin with organization_id in the body), so the handler's own
all_users PROXY_ADMIN gate is never the deciding check at the boundary."""
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
resp = await proxy_client.post(
"/team/bulk_member_add",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"team_id": scratch.prefix, "all_users": True},
)
assert resp.status_code == 200, resp.text
row = await prisma.db.litellm_teamtable.find_unique(
where={"team_id": scratch.prefix}
)
assert row is not None
member_ids = _member_ids(row)
# every world actor is a user in the DB, so all are now team members
assert world.keys[Actor.INTERNAL_USER].user_id in member_ids

197
tests/proxy_behavior/management/test_team_member_reset_budget.py
View file

@ -1,197 +0,0 @@
import uuid
import pytest
from .actors import Actor
from .conftest import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
_SEED_SPEND = 5.0
_TEAM_DEFAULT_MAX_BUDGET = 100.0
_CUSTOM_MAX_BUDGET = 50.0
_MATRIX = [
("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200),
("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200),
("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 200),
("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403),
("alpha/owner", Actor.OWNER, "alpha", 403),
("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403),
("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403),
("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403),
("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403),
("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200),
("beta/org_admin", Actor.ORG_ADMIN, "beta", 403),
("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403),
("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200),
]
async def _seed_budget(prisma, budget_id: str, max_budget: float) -> str:
await prisma.db.litellm_budgettable.create(
data={
"budget_id": budget_id,
"max_budget": max_budget,
"created_by": "phase4-scratch",
"updated_by": "phase4-scratch",
}
)
return budget_id
async def _seed_team_with_default_budget(prisma, world, shape: str, team_id: str, scratch) -> str:
default_budget_id = await _seed_budget(prisma, scratch.tag("team-default-budget"), _TEAM_DEFAULT_MAX_BUDGET)
metadata = {"team_member_budget_id": default_budget_id}
if shape == "alpha":
await create_scratch_team(
prisma,
team_id,
organization_id=world.org_a_id,
admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id],
metadata=metadata,
)
elif shape == "beta":
await create_scratch_team(prisma, team_id, organization_id=world.org_b_id, metadata=metadata)
else: # pragma: no cover - guard
pytest.fail(f"unknown shape={shape}")
return default_budget_id
async def _seed_custom_member(prisma, team_id: str, member_id: str, scratch) -> str:
custom_budget_id = await _seed_budget(prisma, scratch.tag("custom-budget"), _CUSTOM_MAX_BUDGET)
await prisma.db.litellm_teammembership.create(
data={
"user_id": member_id,
"team_id": team_id,
"spend": _SEED_SPEND,
"litellm_budget_table": {"connect": {"budget_id": custom_budget_id}},
}
)
return custom_budget_id
async def _membership(prisma, team_id: str, member_id: str):
row = await prisma.db.litellm_teammembership.find_unique(
where={"user_id_team_id": {"user_id": member_id, "team_id": team_id}}
)
assert row is not None
return row
@pytest.mark.parametrize(
"actor,shape,expected_status",
[(a, sh, s) for (_id, a, sh, s) in _MATRIX],
ids=[s[0] for s in _MATRIX],
)
async def test_team_member_reset_budget_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
):
member_id = scratch.tag("member")
default_budget_id = await _seed_team_with_default_budget(prisma, world, shape, scratch.prefix, scratch)
custom_budget_id = await _seed_custom_member(prisma, scratch.prefix, member_id, scratch)
caller = world.keys[actor]
resp = await proxy_client.post(
f"/team/{scratch.prefix}/member/{member_id}/reset_budget",
headers={"Authorization": f"Bearer {caller.cleartext}"},
)
assert resp.status_code == expected_status, f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await _membership(prisma, scratch.prefix, member_id)
assert row.spend == _SEED_SPEND, "reset_budget must never touch spend"
if expected_status == 200:
assert row.budget_id == default_budget_id
body = resp.json()
assert body["budget_id"] == default_budget_id
assert body["previous_budget_id"] == custom_budget_id
assert body["budget_source"] == "team_default"
else:
assert row.budget_id == custom_budget_id, "denied but budget relinked"
async def test_team_member_reset_budget_leaves_shared_default_row_untouched(proxy_client, prisma, scratch, world):
member_id = scratch.tag("member")
default_budget_id = await _seed_team_with_default_budget(prisma, world, "alpha", scratch.prefix, scratch)
await _seed_custom_member(prisma, scratch.prefix, member_id, scratch)
resp = await proxy_client.post(
f"/team/{scratch.prefix}/member/{member_id}/reset_budget",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 200, resp.text
default_row = await prisma.db.litellm_budgettable.find_unique(where={"budget_id": default_budget_id})
assert default_row is not None and default_row.max_budget == _TEAM_DEFAULT_MAX_BUDGET
info = await proxy_client.get(
f"/team/info?team_id={scratch.prefix}",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert info.status_code == 200, info.text
memberships = {tm["user_id"]: tm for tm in info.json()["team_memberships"]}
assert memberships[member_id]["budget_source"] == "team_default"
assert memberships[member_id]["litellm_budget_table"]["max_budget"] == _TEAM_DEFAULT_MAX_BUDGET
async def test_team_member_reset_budget_without_team_default_detaches_member(proxy_client, prisma, scratch, world):
member_id = scratch.tag("member")
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
await _seed_custom_member(prisma, scratch.prefix, member_id, scratch)
resp = await proxy_client.post(
f"/team/{scratch.prefix}/member/{member_id}/reset_budget",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 200, resp.text
assert resp.json()["budget_id"] is None
assert resp.json()["budget_source"] == "none"
row = await _membership(prisma, scratch.prefix, member_id)
assert row.budget_id is None
assert row.spend == _SEED_SPEND
async def test_team_member_reset_budget_with_deleted_team_default_detaches_member(proxy_client, prisma, scratch, world):
member_id = scratch.tag("member")
await create_scratch_team(
prisma,
scratch.prefix,
organization_id=world.org_a_id,
metadata={"team_member_budget_id": scratch.tag("deleted-budget")},
)
await _seed_custom_member(prisma, scratch.prefix, member_id, scratch)
resp = await proxy_client.post(
f"/team/{scratch.prefix}/member/{member_id}/reset_budget",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 200, resp.text
assert resp.json()["budget_id"] is None
assert resp.json()["budget_source"] == "none"
row = await _membership(prisma, scratch.prefix, member_id)
assert row.budget_id is None
async def test_team_member_reset_budget_missing_team_is_404(proxy_client, world):
resp = await proxy_client.post(
f"/team/behavior-pin-no-such-team/member/{uuid.uuid4().hex}/reset_budget",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 404, resp.text
async def test_team_member_reset_budget_missing_membership_is_404(proxy_client, prisma, scratch, world):
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
resp = await proxy_client.post(
f"/team/{scratch.prefix}/member/{uuid.uuid4().hex}/reset_budget",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
)
assert resp.status_code == 404, resp.text

30
tests/proxy_behavior/management/test_world_seed.py
View file

@ -1,30 +0,0 @@
import pytest
from .actors import Actor
pytestmark = pytest.mark.asyncio(loop_scope="session")
@pytest.mark.parametrize("actor", list(Actor), ids=[a.value for a in Actor])
async def test_each_actor_can_self_info(actor, proxy_client, world):
seeded = world.keys[actor]
resp = await proxy_client.get(
"/key/info",
headers={"Authorization": f"Bearer {seeded.cleartext}"},
)
assert resp.status_code == 200, f"{actor.value}: {resp.text}"
body = resp.json()
assert body.get("key") == seeded.hashed
assert body["info"].get("user_id") == seeded.user_id
async def test_proxy_admin_actor_can_create_keys_for_others(proxy_client, world):
seeder = world.keys[Actor.PROXY_ADMIN]
target_user_id = world.keys[Actor.OWNER].user_id
resp = await proxy_client.post(
"/key/generate",
headers={"Authorization": f"Bearer {seeder.cleartext}"},
json={"key_alias": "smoke-proxy-admin-bypass", "user_id": target_user_id},
)
assert resp.status_code == 200, resp.text

0
tests/proxy_behavior/spend/__init__.py
View file

12
tests/proxy_behavior/spend/conftest.py
View file

@ -1,12 +0,0 @@
"""Session-scoped Prisma client for spend-rollup behavior tests against a real Postgres."""
import pytest_asyncio
from prisma import Prisma
@pytest_asyncio.fixture(scope="session", loop_scope="session")
async def db():
client = Prisma()
await client.connect()
yield client
await client.disconnect()

2
tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py
View file

@ -119,7 +119,7 @@ async def test_filters_are_passed_to_sql_as_json_arrays(mock_prisma: MagicMock):
async def test_every_query_excludes_the_same_info_routes(mock_prisma: MagicMock):
"""Regression for LIT-5884: failed info-route calls are spend-logged but are not inference traffic, so
the groups, error breakdown and both filter-option queries all receive the same exclusion list. What
the SQL does with it is covered against Postgres in tests/proxy_behavior/spend/test_cache_activity.py."""
the SQL does with it is covered against Postgres in tests/integration/spend/test_cache_activity.py."""
await get_global_activity(start_date="2026-07-01", end_date="2026-07-27", key_aliases=[], models=[])
exclusions_by_query = {call.args[0]: json.loads(call.args[-1]) for call in mock_prisma.db.query_raw.call_args_list}

2
tests/unit/proxy/db/test_autorouter_session_rollup.py
View file

@ -2,7 +2,7 @@
Unit tests for the auto-router per-session benchmarks rollup writer.
The classification SQL itself runs against a real Postgres in
tests/proxy_behavior/spend/test_autorouter_session_rollup.py; these tests cover the
tests/integration/spend/test_autorouter_session_rollup.py; these tests cover the
request-time transaction builder and the flush contract with an injected fake client.
"""

2
tests/unit/proxy/management_endpoints/test_encryption_endpoints.py
View file

@ -3,7 +3,7 @@
The endpoint bodies are exercised directly with the migration engine mocked, so
the admin guard, db-not-connected guard, and success path are all covered
without touching a live DB. The live ASGI/auth contract is covered separately in
``tests/proxy_behavior/management/test_credential_migration_endpoint.py``.
``tests/integration/management/test_credential_migration_endpoint.py``.
"""
from types import SimpleNamespace

88
tests/unit/proxy/management_endpoints/test_management_route_coverage.py Normal file
View file

@ -0,0 +1,88 @@
import ast
from collections.abc import Iterator
from itertools import chain
from pathlib import Path
import re
from typing import Final
REPO_ROOT: Final = Path(__file__).resolve().parents[4]
SOURCE_FILES: Final = (
REPO_ROOT / "litellm/proxy/management_endpoints/key_management_endpoints.py",
REPO_ROOT / "litellm/proxy/management_endpoints/team_endpoints.py",
)
TEST_DIRS: Final = (
REPO_ROOT / "tests/integration/management",
REPO_ROOT / "tests/integration/authorization",
)
SELF: Final = Path(__file__).resolve()
_ROUTE_DECORATOR = re.compile(
r"@router\.(?:get|post|put|delete|patch)\(\s*[\"']([^\"']+)[\"']"
)
def _source_routes() -> frozenset[str]:
return frozenset(chain.from_iterable(_ROUTE_DECORATOR.findall(path.read_text()) for path in SOURCE_FILES))
def _route_to_regex(route: str) -> re.Pattern[str]:
def route_part(part: str) -> str:
if part.startswith("{") and part.endswith("}"):
return "[^?]+" if part.endswith(":path}") else "[^/?]+"
return re.escape(part)
pattern: Final = "".join(route_part(part) for part in re.split(r"(\{[^}]+\})", route))
return re.compile(f"^{pattern}$")
def _url_literal(node: ast.AST) -> str | None:
if isinstance(node, ast.Constant) and isinstance(node.value, str):
return node.value
if isinstance(node, ast.JoinedStr):
return "".join(
value.value if isinstance(value, ast.Constant) and isinstance(value.value, str) else "X"
for value in node.values
)
return None
def _urls(path: Path) -> Iterator[str]:
tree: Final = ast.parse(path.read_text())
return (
literal.split("?", 1)[0]
for node in ast.walk(tree)
if (literal := _url_literal(node)) is not None and literal.startswith("/")
)
def _directory_urls(directory: Path) -> Iterator[str]:
return chain.from_iterable(
_urls(path)
for path in sorted(directory.glob("test_*.py"))
if path.resolve() != SELF
)
def _test_urls() -> frozenset[str]:
return frozenset(
chain.from_iterable(_directory_urls(directory) for directory in TEST_DIRS)
)
def _uncovered_routes(routes: frozenset[str], urls: frozenset[str]) -> Iterator[str]:
return (
route
for route in routes
if not any(_route_to_regex(route).match(url) for url in urls)
)
def test_every_management_route_has_a_behavior_scenario() -> None:
routes: Final = _source_routes()
assert routes, "no @router routes parsed"
urls: Final = _test_urls()
uncovered: Final = tuple(sorted(_uncovered_routes(routes, urls)))
assert (
not uncovered
), "management routes with no behavior-suite scenario:\n " + "\n ".join(uncovered)

1
tests/unit/test_assert_ci_coverage.py
View file

@ -124,7 +124,6 @@ def test_every_circleci_postgres_suite_is_credited_on_the_repo_as_it_stands() ->
if isinstance(job, dict) and "postgres_suite" in job
)
assert paths == (
"tests/proxy_behavior",
"tests/proxy_security_tests",
"tests/proxy_migration_tests",
"tests/integration/database/test_roi_observed.py",