mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
chore: Convert MCPToolset to McpToolset in unittests
MCPToolset has been deprecated, use McpToolset instead. Co-authored-by: Kathy Wu <wukathy@google.com> PiperOrigin-RevId: 862457868
This commit is contained in:
Kathy Wu
committed by
Copybara-Service
Copybara-Service
parent
32f9f92042
commit
4341839420
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user