diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index cbc79ee2c2a..def9b9e30c2 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -307,6 +307,7 @@ class McpClient: server_id: str, name: str, arguments: McpToolArguments, + headers: AuthHeaders | None = None, ) -> McpCallToolResponse: """Poll tools/call until the result is not a multi-worker registry miss. @@ -317,7 +318,7 @@ class McpClient: deadline = time.monotonic() + self.proxy.poll_timeout last: Result[McpCallToolResponse] | None = None while True: - last = self.call_tool(key, server_id=server_id, name=name, arguments=arguments) + last = self.call_tool(key, server_id=server_id, name=name, arguments=arguments, headers=headers) if not _is_mcp_not_synced(last, tool_name=name): return unwrap(last) if time.monotonic() >= deadline: diff --git a/tests/e2e/mcp/test_mcp_jwt_auth_e2e.py b/tests/e2e/mcp/test_mcp_jwt_auth_e2e.py index 0b25ffd1f69..c0c39fc1858 100644 --- a/tests/e2e/mcp/test_mcp_jwt_auth_e2e.py +++ b/tests/e2e/mcp/test_mcp_jwt_auth_e2e.py @@ -12,7 +12,7 @@ from idp import SHORT_LIVED_CLIENT_ID, Identity, Keycloak, token_claims from lifecycle import ResourceManager from management.management_client import build_client as build_management_client from mcp_client import McpClient -from models import KeyGenerateBody, ObjectPermission, UserScimMetadata, UserUpdateBody +from models import KeyGenerateBody, ObjectPermission, TeamUpdateBody, UserScimMetadata, UserUpdateBody pytestmark = pytest.mark.e2e @@ -29,14 +29,12 @@ def _allowed(client: McpClient, access: GrantedMcp, headers: AuthHeaders) -> Non listing: Final = unwrap(client.list_tools(access.token, headers=headers)) assert listing.tool_names_for_server(access.server_id) == frozenset((access.tool,)) assert listing.tool_names_for_server(access.other_server_id) == frozenset() - result: Final = unwrap( - client.call_tool( - access.token, - server_id=access.server_id, - name=access.tool, - arguments={"query": f"lit4506-{unique_marker()}", "from": DD_SEARCH_FROM, "to": "now", "max_tokens": 1000}, - headers=headers, - ) + result: Final = client.await_call_tool( + access.token, + server_id=access.server_id, + name=access.tool, + arguments={"query": f"lit4506-{unique_marker()}", "from": DD_SEARCH_FROM, "to": "now", "max_tokens": 1000}, + headers=headers, ) assert result.is_error is False, result assert result.all_text.strip(), "The successful control must return a real tool result" @@ -64,6 +62,13 @@ def granted_mcp(client: McpClient, resources: ResourceManager, jwt_identity: Ide client.await_registered(server_id) client.await_registered(other_server_id) management: Final = build_management_client(client.proxy) + management.update_team( + TeamUpdateBody( + team_id=jwt_identity.group, + team_alias=jwt_identity.group, + object_permission=ObjectPermission(mcp_servers=[server_id]), + ) + ) management.update_user( UserUpdateBody( user_id=jwt_identity.user_id, diff --git a/tests/e2e/models.py b/tests/e2e/models.py index eaf6de0bdd1..502ee72f755 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1174,6 +1174,7 @@ class TeamUpdateBody(BaseModel): team_id: str team_alias: str models: list[str] | None = None + object_permission: ObjectPermission | None = None class TeamInfoParams(BaseModel):