diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 5520b44993d..f8b5aa42900 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -11,6 +11,11 @@ from dataclasses import dataclass from e2e_gateway import Gateway, build_gateway from e2e_http import NoBody, ProbeResult, StreamingResponse, unwrap from models import ( + BudgetData, + BudgetDeleteBody, + BudgetInfoBody, + BudgetInfoResponse, + BudgetNewBody, ChatBody, ChatMessage, KeyDeleteBody, @@ -43,6 +48,8 @@ from models import ( MODEL_ACCESS_DENIED_MARKER = "key_model_access_denied" ROUTE_NOT_ALLOWED_MARKER = "not allowed to call this route" +PROXY_ADMIN_REQUIRED_MARKER = "Only proxy admin can be used" +TEAM_ADMIN_REQUIRED_MARKER = "not proxy admin OR team admin" @dataclass(frozen=True, slots=True) @@ -201,6 +208,35 @@ class ManagementClient: ) ) + def create_budget(self, body: BudgetNewBody) -> str: + return unwrap( + self.gateway.transport.post( + "/budget/new", + headers=self.gateway.transport.master, + json=body, + response_type=BudgetData, + ) + ).budget_id + + def delete_budget(self, budget_id: str) -> None: + _ = self.gateway.transport.post( + "/budget/delete", + headers=self.gateway.transport.master, + json=BudgetDeleteBody(id=budget_id), + response_type=NoBody, + ) + + def budget_row(self, budget_id: str) -> BudgetData | None: + rows = unwrap( + self.gateway.transport.post( + "/budget/info", + headers=self.gateway.transport.master, + json=BudgetInfoBody(budgets=[budget_id]), + response_type=BudgetInfoResponse, + ) + ).root + return next((row for row in rows if row.budget_id == budget_id), None) + def chat_status(self, key: str, model: str, content: str) -> StreamingResponse: return self.gateway.transport.send( "/chat/completions", @@ -217,6 +253,15 @@ class ManagementClient: def user_new_status(self, key: str, body: UserNewBody) -> StreamingResponse: return self.gateway.transport.send("/user/new", headers=self.gateway.transport.bearer(key), json=body) + def budget_new_status(self, key: str, body: BudgetNewBody) -> StreamingResponse: + return self.gateway.transport.send("/budget/new", headers=self.gateway.transport.bearer(key), json=body) + + def team_member_add_status(self, key: str, body: TeamMemberAddBody) -> StreamingResponse: + return self.gateway.transport.send("/team/member_add", headers=self.gateway.transport.bearer(key), json=body) + + def team_member_delete_status(self, key: str, body: TeamMemberDeleteBody) -> StreamingResponse: + return self.gateway.transport.send("/team/member_delete", headers=self.gateway.transport.bearer(key), json=body) + def build_client() -> ManagementClient: return ManagementClient(gateway=build_gateway()) diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index 3beb039b8bd..b1bf84b3ad1 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -19,10 +19,21 @@ from e2e_http import StreamingResponse from lifecycle import ResourceManager from management_client import ( MODEL_ACCESS_DENIED_MARKER, + PROXY_ADMIN_REQUIRED_MARKER, ROUTE_NOT_ALLOWED_MARKER, + TEAM_ADMIN_REQUIRED_MARKER, ManagementClient, ) -from models import KeyGenerateBody, OrgNewBody, TeamNewBody, UserNewBody +from models import ( + BudgetNewBody, + KeyGenerateBody, + OrgNewBody, + TeamMemberAddBody, + TeamMemberDeleteBody, + TeamMemberEntry, + TeamNewBody, + UserNewBody, +) pytestmark = pytest.mark.e2e @@ -54,6 +65,36 @@ def _create_user(client: ManagementClient, resources: ResourceManager, body: Use return user_id +def _create_budget(client: ManagementClient, resources: ResourceManager, body: BudgetNewBody) -> str: + budget_id = client.create_budget(body) + resources.defer(lambda: client.delete_budget(budget_id)) + return budget_id + + +def _create_nonadmin_key(client: ManagementClient, resources: ResourceManager) -> str: + """A key bound to a fresh internal_user: it inherits that user's non-admin + role, so route_checks denies it the proxy-admin-only management writes.""" + user_id = _create_user( + client, + resources, + UserNewBody(user_email=f"e2e-mgmt-nonadmin-{unique_marker()}@example.com", user_role="internal_user"), + ) + return _generate_key(client, resources, KeyGenerateBody(user_id=user_id)) + + +def _add_team_member(client: ManagementClient, resources: ResourceManager, team_id: str) -> tuple[str, str]: + """Create an internal user, add them to `team_id` as a plain member (role + 'user', not team admin), and return their (user_id, key scoped to the team).""" + user_id = _create_user( + client, + resources, + UserNewBody(user_email=f"e2e-mgmt-member-{unique_marker()}@example.com", user_role="internal_user"), + ) + client.add_team_member(team_id, user_id) + key = _generate_key(client, resources, KeyGenerateBody(user_id=user_id, team_id=team_id)) + return user_id, key + + def _is_model_denial(outcome: StreamingResponse) -> bool: return outcome.status_code == 403 and MODEL_ACCESS_DENIED_MARKER in outcome.body @@ -104,6 +145,7 @@ def _poll_model_access_granted(client: ManagementClient, key: str, model: str) - class TestKeyRoutes: @pytest.mark.covers("mgmt.key.generate.persists") + @pytest.mark.covers("mgmt.key.info.persists") def test_generate_persists_to_key_info_and_scopes_chat( self, client: ManagementClient, resources: ResourceManager ) -> None: @@ -167,7 +209,7 @@ class TestKeyRoutes: class TestTeamRoutes: - @pytest.mark.covers("management.team.new.persists") + @pytest.mark.covers("mgmt.team.new.persists") def test_new_persists_to_team_info_and_binds_keys( self, client: ManagementClient, resources: ResourceManager ) -> None: @@ -187,6 +229,7 @@ class TestTeamRoutes: ) @pytest.mark.covers("mgmt.team.member_add.persists") + @pytest.mark.covers("mgmt.team.member_delete.persists") def test_member_add_and_delete_persist_to_team_info( self, client: ManagementClient, resources: ResourceManager ) -> None: @@ -210,9 +253,57 @@ class TestTeamRoutes: f"/team/info still lists {user_id} after /team/member_delete" ) + @pytest.mark.covers("mgmt.team.new.admin_only") + def test_new_forbidden_for_non_admin(self, client: ManagementClient, resources: ResourceManager) -> None: + key = _create_nonadmin_key(client, resources) + team_id = f"e2e-mgmt-team-forbidden-{unique_marker()}" + + _assert_admin_only( + "/team/new", client.team_new_status(key, TeamNewBody(team_alias=team_id, team_id=team_id)) + ) + + probe = client.team_info_status(team_id) + assert probe.status_code == 404, ( + f"team {team_id} was created despite the 401 admin-only denial: " + f"/team/info returned {probe.status_code}: {probe.body[:300]}" + ) + + @pytest.mark.covers("mgmt.team.member_add.member_forbidden") + @pytest.mark.covers("mgmt.team.member_delete.member_forbidden") + def test_member_writes_forbidden_for_non_admin_member( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + team_id = _create_team(client, resources, f"e2e-mgmt-team-{unique_marker()}", ["gemini-2.5-flash"]) + member_id, member_key = _add_team_member(client, resources, team_id) + outsider_id = _create_user( + client, + resources, + UserNewBody(user_email=f"e2e-mgmt-outsider-{unique_marker()}@example.com", user_role="internal_user"), + ) + + _assert_team_admin_required( + "/team/member_add", + client.team_member_add_status( + member_key, + TeamMemberAddBody(team_id=team_id, member=TeamMemberEntry(role="user", user_id=outsider_id)), + ), + ) + _assert_team_admin_required( + "/team/member_delete", + client.team_member_delete_status(member_key, TeamMemberDeleteBody(team_id=team_id, user_id=member_id)), + ) + + roster = {entry.user_id for entry in client.team_info(team_id).members_with_roles} + assert outsider_id not in roster, ( + f"outsider {outsider_id} was added to team {team_id} despite the 403 team-admin-required denial" + ) + assert member_id in roster, ( + f"member {member_id} was removed from team {team_id} despite the 403 team-admin-required denial" + ) + class TestUserRoutes: - @pytest.mark.covers("mgmt.user.new.persists") + @pytest.mark.covers("mgmt.user.new.happy_path") def test_new_persists_to_user_info(self, client: ManagementClient, resources: ResourceManager) -> None: email = f"e2e-mgmt-{unique_marker()}@example.com" user_id = _create_user(client, resources, UserNewBody(user_email=email, user_role="internal_user")) @@ -225,7 +316,7 @@ class TestUserRoutes: class TestOrganizationRoutes: - @pytest.mark.covers("mgmt.organization.new.persists") + @pytest.mark.covers("mgmt.organization.new.happy_path") def test_new_persists_to_organization_info( self, client: ManagementClient, resources: ResourceManager ) -> None: @@ -242,6 +333,41 @@ class TestOrganizationRoutes: ) +class TestBudgetRoutes: + @pytest.mark.covers("mgmt.budget.new.persists") + def test_new_persists_to_budget_info(self, client: ManagementClient, resources: ResourceManager) -> None: + budget_id = f"e2e-mgmt-budget-{unique_marker()}" + _create_budget( + client, + resources, + BudgetNewBody(budget_id=budget_id, max_budget=42.5, soft_budget=21.25, budget_duration="30d"), + ) + + row = client.budget_row(budget_id) + assert row is not None, f"/budget/info does not list {budget_id} after /budget/new" + assert row.max_budget == 42.5, f"/budget/info reports max_budget {row.max_budget}, configured 42.5" + assert row.soft_budget == 21.25, f"/budget/info reports soft_budget {row.soft_budget}, configured 21.25" + assert row.budget_duration == "30d", ( + f"/budget/info reports budget_duration {row.budget_duration!r}, configured '30d'" + ) + assert row.budget_reset_at is not None, ( + "a budget with a 30d duration must persist a computed budget_reset_at, got None" + ) + + @pytest.mark.covers("mgmt.budget.new.admin_only") + def test_new_forbidden_for_non_admin(self, client: ManagementClient, resources: ResourceManager) -> None: + key = _create_nonadmin_key(client, resources) + budget_id = f"e2e-mgmt-budget-forbidden-{unique_marker()}" + + _assert_admin_only( + "/budget/new", client.budget_new_status(key, BudgetNewBody(budget_id=budget_id, max_budget=1)) + ) + + assert client.budget_row(budget_id) is None, ( + f"budget {budget_id} was created despite the 401 admin-only denial" + ) + + def _assert_route_forbidden(route: str, outcome: StreamingResponse) -> None: assert outcome.status_code == 403, ( f"llm-only key POSTing {route} must be denied exactly 403, got {outcome.status_code}: {outcome.body[:300]}" @@ -251,8 +377,26 @@ def _assert_route_forbidden(route: str, outcome: StreamingResponse) -> None: ) +def _assert_admin_only(route: str, outcome: StreamingResponse) -> None: + assert outcome.status_code == 401, ( + f"non-admin key POSTing {route} must be denied 401, got {outcome.status_code}: {outcome.body[:300]}" + ) + assert PROXY_ADMIN_REQUIRED_MARKER in outcome.body, ( + f"{route} denial body must be a proxy-admin-only error, got: {outcome.body[:300]}" + ) + + +def _assert_team_admin_required(route: str, outcome: StreamingResponse) -> None: + assert outcome.status_code == 403, ( + f"non-admin team member POSTing {route} must be denied 403, got {outcome.status_code}: {outcome.body[:300]}" + ) + assert TEAM_ADMIN_REQUIRED_MARKER in outcome.body, ( + f"{route} denial body must be a team-admin-required error, got: {outcome.body[:300]}" + ) + + class TestManagementRoutePermissions: - @pytest.mark.covers("mgmt.key.generate.member_forbidden") + @pytest.mark.covers("other.auth.virtual_key.route_permission_enforced") def test_llm_only_key_forbidden_from_management_writes( self, client: ManagementClient, resources: ResourceManager ) -> None: diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 0490db286ea..48c783400ce 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -535,3 +535,34 @@ class OrgInfoResponse(BaseModel): class OrgDeleteBody(BaseModel): organization_ids: list[str] + + +# ---------- budget management ---------- + + +class BudgetNewBody(BaseModel): + budget_id: str + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + + +class BudgetInfoBody(BaseModel): + budgets: list[str] + + +class BudgetData(BaseModel): + budget_id: str + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + budget_reset_at: str | None = None + + +class BudgetInfoResponse(RootModel[list[BudgetData]]): + """POST /budget/info answers with a bare array of budget rows (one per queried + id), not an object wrapping them. Read the rows off .root.""" + + +class BudgetDeleteBody(BaseModel): + id: str