mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
George WealeandCopybara-Service
8c9105bf14
Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 839799108
72 lines
2.5 KiB
Python
72 lines
2.5 KiB
Python
# 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"
|