From 854266d56282af174b66298b127d29fdffe2186f Mon Sep 17 00:00:00 2001
From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Date: Thu, 17 Sep 2026 17:41:03 -0700
Subject: [PATCH] test(mcp): reuse the successful catalog snapshot for toolset
 controls

---
 tests/e2e/mcp/mcp_client.py                       |  9 +++++++--
 tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py | 12 ++++++++----
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py
index 2f6f905fb51..71e2ad43c67 100644
--- a/tests/e2e/mcp/mcp_client.py
+++ b/tests/e2e/mcp/mcp_client.py
@@ -278,8 +278,13 @@ class McpClient:
         return self.await_tool_entry(key, server_id, needle).name
 
     def await_tool_entry(self, key: str, server_id: str, needle: str) -> McpToolEntry:
+        tool = self.await_tool_catalog(key, server_id, needle).tool_containing(server_id, needle)
+        assert tool is not None
+        return tool
+
+    def await_tool_catalog(self, key: str, server_id: str, needle: str) -> McpToolsListResponse:
         """Poll tools/list until `server_id` serves a tool matching `needle`, and
-        return the successful tool snapshot. Fails at poll_timeout.
+        return that catalog snapshot. Fails at poll_timeout.
 
         /v1/mcp/server returns as soon as the DB row is written, but the gateway
         runs the initialize + tools/list handshake against the upstream lazily on
@@ -293,7 +298,7 @@ class McpClient:
             if isinstance(result, Success):
                 tool = result.data.tool_containing(server_id, needle)
                 if tool is not None:
-                    return tool
+                    return result.data
             if time.monotonic() >= deadline:
                 raise AssertionError(
                     f"server {server_id} never served a tool matching {needle!r} within "
diff --git a/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py
index 1c1f725fff4..1deafdf0ecc 100644
--- a/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py
+++ b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py
@@ -75,8 +75,10 @@ class TestMcpToolsetEnforcement:
         client.await_registered(server_id)
 
         catalog_key: Final = _key(client, resources, "catalog", server_id=server_id)
-        known_wire: Final = client.await_tool(catalog_key, server_id, SEARCH_LOGS_TOOL)
-        catalog: Final = unwrap(client.list_tools(catalog_key)).tool_names_for_server(server_id)
+        snapshot: Final = client.await_tool_catalog(catalog_key, server_id, SEARCH_LOGS_TOOL)
+        known_wire: Final = snapshot.tool_name_containing(server_id, SEARCH_LOGS_TOOL)
+        assert known_wire is not None
+        catalog: Final = snapshot.tool_names_for_server(server_id)
         assert len(catalog) > 2, (
             f"the Datadog core toolset must serve more tools than the toolset names, or the "
             f"restriction has nothing to hide; got {sorted(catalog)}"
@@ -149,8 +151,10 @@ def _assert_principal_toolset(client: McpClient, resources: ResourceManager, pri
 
     for transport in client.proxy.replicas_for("/mcp-rest/tools/list").values():
         replica = McpClient(proxy=replace(client.proxy, transport=transport))
-        granted = replica.await_tool_entry(control_key, server_id, SEARCH_LOGS_TOOL)
-        catalog = unwrap(replica.list_tools(control_key)).tool_names_for_server(server_id)
+        snapshot = replica.await_tool_catalog(control_key, server_id, SEARCH_LOGS_TOOL)
+        granted = snapshot.tool_containing(server_id, SEARCH_LOGS_TOOL)
+        assert granted is not None
+        catalog = snapshot.tool_names_for_server(server_id)
         outside = sorted(catalog - {granted.name})
         assert outside, f"uncapped upstream must expose a tool outside the grant: {catalog}"
         expected = frozenset({granted.name})