From 43418394202c2d01b0d37f6424bd601148077e27 Mon Sep 17 00:00:00 2001 From: Kathy Wu Date: Wed, 28 Jan 2026 17:03:44 -0800 Subject: [PATCH] chore: Convert MCPToolset to McpToolset in unittests MCPToolset has been deprecated, use McpToolset instead. Co-authored-by: Kathy Wu PiperOrigin-RevId: 862457868 --- .../tools/mcp_tool/test_mcp_toolset.py | 43 +++++++++---------- 1 file changed, 21 insertions(+), 22 deletions(-) diff --git a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py index d2fece0e..342f2e77 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py @@ -30,7 +30,6 @@ from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams from google.adk.tools.mcp_tool.mcp_tool import MCPTool -from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset from google.adk.tools.mcp_tool.mcp_toolset import McpToolset from mcp import StdioServerParameters from mcp.types import ListResourcesResult @@ -57,8 +56,8 @@ class MockListToolsResult: self.tools = tools -class TestMCPToolset: - """Test suite for MCPToolset class.""" +class TestMcpToolset: + """Test suite for McpToolset class.""" def setup_method(self): """Set up test fixtures.""" @@ -73,7 +72,7 @@ class TestMCPToolset: def test_init_basic(self): """Test basic initialization with StdioServerParameters.""" - toolset = MCPToolset(connection_params=self.mock_stdio_params) + toolset = McpToolset(connection_params=self.mock_stdio_params) # Note: StdioServerParameters gets converted to StdioConnectionParams internally assert toolset._errlog == sys.stderr @@ -85,7 +84,7 @@ class TestMCPToolset: stdio_params = StdioConnectionParams( server_params=self.mock_stdio_params, timeout=10.0 ) - toolset = MCPToolset(connection_params=stdio_params) + toolset = McpToolset(connection_params=stdio_params) assert toolset._connection_params == stdio_params @@ -94,7 +93,7 @@ class TestMCPToolset: sse_params = SseConnectionParams( url="https://example.com/mcp", headers={"Authorization": "Bearer token"} ) - toolset = MCPToolset(connection_params=sse_params) + toolset = McpToolset(connection_params=sse_params) assert toolset._connection_params == sse_params @@ -104,14 +103,14 @@ class TestMCPToolset: url="https://example.com/mcp", headers={"Content-Type": "application/json"}, ) - toolset = MCPToolset(connection_params=http_params) + toolset = McpToolset(connection_params=http_params) assert toolset._connection_params == http_params def test_init_with_tool_filter_list(self): """Test initialization with tool filter as list.""" tool_filter = ["tool1", "tool2"] - toolset = MCPToolset( + toolset = McpToolset( connection_params=self.mock_stdio_params, tool_filter=tool_filter ) @@ -132,7 +131,7 @@ class TestMCPToolset: oauth2=OAuth2Auth(client_id="test_id", client_secret="test_secret"), ) - toolset = MCPToolset( + toolset = McpToolset( connection_params=self.mock_stdio_params, auth_scheme=auth_scheme, auth_credential=auth_credential, @@ -144,7 +143,7 @@ class TestMCPToolset: def test_init_missing_connection_params(self): """Test initialization with missing connection params raises error.""" with pytest.raises(ValueError, match="Missing connection params"): - MCPToolset(connection_params=None) + McpToolset(connection_params=None) @pytest.mark.asyncio async def test_get_tools_basic(self): @@ -159,7 +158,7 @@ class TestMCPToolset: return_value=MockListToolsResult(mock_tools) ) - toolset = MCPToolset(connection_params=self.mock_stdio_params) + toolset = McpToolset(connection_params=self.mock_stdio_params) toolset._mcp_session_manager = self.mock_session_manager tools = await toolset.get_tools() @@ -185,7 +184,7 @@ class TestMCPToolset: ) tool_filter = ["tool1", "tool3"] - toolset = MCPToolset( + toolset = McpToolset( connection_params=self.mock_stdio_params, tool_filter=tool_filter ) toolset._mcp_session_manager = self.mock_session_manager @@ -213,7 +212,7 @@ class TestMCPToolset: """Filter for file-related tools only.""" return "file" in tool.name - toolset = MCPToolset( + toolset = McpToolset( connection_params=self.mock_stdio_params, tool_filter=file_tools_filter ) toolset._mcp_session_manager = self.mock_session_manager @@ -235,7 +234,7 @@ class TestMCPToolset: expected_headers = {"X-Tenant-ID": "test-tenant"} header_provider = Mock(return_value=expected_headers) - toolset = MCPToolset( + toolset = McpToolset( connection_params=self.mock_stdio_params, header_provider=header_provider, ) @@ -252,7 +251,7 @@ class TestMCPToolset: @pytest.mark.asyncio async def test_close_success(self): """Test successful cleanup.""" - toolset = MCPToolset(connection_params=self.mock_stdio_params) + toolset = McpToolset(connection_params=self.mock_stdio_params) toolset._mcp_session_manager = self.mock_session_manager await toolset.close() @@ -262,7 +261,7 @@ class TestMCPToolset: @pytest.mark.asyncio async def test_close_with_exception(self): """Test cleanup when session manager raises exception.""" - toolset = MCPToolset(connection_params=self.mock_stdio_params) + toolset = McpToolset(connection_params=self.mock_stdio_params) toolset._mcp_session_manager = self.mock_session_manager # Mock close to raise an exception @@ -287,7 +286,7 @@ class TestMCPToolset: stdio_params = StdioConnectionParams( server_params=self.mock_stdio_params, timeout=0.01 ) - toolset = MCPToolset(connection_params=stdio_params) + toolset = McpToolset(connection_params=stdio_params) toolset._mcp_session_manager = self.mock_session_manager async def long_running_list_tools(): @@ -304,7 +303,7 @@ class TestMCPToolset: @pytest.mark.asyncio async def test_get_tools_retry_decorator(self): """Test that get_tools has retry decorator applied.""" - toolset = MCPToolset(connection_params=self.mock_stdio_params) + toolset = McpToolset(connection_params=self.mock_stdio_params) # Check that the method has the retry decorator assert hasattr(toolset.get_tools, "__wrapped__") @@ -376,7 +375,7 @@ class TestMCPToolset: return_value=list_resources_result ) - toolset = MCPToolset(connection_params=self.mock_stdio_params) + toolset = McpToolset(connection_params=self.mock_stdio_params) toolset._mcp_session_manager = self.mock_session_manager result = await toolset.list_resources() @@ -402,7 +401,7 @@ class TestMCPToolset: return_value=list_resources_result ) - toolset = MCPToolset(connection_params=self.mock_stdio_params) + toolset = McpToolset(connection_params=self.mock_stdio_params) toolset._mcp_session_manager = self.mock_session_manager result = await toolset.get_resource_info("data.json") @@ -427,7 +426,7 @@ class TestMCPToolset: return_value=list_resources_result ) - toolset = MCPToolset(connection_params=self.mock_stdio_params) + toolset = McpToolset(connection_params=self.mock_stdio_params) toolset._mcp_session_manager = self.mock_session_manager with pytest.raises( @@ -482,7 +481,7 @@ class TestMCPToolset: get_resource_result.encoding = encoding self.mock_session.get_resource = AsyncMock(return_value=get_resource_result) - toolset = MCPToolset(connection_params=self.mock_stdio_params) + toolset = McpToolset(connection_params=self.mock_stdio_params) toolset._mcp_session_manager = self.mock_session_manager result = await toolset.read_resource(name)