From 780c504e001200af8d0143d9b06f4b71a7ae58fd Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 11 Oct 2026 08:28:42 -0700 Subject: [PATCH] test: migrate tests/proxy_behavior into integration and unit tiers (#45939) * 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 --- .circleci/config.yml | 7 - codecov.yaml | 2 + tests/integration/_support/actors.py | 443 ++++++++++++++ tests/integration/_support/prisma_clients.py | 35 ++ .../authorization}/__init__.py | 0 .../test_auth_object_prefetch.py | 10 + tests/integration/conftest.py | 59 +- .../management}/__init__.py | 0 .../test_credential_migration_endpoint.py | 29 +- .../management/test_f7_coverage_closeout.py | 199 +++---- .../management/test_f7_key_coverage_push.py | 550 +++++++++--------- .../management/test_key_aliases.py | 50 +- .../management/test_key_block_unblock.py | 60 +- .../management/test_key_budget_limits.py | 88 +-- .../management/test_key_bulk_update.py | 44 +- .../management/test_key_delete.py | 46 +- .../management/test_key_generate.py | 20 +- .../management/test_key_health.py | 14 +- .../management/test_key_info.py | 10 +- .../management/test_key_info_v2.py | 28 +- .../management/test_key_list.py | 104 ++-- .../management/test_key_regenerate.py | 107 ++-- .../management/test_key_reset_spend.py | 54 +- .../test_key_service_account_generate.py | 36 +- .../management/test_key_team_change.py | 84 +-- .../management/test_key_update.py | 78 +-- tests/integration/management/test_smoke.py | 22 + .../management/test_team_available.py | 10 +- .../management/test_team_block_unblock.py | 64 +- .../management/test_team_budget_limits.py | 142 ++--- .../management/test_team_bulk_member_add.py | 105 ++++ .../test_team_bulk_member_delete.py | 140 ++--- .../management/test_team_daily_activity.py | 32 +- .../management/test_team_delete.py | 40 +- .../management/test_team_filter_ui.py | 22 +- .../management/test_team_info.py | 22 +- .../management/test_team_key_bulk_update.py | 96 +-- .../management/test_team_list.py | 28 +- .../management/test_team_list_v2.py | 54 +- .../management/test_team_member_add.py | 135 ++--- .../management/test_team_member_delete.py | 40 +- .../test_team_member_info_validation.py | 87 +-- .../management/test_team_member_me.py | 28 +- .../test_team_member_reset_budget.py | 200 +++++++ .../test_team_member_reset_spend.py | 94 +-- .../management/test_team_member_update.py | 40 +- .../management/test_team_metadata_schema.py | 14 +- .../management/test_team_model.py | 30 +- .../management/test_team_new.py | 56 +- .../management/test_team_permissions.py | 104 ++-- .../test_team_permissions_bulk_update.py | 136 ++--- .../management/test_team_spend_by_user.py | 10 +- .../management/test_team_update.py | 130 ++--- .../management/test_users_bulk_delete.py | 70 +-- .../spend}/__init__.py | 0 .../spend/test_autorouter_session_rollup.py | 179 +++--- .../spend/test_baseline_accounting.py | 28 +- .../spend/test_cache_activity.py | 10 + tests/proxy_behavior/auth/conftest.py | 20 - tests/proxy_behavior/management/actors.py | 279 --------- tests/proxy_behavior/management/conftest.py | 396 ------------- .../management/test_no_management_imports.py | 46 -- .../management/test_route_coverage.py | 91 --- .../management/test_scratch_teardown.py | 61 -- tests/proxy_behavior/management/test_smoke.py | 28 - .../management/test_team_bulk_member_add.py | 105 ---- .../test_team_member_reset_budget.py | 197 ------- .../management/test_world_seed.py | 30 - tests/proxy_behavior/spend/__init__.py | 0 tests/proxy_behavior/spend/conftest.py | 12 - .../test_analytics_endpoints.py | 2 +- .../db/test_autorouter_session_rollup.py | 2 +- .../test_encryption_endpoints.py | 2 +- .../test_management_route_coverage.py | 88 +++ tests/unit/test_assert_ci_coverage.py | 1 - 75 files changed, 2705 insertions(+), 2980 deletions(-) create mode 100644 tests/integration/_support/actors.py create mode 100644 tests/integration/_support/prisma_clients.py rename tests/{proxy_behavior => integration/authorization}/__init__.py (100%) rename tests/{proxy_behavior/auth => integration/authorization}/test_auth_object_prefetch.py (96%) rename tests/{proxy_behavior/auth => integration/management}/__init__.py (100%) rename tests/{proxy_behavior => integration}/management/test_credential_migration_endpoint.py (69%) rename tests/{proxy_behavior => integration}/management/test_f7_coverage_closeout.py (66%) rename tests/{proxy_behavior => integration}/management/test_f7_key_coverage_push.py (63%) rename tests/{proxy_behavior => integration}/management/test_key_aliases.py (70%) rename tests/{proxy_behavior => integration}/management/test_key_block_unblock.py (76%) rename tests/{proxy_behavior => integration}/management/test_key_budget_limits.py (87%) rename tests/{proxy_behavior => integration}/management/test_key_bulk_update.py (73%) rename tests/{proxy_behavior => integration}/management/test_key_delete.py (76%) rename tests/{proxy_behavior => integration}/management/test_key_generate.py (86%) rename tests/{proxy_behavior => integration}/management/test_key_health.py (56%) rename tests/{proxy_behavior => integration}/management/test_key_info.py (92%) rename tests/{proxy_behavior => integration}/management/test_key_info_v2.py (71%) rename tests/{proxy_behavior => integration}/management/test_key_list.py (63%) rename tests/{proxy_behavior => integration}/management/test_key_regenerate.py (60%) rename tests/{proxy_behavior => integration}/management/test_key_reset_spend.py (75%) rename tests/{proxy_behavior => integration}/management/test_key_service_account_generate.py (80%) rename tests/{proxy_behavior => integration}/management/test_key_team_change.py (74%) rename tests/{proxy_behavior => integration}/management/test_key_update.py (75%) create mode 100644 tests/integration/management/test_smoke.py rename tests/{proxy_behavior => integration}/management/test_team_available.py (70%) rename tests/{proxy_behavior => integration}/management/test_team_block_unblock.py (62%) rename tests/{proxy_behavior => integration}/management/test_team_budget_limits.py (77%) create mode 100644 tests/integration/management/test_team_bulk_member_add.py rename tests/{proxy_behavior => integration}/management/test_team_bulk_member_delete.py (54%) rename tests/{proxy_behavior => integration}/management/test_team_daily_activity.py (61%) rename tests/{proxy_behavior => integration}/management/test_team_delete.py (66%) rename tests/{proxy_behavior => integration}/management/test_team_filter_ui.py (67%) rename tests/{proxy_behavior => integration}/management/test_team_info.py (85%) rename tests/{proxy_behavior => integration}/management/test_team_key_bulk_update.py (68%) rename tests/{proxy_behavior => integration}/management/test_team_list.py (84%) rename tests/{proxy_behavior => integration}/management/test_team_list_v2.py (71%) rename tests/{proxy_behavior => integration}/management/test_team_member_add.py (69%) rename tests/{proxy_behavior => integration}/management/test_team_member_delete.py (73%) rename tests/{proxy_behavior => integration}/management/test_team_member_info_validation.py (62%) rename tests/{proxy_behavior => integration}/management/test_team_member_me.py (74%) create mode 100644 tests/integration/management/test_team_member_reset_budget.py rename tests/{proxy_behavior => integration}/management/test_team_member_reset_spend.py (58%) rename tests/{proxy_behavior => integration}/management/test_team_member_update.py (74%) rename tests/{proxy_behavior => integration}/management/test_team_metadata_schema.py (57%) rename tests/{proxy_behavior => integration}/management/test_team_model.py (75%) rename tests/{proxy_behavior => integration}/management/test_team_new.py (70%) rename tests/{proxy_behavior => integration}/management/test_team_permissions.py (52%) rename tests/{proxy_behavior => integration}/management/test_team_permissions_bulk_update.py (57%) rename tests/{proxy_behavior => integration}/management/test_team_spend_by_user.py (84%) rename tests/{proxy_behavior => integration}/management/test_team_update.py (65%) rename tests/{proxy_behavior => integration}/management/test_users_bulk_delete.py (61%) rename tests/{proxy_behavior/management => integration/spend}/__init__.py (100%) rename tests/{proxy_behavior => integration}/spend/test_autorouter_session_rollup.py (87%) rename tests/{proxy_behavior => integration}/spend/test_baseline_accounting.py (92%) rename tests/{proxy_behavior => integration}/spend/test_cache_activity.py (93%) delete mode 100644 tests/proxy_behavior/auth/conftest.py delete mode 100644 tests/proxy_behavior/management/actors.py delete mode 100644 tests/proxy_behavior/management/conftest.py delete mode 100644 tests/proxy_behavior/management/test_no_management_imports.py delete mode 100644 tests/proxy_behavior/management/test_route_coverage.py delete mode 100644 tests/proxy_behavior/management/test_scratch_teardown.py delete mode 100644 tests/proxy_behavior/management/test_smoke.py delete mode 100644 tests/proxy_behavior/management/test_team_bulk_member_add.py delete mode 100644 tests/proxy_behavior/management/test_team_member_reset_budget.py delete mode 100644 tests/proxy_behavior/management/test_world_seed.py delete mode 100644 tests/proxy_behavior/spend/__init__.py delete mode 100644 tests/proxy_behavior/spend/conftest.py create mode 100644 tests/unit/proxy/management_endpoints/test_management_route_coverage.py diff --git a/.circleci/config.yml b/.circleci/config.yml index f8c7a9e60be..40ed959ff29 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -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 diff --git a/codecov.yaml b/codecov.yaml index ae72f6d5286..5b30e78a8fb 100644 --- a/codecov.yaml +++ b/codecov.yaml @@ -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 diff --git a/tests/integration/_support/actors.py b/tests/integration/_support/actors.py new file mode 100644 index 00000000000..6a1e53bda31 --- /dev/null +++ b/tests/integration/_support/actors.py @@ -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}}, + ) diff --git a/tests/integration/_support/prisma_clients.py b/tests/integration/_support/prisma_clients.py new file mode 100644 index 00000000000..5f765e2a6db --- /dev/null +++ b/tests/integration/_support/prisma_clients.py @@ -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() diff --git a/tests/proxy_behavior/__init__.py b/tests/integration/authorization/__init__.py similarity index 100% rename from tests/proxy_behavior/__init__.py rename to tests/integration/authorization/__init__.py diff --git a/tests/proxy_behavior/auth/test_auth_object_prefetch.py b/tests/integration/authorization/test_auth_object_prefetch.py similarity index 96% rename from tests/proxy_behavior/auth/test_auth_object_prefetch.py rename to tests/integration/authorization/test_auth_object_prefetch.py index 2d3b6da2a45..b37399b5338 100644 --- a/tests/proxy_behavior/auth/test_auth_object_prefetch.py +++ b/tests/integration/authorization/test_auth_object_prefetch.py @@ -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) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index dab7c1bbd1d..047347ae962 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -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": diff --git a/tests/proxy_behavior/auth/__init__.py b/tests/integration/management/__init__.py similarity index 100% rename from tests/proxy_behavior/auth/__init__.py rename to tests/integration/management/__init__.py diff --git a/tests/proxy_behavior/management/test_credential_migration_endpoint.py b/tests/integration/management/test_credential_migration_endpoint.py similarity index 69% rename from tests/proxy_behavior/management/test_credential_migration_endpoint.py rename to tests/integration/management/test_credential_migration_endpoint.py index b0428195674..795c76370f7 100644 --- a/tests/proxy_behavior/management/test_credential_migration_endpoint.py +++ b/tests/integration/management/test_credential_migration_endpoint.py @@ -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}"}, ) diff --git a/tests/proxy_behavior/management/test_f7_coverage_closeout.py b/tests/integration/management/test_f7_coverage_closeout.py similarity index 66% rename from tests/proxy_behavior/management/test_f7_coverage_closeout.py rename to tests/integration/management/test_f7_coverage_closeout.py index c78adbbb79d..527b230e543 100644 --- a/tests/proxy_behavior/management/test_f7_coverage_closeout.py +++ b/tests/integration/management/test_f7_coverage_closeout.py @@ -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--budget-`), 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 diff --git a/tests/proxy_behavior/management/test_f7_key_coverage_push.py b/tests/integration/management/test_f7_key_coverage_push.py similarity index 63% rename from tests/proxy_behavior/management/test_f7_key_coverage_push.py rename to tests/integration/management/test_f7_key_coverage_push.py index 1bf3d1d394e..84b186abd0d 100644 --- a/tests/proxy_behavior/management/test_f7_key_coverage_push.py +++ b/tests/integration/management/test_f7_key_coverage_push.py @@ -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}"}, ) diff --git a/tests/proxy_behavior/management/test_key_aliases.py b/tests/integration/management/test_key_aliases.py similarity index 70% rename from tests/proxy_behavior/management/test_key_aliases.py rename to tests/integration/management/test_key_aliases.py index 38ce5cdfaf3..cab06c33f3c 100644 --- a/tests/proxy_behavior/management/test_key_aliases.py +++ b/tests/integration/management/test_key_aliases.py @@ -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"]} diff --git a/tests/proxy_behavior/management/test_key_block_unblock.py b/tests/integration/management/test_key_block_unblock.py similarity index 76% rename from tests/proxy_behavior/management/test_key_block_unblock.py rename to tests/integration/management/test_key_block_unblock.py index 37aa0c0219a..b8d197e4218 100644 --- a/tests/proxy_behavior/management/test_key_block_unblock.py +++ b/tests/integration/management/test_key_block_unblock.py @@ -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}, diff --git a/tests/proxy_behavior/management/test_key_budget_limits.py b/tests/integration/management/test_key_budget_limits.py similarity index 87% rename from tests/proxy_behavior/management/test_key_budget_limits.py rename to tests/integration/management/test_key_budget_limits.py index f5066d60a34..52919d86da7 100644 --- a/tests/proxy_behavior/management/test_key_budget_limits.py +++ b/tests/integration/management/test_key_budget_limits.py @@ -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) diff --git a/tests/proxy_behavior/management/test_key_bulk_update.py b/tests/integration/management/test_key_bulk_update.py similarity index 73% rename from tests/proxy_behavior/management/test_key_bulk_update.py rename to tests/integration/management/test_key_bulk_update.py index 1a57998cece..ba566c19f8a 100644 --- a/tests/proxy_behavior/management/test_key_bulk_update.py +++ b/tests/integration/management/test_key_bulk_update.py @@ -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 diff --git a/tests/proxy_behavior/management/test_key_delete.py b/tests/integration/management/test_key_delete.py similarity index 76% rename from tests/proxy_behavior/management/test_key_delete.py rename to tests/integration/management/test_key_delete.py index 0b483edc056..2920692bd13 100644 --- a/tests/proxy_behavior/management/test_key_delete.py +++ b/tests/integration/management/test_key_delete.py @@ -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 diff --git a/tests/proxy_behavior/management/test_key_generate.py b/tests/integration/management/test_key_generate.py similarity index 86% rename from tests/proxy_behavior/management/test_key_generate.py rename to tests/integration/management/test_key_generate.py index 851de33d3ff..9d6cd60fe7b 100644 --- a/tests/proxy_behavior/management/test_key_generate.py +++ b/tests/integration/management/test_key_generate.py @@ -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"] diff --git a/tests/proxy_behavior/management/test_key_health.py b/tests/integration/management/test_key_health.py similarity index 56% rename from tests/proxy_behavior/management/test_key_health.py rename to tests/integration/management/test_key_health.py index 62147e7fa13..af0d985d64b 100644 --- a/tests/proxy_behavior/management/test_key_health.py +++ b/tests/integration/management/test_key_health.py @@ -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 diff --git a/tests/proxy_behavior/management/test_key_info.py b/tests/integration/management/test_key_info.py similarity index 92% rename from tests/proxy_behavior/management/test_key_info.py rename to tests/integration/management/test_key_info.py index ddcef9fd27b..78226dc1510 100644 --- a/tests/proxy_behavior/management/test_key_info.py +++ b/tests/integration/management/test_key_info.py @@ -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}"}, ) diff --git a/tests/proxy_behavior/management/test_key_info_v2.py b/tests/integration/management/test_key_info_v2.py similarity index 71% rename from tests/proxy_behavior/management/test_key_info_v2.py rename to tests/integration/management/test_key_info_v2.py index b0fb27a19fa..777a1ab0a25 100644 --- a/tests/proxy_behavior/management/test_key_info_v2.py +++ b/tests/integration/management/test_key_info_v2.py @@ -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 diff --git a/tests/proxy_behavior/management/test_key_list.py b/tests/integration/management/test_key_list.py similarity index 63% rename from tests/proxy_behavior/management/test_key_list.py rename to tests/integration/management/test_key_list.py index 0d3f329950c..b70a3782eda 100644 --- a/tests/proxy_behavior/management/test_key_list.py +++ b/tests/integration/management/test_key_list.py @@ -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 diff --git a/tests/proxy_behavior/management/test_key_regenerate.py b/tests/integration/management/test_key_regenerate.py similarity index 60% rename from tests/proxy_behavior/management/test_key_regenerate.py rename to tests/integration/management/test_key_regenerate.py index 724b8b6d65b..3cd81520047 100644 --- a/tests/proxy_behavior/management/test_key_regenerate.py +++ b/tests/integration/management/test_key_regenerate.py @@ -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 diff --git a/tests/proxy_behavior/management/test_key_reset_spend.py b/tests/integration/management/test_key_reset_spend.py similarity index 75% rename from tests/proxy_behavior/management/test_key_reset_spend.py rename to tests/integration/management/test_key_reset_spend.py index fb1c266f655..c43ff1720bd 100644 --- a/tests/proxy_behavior/management/test_key_reset_spend.py +++ b/tests/integration/management/test_key_reset_spend.py @@ -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}, diff --git a/tests/proxy_behavior/management/test_key_service_account_generate.py b/tests/integration/management/test_key_service_account_generate.py similarity index 80% rename from tests/proxy_behavior/management/test_key_service_account_generate.py rename to tests/integration/management/test_key_service_account_generate.py index 3b5bbe39754..af33ebc1585 100644 --- a/tests/proxy_behavior/management/test_key_service_account_generate.py +++ b/tests/integration/management/test_key_service_account_generate.py @@ -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 == [] diff --git a/tests/proxy_behavior/management/test_key_team_change.py b/tests/integration/management/test_key_team_change.py similarity index 74% rename from tests/proxy_behavior/management/test_key_team_change.py rename to tests/integration/management/test_key_team_change.py index 3bd4a0af0e2..dc6a17214c9 100644 --- a/tests/proxy_behavior/management/test_key_team_change.py +++ b/tests/integration/management/test_key_team_change.py @@ -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" diff --git a/tests/proxy_behavior/management/test_key_update.py b/tests/integration/management/test_key_update.py similarity index 75% rename from tests/proxy_behavior/management/test_key_update.py rename to tests/integration/management/test_key_update.py index 7b7f6f5558b..830647e749b 100644 --- a/tests/proxy_behavior/management/test_key_update.py +++ b/tests/integration/management/test_key_update.py @@ -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 diff --git a/tests/integration/management/test_smoke.py b/tests/integration/management/test_smoke.py new file mode 100644 index 00000000000..7848b702bb6 --- /dev/null +++ b/tests/integration/management/test_smoke.py @@ -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 diff --git a/tests/proxy_behavior/management/test_team_available.py b/tests/integration/management/test_team_available.py similarity index 70% rename from tests/proxy_behavior/management/test_team_available.py rename to tests/integration/management/test_team_available.py index 874c8dd4df7..f37d70b71a5 100644 --- a/tests/proxy_behavior/management/test_team_available.py +++ b/tests/integration/management/test_team_available.py @@ -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() == [] diff --git a/tests/proxy_behavior/management/test_team_block_unblock.py b/tests/integration/management/test_team_block_unblock.py similarity index 62% rename from tests/proxy_behavior/management/test_team_block_unblock.py rename to tests/integration/management/test_team_block_unblock.py index f90ee6ee6d8..106c6ccbdbd 100644 --- a/tests/proxy_behavior/management/test_team_block_unblock.py +++ b/tests/integration/management/test_team_block_unblock.py @@ -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 diff --git a/tests/proxy_behavior/management/test_team_budget_limits.py b/tests/integration/management/test_team_budget_limits.py similarity index 77% rename from tests/proxy_behavior/management/test_team_budget_limits.py rename to tests/integration/management/test_team_budget_limits.py index a172f625e91..e886f729d5a 100644 --- a/tests/proxy_behavior/management/test_team_budget_limits.py +++ b/tests/integration/management/test_team_budget_limits.py @@ -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" diff --git a/tests/integration/management/test_team_bulk_member_add.py b/tests/integration/management/test_team_bulk_member_add.py new file mode 100644 index 00000000000..90bb9670cf2 --- /dev/null +++ b/tests/integration/management/test_team_bulk_member_add.py @@ -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 diff --git a/tests/proxy_behavior/management/test_team_bulk_member_delete.py b/tests/integration/management/test_team_bulk_member_delete.py similarity index 54% rename from tests/proxy_behavior/management/test_team_bulk_member_delete.py rename to tests/integration/management/test_team_bulk_member_delete.py index 0818fb25dfe..49108a7cf4e 100644 --- a/tests/proxy_behavior/management/test_team_bulk_member_delete.py +++ b/tests/integration/management/test_team_bulk_member_delete.py @@ -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" diff --git a/tests/proxy_behavior/management/test_team_daily_activity.py b/tests/integration/management/test_team_daily_activity.py similarity index 61% rename from tests/proxy_behavior/management/test_team_daily_activity.py rename to tests/integration/management/test_team_daily_activity.py index d84cc4c94af..1b71db993ef 100644 --- a/tests/proxy_behavior/management/test_team_daily_activity.py +++ b/tests/integration/management/test_team_daily_activity.py @@ -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 diff --git a/tests/proxy_behavior/management/test_team_delete.py b/tests/integration/management/test_team_delete.py similarity index 66% rename from tests/proxy_behavior/management/test_team_delete.py rename to tests/integration/management/test_team_delete.py index 2fa1ba09883..a4b3a4121b5 100644 --- a/tests/proxy_behavior/management/test_team_delete.py +++ b/tests/integration/management/test_team_delete.py @@ -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" diff --git a/tests/proxy_behavior/management/test_team_filter_ui.py b/tests/integration/management/test_team_filter_ui.py similarity index 67% rename from tests/proxy_behavior/management/test_team_filter_ui.py rename to tests/integration/management/test_team_filter_ui.py index 69cbabf72a1..8b2e35c0ebc 100644 --- a/tests/proxy_behavior/management/test_team_filter_ui.py +++ b/tests/integration/management/test_team_filter_ui.py @@ -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 diff --git a/tests/proxy_behavior/management/test_team_info.py b/tests/integration/management/test_team_info.py similarity index 85% rename from tests/proxy_behavior/management/test_team_info.py rename to tests/integration/management/test_team_info.py index eecb22cf731..bad6ccb3ceb 100644 --- a/tests/proxy_behavior/management/test_team_info.py +++ b/tests/integration/management/test_team_info.py @@ -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. diff --git a/tests/proxy_behavior/management/test_team_key_bulk_update.py b/tests/integration/management/test_team_key_bulk_update.py similarity index 68% rename from tests/proxy_behavior/management/test_team_key_bulk_update.py rename to tests/integration/management/test_team_key_bulk_update.py index 5acf0c8185c..8ad3d3db6ee 100644 --- a/tests/proxy_behavior/management/test_team_key_bulk_update.py +++ b/tests/integration/management/test_team_key_bulk_update.py @@ -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 diff --git a/tests/proxy_behavior/management/test_team_list.py b/tests/integration/management/test_team_list.py similarity index 84% rename from tests/proxy_behavior/management/test_team_list.py rename to tests/integration/management/test_team_list.py index 2bd106dd2d0..df8660b6502 100644 --- a/tests/proxy_behavior/management/test_team_list.py +++ b/tests/integration/management/test_team_list.py @@ -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)}" ) diff --git a/tests/proxy_behavior/management/test_team_list_v2.py b/tests/integration/management/test_team_list_v2.py similarity index 71% rename from tests/proxy_behavior/management/test_team_list_v2.py rename to tests/integration/management/test_team_list_v2.py index 81178ad73c0..d608d054daa 100644 --- a/tests/proxy_behavior/management/test_team_list_v2.py +++ b/tests/integration/management/test_team_list_v2.py @@ -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 diff --git a/tests/proxy_behavior/management/test_team_member_add.py b/tests/integration/management/test_team_member_add.py similarity index 69% rename from tests/proxy_behavior/management/test_team_member_add.py rename to tests/integration/management/test_team_member_add.py index 3309aa11cb0..3f7c686fca3 100644 --- a/tests/proxy_behavior/management/test_team_member_add.py +++ b/tests/integration/management/test_team_member_add.py @@ -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: diff --git a/tests/proxy_behavior/management/test_team_member_delete.py b/tests/integration/management/test_team_member_delete.py similarity index 73% rename from tests/proxy_behavior/management/test_team_member_delete.py rename to tests/integration/management/test_team_member_delete.py index 43879d9fd16..47e49e682cd 100644 --- a/tests/proxy_behavior/management/test_team_member_delete.py +++ b/tests/integration/management/test_team_member_delete.py @@ -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: diff --git a/tests/proxy_behavior/management/test_team_member_info_validation.py b/tests/integration/management/test_team_member_info_validation.py similarity index 62% rename from tests/proxy_behavior/management/test_team_member_info_validation.py rename to tests/integration/management/test_team_member_info_validation.py index b22e3579400..131e5a1637e 100644 --- a/tests/proxy_behavior/management/test_team_member_info_validation.py +++ b/tests/integration/management/test_team_member_info_validation.py @@ -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" diff --git a/tests/proxy_behavior/management/test_team_member_me.py b/tests/integration/management/test_team_member_me.py similarity index 74% rename from tests/proxy_behavior/management/test_team_member_me.py rename to tests/integration/management/test_team_member_me.py index bfbbe0504ae..c1bee6c35f3 100644 --- a/tests/proxy_behavior/management/test_team_member_me.py +++ b/tests/integration/management/test_team_member_me.py @@ -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 diff --git a/tests/integration/management/test_team_member_reset_budget.py b/tests/integration/management/test_team_member_reset_budget.py new file mode 100644 index 00000000000..211d5a617f0 --- /dev/null +++ b/tests/integration/management/test_team_member_reset_budget.py @@ -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 diff --git a/tests/proxy_behavior/management/test_team_member_reset_spend.py b/tests/integration/management/test_team_member_reset_spend.py similarity index 58% rename from tests/proxy_behavior/management/test_team_member_reset_spend.py rename to tests/integration/management/test_team_member_reset_spend.py index fa7765ff6c3..5d8be6a5323 100644 --- a/tests/proxy_behavior/management/test_team_member_reset_spend.py +++ b/tests/integration/management/test_team_member_reset_spend.py @@ -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" diff --git a/tests/proxy_behavior/management/test_team_member_update.py b/tests/integration/management/test_team_member_update.py similarity index 74% rename from tests/proxy_behavior/management/test_team_member_update.py rename to tests/integration/management/test_team_member_update.py index 53b245bd1e9..4aaeab906b9 100644 --- a/tests/proxy_behavior/management/test_team_member_update.py +++ b/tests/integration/management/test_team_member_update.py @@ -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: diff --git a/tests/proxy_behavior/management/test_team_metadata_schema.py b/tests/integration/management/test_team_metadata_schema.py similarity index 57% rename from tests/proxy_behavior/management/test_team_metadata_schema.py rename to tests/integration/management/test_team_metadata_schema.py index c1d308eba33..a421a8c0a83 100644 --- a/tests/proxy_behavior/management/test_team_metadata_schema.py +++ b/tests/integration/management/test_team_metadata_schema.py @@ -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 diff --git a/tests/proxy_behavior/management/test_team_model.py b/tests/integration/management/test_team_model.py similarity index 75% rename from tests/proxy_behavior/management/test_team_model.py rename to tests/integration/management/test_team_model.py index 3564e8df83a..51a699df2c7 100644 --- a/tests/proxy_behavior/management/test_team_model.py +++ b/tests/integration/management/test_team_model.py @@ -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 diff --git a/tests/proxy_behavior/management/test_team_new.py b/tests/integration/management/test_team_new.py similarity index 70% rename from tests/proxy_behavior/management/test_team_new.py rename to tests/integration/management/test_team_new.py index 9846566d0b9..4b8f540407d 100644 --- a/tests/proxy_behavior/management/test_team_new.py +++ b/tests/integration/management/test_team_new.py @@ -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 diff --git a/tests/proxy_behavior/management/test_team_permissions.py b/tests/integration/management/test_team_permissions.py similarity index 52% rename from tests/proxy_behavior/management/test_team_permissions.py rename to tests/integration/management/test_team_permissions.py index 5d16702fe6c..4b33082c6d2 100644 --- a/tests/proxy_behavior/management/test_team_permissions.py +++ b/tests/integration/management/test_team_permissions.py @@ -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 diff --git a/tests/proxy_behavior/management/test_team_permissions_bulk_update.py b/tests/integration/management/test_team_permissions_bulk_update.py similarity index 57% rename from tests/proxy_behavior/management/test_team_permissions_bulk_update.py rename to tests/integration/management/test_team_permissions_bulk_update.py index 552c3526eea..365150087d8 100644 --- a/tests/proxy_behavior/management/test_team_permissions_bulk_update.py +++ b/tests/integration/management/test_team_permissions_bulk_update.py @@ -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}" diff --git a/tests/proxy_behavior/management/test_team_spend_by_user.py b/tests/integration/management/test_team_spend_by_user.py similarity index 84% rename from tests/proxy_behavior/management/test_team_spend_by_user.py rename to tests/integration/management/test_team_spend_by_user.py index 1d6aab04003..bd349e23681 100644 --- a/tests/proxy_behavior/management/test_team_spend_by_user.py +++ b/tests/integration/management/test_team_spend_by_user.py @@ -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: diff --git a/tests/proxy_behavior/management/test_team_update.py b/tests/integration/management/test_team_update.py similarity index 65% rename from tests/proxy_behavior/management/test_team_update.py rename to tests/integration/management/test_team_update.py index 50d6ec6ccaa..561191c9722 100644 --- a/tests/proxy_behavior/management/test_team_update.py +++ b/tests/integration/management/test_team_update.py @@ -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" diff --git a/tests/proxy_behavior/management/test_users_bulk_delete.py b/tests/integration/management/test_users_bulk_delete.py similarity index 61% rename from tests/proxy_behavior/management/test_users_bulk_delete.py rename to tests/integration/management/test_users_bulk_delete.py index 049234e737a..85f19cadffe 100644 --- a/tests/proxy_behavior/management/test_users_bulk_delete.py +++ b/tests/integration/management/test_users_bulk_delete.py @@ -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 diff --git a/tests/proxy_behavior/management/__init__.py b/tests/integration/spend/__init__.py similarity index 100% rename from tests/proxy_behavior/management/__init__.py rename to tests/integration/spend/__init__.py diff --git a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py b/tests/integration/spend/test_autorouter_session_rollup.py similarity index 87% rename from tests/proxy_behavior/spend/test_autorouter_session_rollup.py rename to tests/integration/spend/test_autorouter_session_rollup.py index 3425bed22b2..a3b959040b5 100644 --- a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py +++ b/tests/integration/spend/test_autorouter_session_rollup.py @@ -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( diff --git a/tests/proxy_behavior/spend/test_baseline_accounting.py b/tests/integration/spend/test_baseline_accounting.py similarity index 92% rename from tests/proxy_behavior/spend/test_baseline_accounting.py rename to tests/integration/spend/test_baseline_accounting.py index 4298b85cab4..60c5c32511a 100644 --- a/tests/proxy_behavior/spend/test_baseline_accounting.py +++ b/tests/integration/spend/test_baseline_accounting.py @@ -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: diff --git a/tests/proxy_behavior/spend/test_cache_activity.py b/tests/integration/spend/test_cache_activity.py similarity index 93% rename from tests/proxy_behavior/spend/test_cache_activity.py rename to tests/integration/spend/test_cache_activity.py index 528764fbf45..c381d51c402 100644 --- a/tests/proxy_behavior/spend/test_cache_activity.py +++ b/tests/integration/spend/test_cache_activity.py @@ -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( diff --git a/tests/proxy_behavior/auth/conftest.py b/tests/proxy_behavior/auth/conftest.py deleted file mode 100644 index 21982fa25cd..00000000000 --- a/tests/proxy_behavior/auth/conftest.py +++ /dev/null @@ -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() diff --git a/tests/proxy_behavior/management/actors.py b/tests/proxy_behavior/management/actors.py deleted file mode 100644 index 6c2f1a61ce1..00000000000 --- a/tests/proxy_behavior/management/actors.py +++ /dev/null @@ -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, - ) diff --git a/tests/proxy_behavior/management/conftest.py b/tests/proxy_behavior/management/conftest.py deleted file mode 100644 index 74db323e7f2..00000000000 --- a/tests/proxy_behavior/management/conftest.py +++ /dev/null @@ -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.`, 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}}, - ) diff --git a/tests/proxy_behavior/management/test_no_management_imports.py b/tests/proxy_behavior/management/test_no_management_imports.py deleted file mode 100644 index f8c52a1c37e..00000000000 --- a/tests/proxy_behavior/management/test_no_management_imports.py +++ /dev/null @@ -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) - ) diff --git a/tests/proxy_behavior/management/test_route_coverage.py b/tests/proxy_behavior/management/test_route_coverage.py deleted file mode 100644 index 1139e251a59..00000000000 --- a/tests/proxy_behavior/management/test_route_coverage.py +++ /dev/null @@ -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.(""` — `\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) diff --git a/tests/proxy_behavior/management/test_scratch_teardown.py b/tests/proxy_behavior/management/test_scratch_teardown.py deleted file mode 100644 index bcb53935558..00000000000 --- a/tests/proxy_behavior/management/test_scratch_teardown.py +++ /dev/null @@ -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 == [] diff --git a/tests/proxy_behavior/management/test_smoke.py b/tests/proxy_behavior/management/test_smoke.py deleted file mode 100644 index 4e90986ad9f..00000000000 --- a/tests/proxy_behavior/management/test_smoke.py +++ /dev/null @@ -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 diff --git a/tests/proxy_behavior/management/test_team_bulk_member_add.py b/tests/proxy_behavior/management/test_team_bulk_member_add.py deleted file mode 100644 index fc83cd414e5..00000000000 --- a/tests/proxy_behavior/management/test_team_bulk_member_add.py +++ /dev/null @@ -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 diff --git a/tests/proxy_behavior/management/test_team_member_reset_budget.py b/tests/proxy_behavior/management/test_team_member_reset_budget.py deleted file mode 100644 index 1e55b8b6b15..00000000000 --- a/tests/proxy_behavior/management/test_team_member_reset_budget.py +++ /dev/null @@ -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 diff --git a/tests/proxy_behavior/management/test_world_seed.py b/tests/proxy_behavior/management/test_world_seed.py deleted file mode 100644 index 00f9540c9c3..00000000000 --- a/tests/proxy_behavior/management/test_world_seed.py +++ /dev/null @@ -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 diff --git a/tests/proxy_behavior/spend/__init__.py b/tests/proxy_behavior/spend/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/proxy_behavior/spend/conftest.py b/tests/proxy_behavior/spend/conftest.py deleted file mode 100644 index 0b918401eac..00000000000 --- a/tests/proxy_behavior/spend/conftest.py +++ /dev/null @@ -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() diff --git a/tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py b/tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py index ea528878db8..69188ebd5a2 100644 --- a/tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py +++ b/tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py @@ -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} diff --git a/tests/unit/proxy/db/test_autorouter_session_rollup.py b/tests/unit/proxy/db/test_autorouter_session_rollup.py index 686c5a3a970..4217feeb83b 100644 --- a/tests/unit/proxy/db/test_autorouter_session_rollup.py +++ b/tests/unit/proxy/db/test_autorouter_session_rollup.py @@ -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. """ diff --git a/tests/unit/proxy/management_endpoints/test_encryption_endpoints.py b/tests/unit/proxy/management_endpoints/test_encryption_endpoints.py index e92cb3b13a7..9398bd795fc 100644 --- a/tests/unit/proxy/management_endpoints/test_encryption_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_encryption_endpoints.py @@ -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 diff --git a/tests/unit/proxy/management_endpoints/test_management_route_coverage.py b/tests/unit/proxy/management_endpoints/test_management_route_coverage.py new file mode 100644 index 00000000000..64e584b4f70 --- /dev/null +++ b/tests/unit/proxy/management_endpoints/test_management_route_coverage.py @@ -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) diff --git a/tests/unit/test_assert_ci_coverage.py b/tests/unit/test_assert_ci_coverage.py index a3c2e933a46..443adf86d9a 100644 --- a/tests/unit/test_assert_ci_coverage.py +++ b/tests/unit/test_assert_ci_coverage.py @@ -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",