diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index 98e683679d3..fc5fc0b128d 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -161,6 +161,31 @@ class Scenario: self.gateway.post("/budget/delete", {"id": identity}) assert read_rows('SELECT budget_id FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (identity,)) == [] + def organization(self, **fields: JsonValue) -> str: + created: Final = self.gateway.post( + "/organization/new", {"organization_alias": f"integration-{uuid.uuid4().hex}", **fields} + ) + identity: Final = string_value(created["organization_id"]) + self.cleanups.callback(self.delete_organization, identity, string_value(created["budget_id"])) + return identity + + def delete_organization(self, identity: str, budget_id: str) -> None: + response: Final = self.gateway.request("DELETE", "/organization/delete", {"organization_ids": [identity]}) + assert response.status_code == 200, response.text + assert ( + read_rows('SELECT organization_id FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', (identity,)) + == [] + ) + self.delete_budget(budget_id) + + def org_member(self, organization_id: str, role: str) -> str: + user_id: Final = self.user(user_role="internal_user") + self.gateway.post( + "/organization/member_add", + {"organization_id": organization_id, "member": {"role": role, "user_id": user_id}}, + ) + return user_id + def user(self, **fields: JsonValue) -> str: created: Final = self.gateway.post( "/user/new", {"user_id": f"integration-{uuid.uuid4().hex}", "auto_create_key": False, **fields} diff --git a/tests/integration/authorization/test_team_admin_gate.py b/tests/integration/authorization/test_team_admin_gate.py index cce23866035..2b02e90fcc5 100644 --- a/tests/integration/authorization/test_team_admin_gate.py +++ b/tests/integration/authorization/test_team_admin_gate.py @@ -1,7 +1,9 @@ """Status-code matrix for every management route that admits a team admin today. Each route is called as a proxy admin, an admin of the target team, a plain member, an admin of another team -and a teamless user. The expected codes pin current behaviour so the shared team-admin gate can prove parity. +and a teamless user. It is called again on a team that belongs to an organization, as an admin of that +organization and as an admin of another organization. The expected codes pin current behaviour so the shared +team-admin gate can prove parity. """ from __future__ import annotations @@ -28,8 +30,17 @@ from tests.integration._support.client import ( ) from tests.integration._support.database import read_rows -Caller = Literal["proxy_admin", "team_admin", "member", "other_team_admin", "outsider"] -CALLERS: Final[tuple[Caller, ...]] = ("proxy_admin", "team_admin", "member", "other_team_admin", "outsider") +Caller = Literal["proxy_admin", "team_admin", "member", "other_team_admin", "outsider", "org_admin", "other_org_admin"] +CALLERS: Final[tuple[Caller, ...]] = ( + "proxy_admin", + "team_admin", + "member", + "other_team_admin", + "outsider", + "org_admin", + "other_org_admin", +) +ORG_CALLERS: Final[frozenset[Caller]] = frozenset({"org_admin", "other_org_admin"}) @dataclass(frozen=True, slots=True) @@ -45,7 +56,6 @@ class TeamScenario: scenario: Scenario team_id: str - other_team_id: str keys: Mapping[Caller, str] request_id: str since: datetime @@ -104,6 +114,8 @@ class Route: member: int | None = None other_team_admin: int | None = None outsider: int | None = None + org_admin: int | None = None + other_org_admin: int | None = None permission: str = "" cleanup: Callable[[TeamScenario, dict[str, JsonValue]], None] | None = None @@ -119,6 +131,10 @@ class Route: return self.others if self.other_team_admin is None else self.other_team_admin case "outsider": return self.others if self.outsider is None else self.outsider + case "org_admin": + return self.others if self.org_admin is None else self.org_admin + case "other_org_admin": + return self.others if self.other_org_admin is None else self.other_org_admin case _: assert_never(caller) @@ -186,48 +202,48 @@ def _delete_model(s: TeamScenario, created: dict[str, JsonValue]) -> None: ROUTES: Final[tuple[Route, ...]] = ( Route("member_add_user", lambda s: Call("POST", "/team/member_add", {"team_id": s.team_id, "member": {"role": "user", "user_id": s.user()}}), - team_admin=200, others=403), + team_admin=200, others=403, org_admin=200), Route("member_add_admin", lambda s: Call("POST", "/team/member_add", {"team_id": s.team_id, "member": {"role": "admin", "user_id": s.user()}}), - team_admin=200, others=403), + team_admin=200, others=403, org_admin=200), Route("member_update_budget", lambda s: Call("POST", "/team/member_update", {"team_id": s.team_id, "user_id": s.member(), "max_budget_in_team": 5}), - team_admin=200, others=403), + team_admin=200, others=403, org_admin=200), Route("member_update_role_admin", lambda s: Call("POST", "/team/member_update", {"team_id": s.team_id, "user_id": s.member(), "role": "admin"}), - team_admin=200, others=403), + team_admin=200, others=403, org_admin=200), Route("member_delete", lambda s: Call("POST", "/team/member_delete", {"team_id": s.team_id, "user_id": s.member()}), - team_admin=200, others=403), + team_admin=200, others=403, org_admin=200), Route("members_bulk_delete", lambda s: Call("POST", f"/management/v1/teams/{s.team_id}/members/bulk_delete", {"members": [{"user_id": s.member()}]}), - team_admin=200, others=403), + team_admin=200, others=403, org_admin=200), Route("members_bulk_update", lambda s: Call("POST", f"/management/v1/teams/{s.team_id}/members/bulk_update", {"members": [{"user_id": s.member(), "max_budget_in_team": 10}]}), - team_admin=200, others=403), + team_admin=200, others=403, org_admin=200), Route("member_reset_spend", lambda s: Call("POST", f"/team/{s.team_id}/member/{s.member()}/reset_spend", {"reset_to": 0}), - team_admin=200, others=403), + team_admin=200, others=403, org_admin=200), Route("member_reset_budget", lambda s: Call("POST", f"/team/{s.team_id}/member/{s.member()}/reset_budget"), - team_admin=200, others=403), + team_admin=200, others=403, org_admin=200), Route("invitation_new", lambda s: Call("POST", "/invitation/new", {"user_id": s.member()}), team_admin=200, others=400), Route("invitation_delete", lambda s: Call("POST", "/invitation/delete", {"invitation_id": s.invitation()}), - team_admin=200, others=400, other_team_admin=403), + team_admin=200, others=400, other_team_admin=403, org_admin=403, other_org_admin=403), Route("user_info_v2", lambda s: Call("GET", f"/v2/user/info?user_id={s.member()}"), team_admin=200, others=404), Route("permissions_update", lambda s: Call("POST", "/team/permissions_update", {"team_id": s.team_id, "team_member_permissions": ["/key/info", "/key/health"]}), - team_admin=200, others=403), + team_admin=200, others=403, org_admin=200), Route("permissions_list", lambda s: Call("GET", f"/team/permissions_list?team_id={s.team_id}"), - team_admin=200, others=403), + team_admin=200, others=403, org_admin=200), Route("key_generate_team", lambda s: Call("POST", "/key/generate", {"team_id": s.team_id}), team_admin=200, others=400, member=401, cleanup=_delete_key), @@ -258,10 +274,10 @@ ROUTES: Final[tuple[Route, ...]] = ( team_admin=200, others=403), Route("key_block", lambda s: Call("POST", "/key/block", {"key": s.member_key()}), - team_admin=200, others=403), + team_admin=200, others=403, org_admin=200), Route("key_unblock", lambda s: Call("POST", "/key/unblock", {"key": s.member_key()}), - team_admin=200, others=403), + team_admin=200, others=403, org_admin=200), Route("key_list_team", lambda s: Call("GET", f"/key/list?team_id={s.team_id}&include_team_keys=true&return_full_object=true"), team_admin=200, others=403, member=200), @@ -291,25 +307,25 @@ ROUTES: Final[tuple[Route, ...]] = ( team_admin=200, others=403), Route("callback_add", lambda s: Call("POST", f"/team/{s.team_id}/callback", _callback_body(s.callback_name())), - team_admin=200, others=403), + team_admin=200, others=403, org_admin=200), Route("callback_get", lambda s: Call("GET", f"/team/{s.team_id}/callback"), - team_admin=200, others=403), + team_admin=200, others=403, org_admin=200), Route("callback_delete", lambda s: Call("DELETE", f"/team/{s.team_id}/callback/{s.callback()}"), - team_admin=200, others=403), + team_admin=200, others=403, org_admin=200), Route("disable_logging", lambda s: Call("POST", f"/team/{s.team_id}/disable_logging"), team_admin=401, others=401), Route("team_info", lambda s: Call("GET", f"/team/info?team_id={s.team_id}"), - team_admin=200, others=403, member=200), + team_admin=200, others=403, member=200, org_admin=200), Route("team_update_budget", lambda s: Call("POST", "/team/update", {"team_id": s.team_id, "max_budget": 5}), - team_admin=403, others=403), + team_admin=403, others=403, org_admin=200), Route("team_update_budget_permitted", lambda s: Call("POST", "/team/update", {"team_id": s.team_id, "max_budget": 7}), - team_admin=200, others=403, permission="max_budget"), + team_admin=200, others=403, org_admin=200, permission="max_budget"), Route("project_new", lambda s: Call("POST", "/project/new", {"team_id": s.team_id, "project_alias": f"matrix-{uuid.uuid4().hex}"}), team_admin=403, others=403, cleanup=_delete_project), @@ -336,6 +352,23 @@ def _cases() -> Iterator[tuple[Route, Caller]]: CASES: Final = tuple(_cases()) +def _team_scenario(scenario: Scenario, team_id: str, keys: Mapping[Caller, str]) -> TeamScenario: + since: Final = datetime.now(timezone.utc) - timedelta(days=1) + until: Final = since + timedelta(days=2) + scenario.gateway.chat(scenario.model(), key=keys["team_admin"]) + rows: Final = eventually( + lambda: _spend_rows(scenario.gateway, team_id, since, until), lambda found: len(found) > 0, seconds=30 + ) + return TeamScenario( + scenario=scenario, + team_id=team_id, + keys=keys, + request_id=string_value(object_value(rows[0])["request_id"]), + since=since, + until=until, + ) + + @pytest.fixture(scope="module") def shared() -> Iterator[TeamScenario]: with gateway_from_environment() as gateway, gateway.scenario() as scenario: @@ -354,27 +387,33 @@ def shared() -> Iterator[TeamScenario]: "outsider": scenario.key(user_id=outsider), } ) - since: Final = datetime.now(timezone.utc) - timedelta(days=1) - until: Final = since + timedelta(days=2) - gateway.chat(scenario.model(), key=keys["team_admin"]) - rows: Final = eventually( - lambda: _spend_rows(gateway, team_id, since, until), lambda found: len(found) > 0, seconds=30 - ) - yield TeamScenario( - scenario=scenario, - team_id=team_id, - other_team_id=other_team_id, - keys=keys, - request_id=string_value(object_value(rows[0])["request_id"]), - since=since, - until=until, + yield _team_scenario(scenario, team_id, keys) + + +@pytest.fixture(scope="module") +def org_team() -> Iterator[TeamScenario]: + with gateway_from_environment() as gateway, gateway.scenario() as scenario: + organization_id: Final = scenario.organization() + other_organization_id: Final = scenario.organization() + team_id: Final = scenario.team(organization_id=organization_id) + team_admin: Final = scenario.member(team_id, role="admin") + org_admin: Final = scenario.org_member(organization_id, role="org_admin") + other_org_admin: Final = scenario.org_member(other_organization_id, role="org_admin") + keys: Final[Mapping[Caller, str]] = MappingProxyType( + { + "team_admin": scenario.key(user_id=team_admin, team_id=team_id), + "org_admin": scenario.key(user_id=org_admin), + "other_org_admin": scenario.key(user_id=other_org_admin), + } ) + yield _team_scenario(scenario, team_id, keys) @pytest.mark.parametrize(("route", "caller"), CASES, ids=tuple(f"{route.name}[{caller}]" for route, caller in CASES)) -def test_status_code(shared: TeamScenario, route: Route, caller: Caller) -> None: - with shared.gateway.scenario() as scenario: - s: Final = replace(shared, scenario=scenario) +def test_status_code(shared: TeamScenario, org_team: TeamScenario, route: Route, caller: Caller) -> None: + team: Final = org_team if caller in ORG_CALLERS else shared + with team.gateway.scenario() as scenario: + s: Final = replace(team, scenario=scenario) if route.permission: scenario.cleanups.enter_context(team_admin_permissions(s.gateway, (route.permission,))) call: Final = route.call(s)