From 4c3879a88e5ea261ecc696bf4f6a6190897b109b Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 19 May 2026 21:31:38 -0700 Subject: [PATCH] test(proxy_behavior): seed 8-actor read-world for the authz matrix MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Slice 4 of the management-endpoints behavior-pinning effort. New ``actors.py`` defines the actor enum + seeds an immutable world (2 orgs, 2 teams, 8 users, 8 verification tokens) under the ``behavior-pin-`` prefix so the rows are identifiable in psql and ``_wipe_world`` is targeted. Each actor key is created with its cleartext form generated locally and its hashed form (via ``litellm.proxy.utils.hash_token``) stored in ``LiteLLM_VerificationToken`` — so the real ``user_api_key_auth`` accepts the cleartext bearer token. Roles, ``team_id``, ``organization_id``, and the service-account metadata flag are all set on the seeded rows so the auth layer resolves the same scopes a real proxy would. The session-scoped ``world`` fixture re-seeds at session start (idempotent via wipe-then-create), and the smoke test confirms each of the 8 actor keys can call ``/key/info`` on itself and receive its own row back. Plan: https://www.notion.so/36643b8acdab8128a581ced0f6a4744d --- tests/proxy_behavior/management/actors.py | 273 ++++++++++++++++++ tests/proxy_behavior/management/conftest.py | 24 ++ .../management/test_world_seed.py | 37 +++ 3 files changed, 334 insertions(+) create mode 100644 tests/proxy_behavior/management/actors.py create mode 100644 tests/proxy_behavior/management/test_world_seed.py diff --git a/tests/proxy_behavior/management/actors.py b/tests/proxy_behavior/management/actors.py new file mode 100644 index 00000000000..6ac8ba4e303 --- /dev/null +++ b/tests/proxy_behavior/management/actors.py @@ -0,0 +1,273 @@ +"""Immutable read-world seed for behavior-pinning tests. + +Seeds 8 actor profiles across 2 orgs / 2 teams so authz matrix tests can drive +``(actor, target, expected)`` tuples against the real proxy. Each actor has +exactly one virtual key (cleartext + hashed) so tests pass the cleartext as a +``Bearer`` token and the real auth stack accepts it. + +All seeded rows are namespaced under the ``behavior-pin-`` prefix on their +primary key, so the seed is identifiable in psql and the wipe step is targeted. +""" + +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" + + +# Stable IDs so re-seed is idempotent and the world is identifiable in psql. +PREFIX = "behavior-pin-" +ORG_A = PREFIX + "org-a" +ORG_B = PREFIX + "org-b" +TEAM_ALPHA = PREFIX + "team-alpha" +TEAM_BETA = PREFIX + "team-beta" +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 + keys: Dict[Actor, SeededKey] + + +def _new_clear_key() -> str: + return "sk-" + uuid.uuid4().hex + + +def _actor_profile() -> Dict[Actor, Dict[str, Any]]: + """Per-actor (role, scoping) profile used for both the user row and its key. + + Scoping rules: + - PROXY_ADMIN: global, no team/org scope on its key. + - ORG_ADMIN: org_a, no team scope. + - TEAM_ADMIN / INTERNAL_USER / OWNER / UNRELATED_SAME_ORG / SERVICE_ACCOUNT: + team_alpha within org_a. + - CROSS_ORG_USER: team_beta within org_b. + + The auth layer reads ``user_role`` off the user row pointed at by the key's + ``user_id``, so setting it once on the user is enough — keys do not need a + separate role field. + """ + 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, + }, + } + + +async def _wipe_world(prisma: PrismaClient) -> None: + """Delete prior seed rows so re-seed is idempotent across sessions.""" + 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) + + # Budget — orgs require a non-null budget_id. + await prisma.db.litellm_budgettable.create( + data={ + "budget_id": BUDGET_ID, + "created_by": "behavior-pin-seeder", + "updated_by": "behavior-pin-seeder", + } + ) + + # Orgs. + 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} + + # Users. + 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, + } + ) + + # Teams. + 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], + ], + } + ) + 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]], + } + ) + + # Org memberships (UI lookups + ORG_ADMIN scoping use these). + 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"), + ]: + await prisma.db.litellm_organizationmembership.create( + data={ + "user_id": user_ids[actor], + "organization_id": org_id, + "user_role": role, + } + ) + + # Team memberships. + 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} + ) + + # Verification tokens — one cleartext per actor, hashed via the real + # credential boundary so user_api_key_auth accepts it. + 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 is non-Optional on models — Postgres lets the + # column be NULL, but the pydantic view used by user_api_key_auth rejects None. + "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: + # LiteLLM convention: service-account keys carry an explicit metadata flag. + 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, + keys=keys, + ) diff --git a/tests/proxy_behavior/management/conftest.py b/tests/proxy_behavior/management/conftest.py index 7deab9ba5d7..c16adc39db1 100644 --- a/tests/proxy_behavior/management/conftest.py +++ b/tests/proxy_behavior/management/conftest.py @@ -62,3 +62,27 @@ async def proxy_client(proxy_app) -> AsyncIterator[httpx.AsyncClient]: transport=transport, base_url="http://testserver" ) as client: yield client + + +@pytest_asyncio.fixture(scope="session") +async def prisma(proxy_app): + """The connected PrismaClient the lifespan opened.""" + from litellm.proxy import proxy_server + + assert ( + proxy_server.prisma_client is not None + ), "FastAPI lifespan did not connect prisma — harness is wrong." + return proxy_server.prisma_client + + +@pytest_asyncio.fixture(scope="session") +async def world(prisma): + """The immutable read-world seed. + + Re-seeds at session start so each pytest invocation gets a clean world. + Tests must not mutate these rows; write tests use Slice 5's namespace + + teardown fixtures for scratch entities. + """ + from .actors import seed_world + + return await seed_world(prisma) diff --git a/tests/proxy_behavior/management/test_world_seed.py b/tests/proxy_behavior/management/test_world_seed.py new file mode 100644 index 00000000000..5b6abb12d1c --- /dev/null +++ b/tests/proxy_behavior/management/test_world_seed.py @@ -0,0 +1,37 @@ +"""Slice 4 smoke: every seeded actor key can authenticate and call /key/info on itself. + +This is the minimum proof that the world is reachable via the real auth stack — +the real ``user_api_key_auth`` dependency hashes the cleartext token from the +Bearer header, looks it up, resolves the user/role, and the handler returns the +key info row. + +If any actor 401/403s here, the seed is wrong (user_role mismatch, scope +fields, etc.) before we scale out to the matrix in Slices 7–12. +""" + +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.status_code} {resp.text}" + body = resp.json() + # /key/info returns {"key": , "info": }. + assert body.get("key") == seeded.hashed, ( + f"{actor.value}: /key/info returned the wrong key " + f"(got {body.get('key')!r}, expected {seeded.hashed!r})" + ) + info = body["info"] + assert info.get("user_id") == seeded.user_id, ( + f"{actor.value}: /key/info returned the wrong user_id " + f"(got {info.get('user_id')!r}, expected {seeded.user_id!r})" + )