From 8264211f9823946e01b23bf888c5741d8df9fe43 Mon Sep 17 00:00:00 2001 From: Kathy Wu Date: Wed, 14 Jan 2026 16:36:39 -0800 Subject: [PATCH] chore: Consolidate test_mcp_toolset.py into one file There was an extra test_mcp_toolset in the tools/ directory with only one test; I moved it into the main file. Co-authored-by: Kathy Wu PiperOrigin-RevId: 856419611 --- .../tools/mcp_tool/test_mcp_toolset.py | 51 +++++++++++++ tests/unittests/tools/test_mcp_toolset.py | 71 ------------------- 2 files changed, 51 insertions(+), 71 deletions(-) delete mode 100644 tests/unittests/tools/test_mcp_toolset.py diff --git a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py index 5809efe5..f6d002ed 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py @@ -17,6 +17,7 @@ from io import StringIO import sys import unittest from unittest.mock import AsyncMock +from unittest.mock import MagicMock from unittest.mock import Mock from unittest.mock import patch @@ -28,6 +29,7 @@ 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 import pytest @@ -302,3 +304,52 @@ class TestMCPToolset: # Check that the method has the retry decorator assert hasattr(toolset.get_tools, "__wrapped__") + + @pytest.mark.asyncio + async def test_mcp_toolset_with_prefix(self): + """Test that McpToolset correctly applies the tool_name_prefix.""" + # Mock the connection parameters + mock_connection_params = MagicMock() + mock_connection_params.timeout = None + + # Mock the MCPSessionManager and its create_session method + mock_session_manager = MagicMock() + mock_session = MagicMock() + + # Mock the list_tools response from the MCP server + mock_tool1 = MagicMock() + mock_tool1.name = "tool1" + mock_tool1.description = "tool 1 desc" + mock_tool2 = MagicMock() + mock_tool2.name = "tool2" + mock_tool2.description = "tool 2 desc" + list_tools_result = MagicMock() + list_tools_result.tools = [mock_tool1, mock_tool2] + mock_session.list_tools = AsyncMock(return_value=list_tools_result) + mock_session_manager.create_session = AsyncMock(return_value=mock_session) + + # Create an instance of McpToolset with a prefix + toolset = McpToolset( + connection_params=mock_connection_params, + tool_name_prefix="my_prefix", + ) + + # Replace the internal session manager with our mock + toolset._mcp_session_manager = mock_session_manager + + # Get the tools from the toolset + tools = await toolset.get_tools() + + # The get_tools method in McpToolset returns MCPTool objects, which are + # instances of BaseTool. The prefixing is handled by the BaseToolset, + # so we need to call get_tools_with_prefix to get the prefixed tools. + prefixed_tools = await toolset.get_tools_with_prefix() + + # Assert that the tools are prefixed correctly + assert len(prefixed_tools) == 2 + assert prefixed_tools[0].name == "my_prefix_tool1" + assert prefixed_tools[1].name == "my_prefix_tool2" + + # Assert that the original tools are not modified + assert tools[0].name == "tool1" + assert tools[1].name == "tool2" diff --git a/tests/unittests/tools/test_mcp_toolset.py b/tests/unittests/tools/test_mcp_toolset.py deleted file mode 100644 index 7bfd9126..00000000 --- a/tests/unittests/tools/test_mcp_toolset.py +++ /dev/null @@ -1,71 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Unit tests for McpToolset.""" - -from unittest.mock import AsyncMock -from unittest.mock import MagicMock - -from google.adk.tools.mcp_tool.mcp_toolset import McpToolset -import pytest - - -@pytest.mark.asyncio -async def test_mcp_toolset_with_prefix(): - """Test that McpToolset correctly applies the tool_name_prefix.""" - # Mock the connection parameters - mock_connection_params = MagicMock() - mock_connection_params.timeout = None - - # Mock the MCPSessionManager and its create_session method - mock_session_manager = MagicMock() - mock_session = MagicMock() - - # Mock the list_tools response from the MCP server - mock_tool1 = MagicMock() - mock_tool1.name = "tool1" - mock_tool1.description = "tool 1 desc" - mock_tool2 = MagicMock() - mock_tool2.name = "tool2" - mock_tool2.description = "tool 2 desc" - list_tools_result = MagicMock() - list_tools_result.tools = [mock_tool1, mock_tool2] - mock_session.list_tools = AsyncMock(return_value=list_tools_result) - mock_session_manager.create_session = AsyncMock(return_value=mock_session) - - # Create an instance of McpToolset with a prefix - toolset = McpToolset( - connection_params=mock_connection_params, - tool_name_prefix="my_prefix", - ) - - # Replace the internal session manager with our mock - toolset._mcp_session_manager = mock_session_manager - - # Get the tools from the toolset - tools = await toolset.get_tools() - - # The get_tools method in McpToolset returns MCPTool objects, which are - # instances of BaseTool. The prefixing is handled by the BaseToolset, - # so we need to call get_tools_with_prefix to get the prefixed tools. - prefixed_tools = await toolset.get_tools_with_prefix() - - # Assert that the tools are prefixed correctly - assert len(prefixed_tools) == 2 - assert prefixed_tools[0].name == "my_prefix_tool1" - assert prefixed_tools[1].name == "my_prefix_tool2" - - # Assert that the original tools are not modified - assert tools[0].name == "tool1" - assert tools[1].name == "tool2"