google · pkduongsu · May 10, 2025 · May 10, 2025 · May 10, 2025 · May 11, 2025
diff --git a/.github/workflows/python-unit-tests.yml b/.github/workflows/python-unit-tests.yml
@@ -51,3 +51,11 @@ jobs:
         pytest tests/unittests \
           --ignore=tests/unittests/artifacts/test_artifact_service.py \
           --ignore=tests/unittests/tools/google_api_tool/test_googleapi_to_openapi_converter.py
+          --ignore=tests/unittests/tools/mcp_tool
+
+    - name: Run MCP unit tests without python 3.9
+      if: matrix.python-version != '3.9'
+      run: |
+        source .venv/bin/activate
+        pytest tests/unittests/tools/mcp_tool
+
diff --git a/tests/unittests/tools/mcp_tool/test_conversion_utils.py b/tests/unittests/tools/mcp_tool/test_conversion_utils.py
@@ -0,0 +1,106 @@
+from typing import Any, Dict
+from unittest.mock import MagicMock
+
+import pytest
+from google.genai.types import Schema, Type
+
+from src.google.adk.tools.mcp_tool.conversion_utils import (
+    adk_to_mcp_tool_type,
+    gemini_to_json_schema,
+)
+from src.google.adk.tools.base_tool import BaseTool
+import mcp.types as mcp_types
+
+
+def test_adk_to_mcp_tool_type():
+    """Test adk_to_mcp_tool_type with a mock BaseTool."""
+    mock_tool = MagicMock(spec=BaseTool)
+    mock_tool.name = "test_tool"
+    mock_tool.description = "test_description"
+    mock_tool._get_declaration.return_value = None
+
+    mcp_tool = adk_to_mcp_tool_type(mock_tool)
+
+    assert isinstance(mcp_tool, mcp_types.Tool)
+    assert mcp_tool.name == "test_tool"
+    assert mcp_tool.description == "test_description"
+    assert mcp_tool.inputSchema == {}
+
+    mock_declaration = MagicMock()
+    mock_declaration.parameters = Schema(type=Type.STRING)
+    mock_tool._get_declaration.return_value = mock_declaration
+    mcp_tool = adk_to_mcp_tool_type(mock_tool)
+    assert mcp_tool.inputSchema == {"type": "string"}
+
+
+def test_gemini_to_json_schema_string():
+    """Test gemini_to_json_schema with a STRING schema."""
+    gemini_schema = Schema(type=Type.STRING, title="test_string", description="test string description", default="test", enum=["a", "b"], format="test_format", example="example", pattern="test_pattern", min_length=1, max_length=10)
+    json_schema = gemini_to_json_schema(gemini_schema)
+    assert json_schema == {
+        "type": "string",
+        "title": "test_string",
+        "description": "test string description",
+        "default": "test",
+        "enum": ["a", "b"],
+        "format": "test_format",
+        "example": "example",
+        "pattern": "test_pattern",
+        "minLength": 1,
+        "maxLength": 10,
+    }
+
+
+def test_gemini_to_json_schema_number():
+    """Test gemini_to_json_schema with a NUMBER schema."""
+    gemini_schema = Schema(type=Type.NUMBER, minimum=1.0, maximum=10.0)
+    json_schema = gemini_to_json_schema(gemini_schema)
+    assert json_schema == {"type": "number", "minimum": 1.0, "maximum": 10.0}
+
+
+def test_gemini_to_json_schema_integer():
+    """Test gemini_to_json_schema with an INTEGER schema."""
+    gemini_schema = Schema(type=Type.INTEGER, minimum=1, maximum=10)
+    json_schema = gemini_to_json_schema(gemini_schema)
+    assert json_schema == {"type": "integer", "minimum": 1, "maximum": 10}
+
+
+def test_gemini_to_json_schema_array():
+    """Test gemini_to_json_schema with an ARRAY schema."""
+    gemini_schema = Schema(type=Type.ARRAY, items=Schema(type=Type.STRING), min_items=1, max_items=10)
+    json_schema = gemini_to_json_schema(gemini_schema)
+    assert json_schema == {"type": "array", "items": {"type": "string"}, "minItems": 1, "maxItems": 10}
+
+
+def test_gemini_to_json_schema_object():
+    """Test gemini_to_json_schema with an OBJECT schema."""
+    gemini_schema = Schema(type=Type.OBJECT, properties={"prop1": Schema(type=Type.STRING)}, required=["prop1"], min_properties=1, max_properties=10)
+    json_schema = gemini_to_json_schema(gemini_schema)
+    assert json_schema == {"type": "object", "properties": {"prop1": {"type": "string"}}, "required": ["prop1"], "minProperties": 1, "maxProperties": 10}
+
+
+def test_gemini_to_json_schema_nullable():
+    """Test gemini_to_json_schema with a nullable schema."""
+    gemini_schema = Schema(type=Type.STRING, nullable=True)
+    json_schema = gemini_to_json_schema(gemini_schema)
+    assert json_schema == {"type": "string", "nullable": True}
+
+
+def test_gemini_to_json_schema_any_of():
+    """Test gemini_to_json_schema with an anyOf schema."""
+    gemini_schema = Schema(type=Type.OBJECT, any_of=[Schema(type=Type.STRING), Schema(type=Type.INTEGER)])
+    json_schema = gemini_to_json_schema(gemini_schema)
+    assert json_schema == {"type": "object", "anyOf": [{"type": "string"}, {"type": "integer"}]}
+
+
+def test_gemini_to_json_schema_type_error():
+    """Test gemini_to_json_schema raises TypeError when input is not a Schema."""
+    with pytest.raises(TypeError):
+        gemini_to_json_schema("not a schema") 
+
+
+def test_gemini_to_json_schema_unspecified_type():
+    """Test gemini_to_json_schema with an unspecified type."""
+    gemini_schema = Schema()
+    json_schema = gemini_to_json_schema(gemini_schema)
+    assert json_schema == {"type": "null"}
diff --git a/tests/unittests/tools/mcp_tool/test_mcp_tool.py b/tests/unittests/tools/mcp_tool/test_mcp_tool.py
@@ -0,0 +1,104 @@
+import pytest
+from unittest.mock import MagicMock, AsyncMock
+from mcp import ClientSession
+
+from google.adk.tools.mcp_tool import MCPTool
+from google.adk.tools.mcp_tool.mcp_session_manager import MCPSessionManager
+
+from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import to_gemini_schema
+
+#helpers
+@pytest.fixture(scope="function")
+def mock_mcp_tool():
+    mock_mcp_tool = MagicMock()
+    mock_mcp_tool.name = "test_tool"
+    mock_mcp_tool.description = "test_description"
+    mock_mcp_tool.inputSchema = {"test": "testing", "test2": "testing2"}
+
+    return mock_mcp_tool
+
+@pytest.fixture
+def mock_mcp_session_manager():
+    return MagicMock(spec=MCPSessionManager)
+
+@pytest.fixture
+def mock_mcp_session():
+    return MagicMock(spec=ClientSession)
+
+@pytest.mark.asyncio
+async def test_init_mcp_tool(mock_mcp_tool, mock_mcp_session_manager, mock_mcp_session):
+    """Test that the MCPTool is initialized correctly."""
+    mock_mcp_session.call_tool = AsyncMock(return_value="Tool executed successfully")
+    mcp_tool = MCPTool(mcp_tool=mock_mcp_tool, mcp_session=mock_mcp_session, mcp_session_manager=mock_mcp_session_manager)
+    assert mcp_tool.name == "test_tool"
+    assert mcp_tool.description == "test_description"
+
+
+@pytest.mark.asyncio
+async def test_call_mcp_tool_valid_arguments(mock_mcp_tool):
+    """Test calling the MCP tool with valid arguments."""
+    mock_mcp_session = MagicMock()
+    mock_mcp_session.call_tool = AsyncMock(return_value="Tool executed successfully")
+    mock_mcp_session_manager = MagicMock()
+
+    mcp_tool = MCPTool(mcp_tool=mock_mcp_tool, mcp_session=mock_mcp_session, mcp_session_manager=mock_mcp_session_manager)
+    result = await mcp_tool.run_async(args={"arg1": "value1"}, tool_context=MagicMock())
+    assert result == "Tool executed successfully"
+    mock_mcp_session.call_tool.assert_called_once_with("test_tool", arguments={"arg1": "value1"})
+
+
+def test_init_mcp_tool_with_none_mcp_tool():
+    mock_mcp_session = MagicMock()
+    mock_mcp_session_manager = MagicMock() 
+
+    with pytest.raises(ValueError, match="mcp_tool cannot be None"):
+        MCPTool(mcp_tool=None, mcp_session=mock_mcp_session, mcp_session_manager=mock_mcp_session_manager)
+
+def test_init_mcp_tool_with_none_mcp_session():
+    mock_mcp_tool = MagicMock()
+    mock_mcp_session_manager = MagicMock() 
+
+    with pytest.raises(ValueError, match="mcp_session cannot be None"):
+        MCPTool(mcp_tool=mock_mcp_tool, mcp_session=None, mcp_session_manager=mock_mcp_session_manager)
+
+
+
+@pytest.mark.asyncio
+async def test_call_mcp_tool_invalid_arguments(mock_mcp_tool):
+    """Test calling the MCP tool with invalid arguments."""
+    mock_mcp_tool.inputSchema = {"arg1": {"type": "string", "required": True}}
+    mock_mcp_session = MagicMock()
+    mock_mcp_session.call_tool = AsyncMock(side_effect=ValueError("Missing required argument"))
+    mock_mcp_session_manager = MagicMock()
+
+    mcp_tool = MCPTool(mcp_tool=mock_mcp_tool, mcp_session=mock_mcp_session, mcp_session_manager=mock_mcp_session_manager)
+    with pytest.raises(ValueError):
+        await mcp_tool.run_async(args={}, tool_context=MagicMock())
+
+
+@pytest.mark.asyncio
+async def test_reinitialize_session(mock_mcp_tool, mock_mcp_session, mock_mcp_session_manager):
+    """Test reinitialize session."""
+    mock_mcp_session_manager.create_session.return_value = mock_mcp_session
+
+    tool = MCPTool(mock_mcp_tool, mock_mcp_session, mock_mcp_session_manager)
+
+    await tool._reinitialize_session()
+
+    mock_mcp_session_manager.create_session.assert_called_once()
+
+    assert tool._mcp_session == mock_mcp_session
+
+
+def test_get_declaration(mock_mcp_tool, mock_mcp_session, mock_mcp_session_manager):
+    """Test getting the Gemini function declaration for the tool."""
+    tool = MCPTool(mock_mcp_tool, mock_mcp_session, mock_mcp_session_manager)
+
+    decl = tool._get_declaration()
+
+    assert mock_mcp_tool.name == decl.name
+    assert mock_mcp_tool.description == decl.description
+    assert to_gemini_schema(mock_mcp_tool.inputSchema) == decl.parameters
+
+
+
diff --git a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py
@@ -0,0 +1,127 @@
+import pytest
+from unittest.mock import MagicMock, AsyncMock
+
+from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset
+from google.adk.tools.mcp_tool import MCPTool
+from mcp import StdioServerParameters
+from src.google.adk.tools.mcp_tool.mcp_session_manager import SseServerParams
+from google.adk.tools.base_toolset import ToolPredicate
+
+@pytest.mark.asyncio
+async def test_init_mcp_toolset():
+    """Test that the MCPToolset is initialized correctly."""
+    mock_connection_params = MagicMock()
+    mcp_toolset = MCPToolset(connection_params=mock_connection_params)
+    assert mcp_toolset._connection_params == mock_connection_params
+
+def test_create_mcp_toolset_invalid_connection_params():
+    """Test creating the MCPToolset with invalid connection parameters."""
+    with pytest.raises(ValueError):
+        MCPToolset(connection_params=None)
+
+@pytest.mark.asyncio
+async def test_get_tools():
+    """Test getting tools from the MCPToolset."""
+    mock_connection_params = MagicMock(spec=StdioServerParameters)
+    mock_connection_params.command = "test_command"
+    mock_connection_params.args = []  # Add the missing 'args' attribute
+    mcp_toolset = MCPToolset(connection_params=mock_connection_params)
+    mock_session = AsyncMock()
+    mock_mcp_tool = MagicMock()
+    mock_mcp_tool.name = "test_tool"
+    mock_mcp_tool.description = "test_description"
+    mock_session.list_tools = AsyncMock(return_value=MagicMock(tools=[mock_mcp_tool]))
+    mcp_toolset._session = mock_session  # Assign the mock session
+    tools = await mcp_toolset.get_tools()
+    assert isinstance(tools, list)
+    assert len(tools) == 1
+    assert isinstance(tools[0], MCPTool)
+    assert tools[0].name == "test_tool"
+    assert tools[0].description == "test_description"
+
+@pytest.mark.asyncio
+async def test_close_mcp_toolset():
+    """Test closing connection to the MCP server."""
+    mock_connection_params = MagicMock()
+    mcp_toolset = MCPToolset(connection_params=mock_connection_params)
+    mcp_toolset._exit_stack = AsyncMock()
+    await mcp_toolset.close()
+
+@pytest.mark.asyncio
+async def test_get_tools_error():
+    """Test handling errors during tool listing."""
+    mock_connection_params = MagicMock(spec=StdioServerParameters)
+    mock_connection_params.command = "test_command"
+    mock_connection_params.args = []  
+    mcp_toolset = MCPToolset(connection_params=mock_connection_params)
+    mock_session = AsyncMock()
+    mock_session.list_tools = AsyncMock(side_effect=Exception("Failed to list tools"))
+    mcp_toolset._session = mock_session
+    mcp_toolset._exit_stack = AsyncMock() 
+    mcp_toolset._exit_stack.aclose = AsyncMock()
+    with pytest.raises(Exception, match="Failed to list tools"):  
+        await mcp_toolset.get_tools()
+    await mcp_toolset._exit_stack.aclose()
+
+@pytest.mark.asyncio
+async def test_initialize():
+    """Test the _initialize method."""
+    mock_connection_params = MagicMock()
+    mcp_toolset = MCPToolset(connection_params=mock_connection_params)
+    mock_session_manager = AsyncMock()
+    mock_session = MagicMock()
+    mock_session_manager.create_session = AsyncMock(return_value=mock_session)
+    mcp_toolset._session_manager = mock_session_manager
+    session = await mcp_toolset._initialize()
+    assert session == mock_session
+    mock_session_manager.create_session.assert_called_once()
+
+def test_is_selected_no_filter():
+    """Test _is_selected when tool_filter is None."""
+    mock_connection_params = MagicMock()
+    mcp_toolset = MCPToolset(connection_params=mock_connection_params)
+    mock_tool = MagicMock()
+    assert mcp_toolset._is_selected(mock_tool, None) is True
+
+def test_is_selected_tool_predicate():
+    """Test _is_selected when tool_filter is a ToolPredicate."""
+    mock_connection_params = MagicMock()
+    tool_filter = MagicMock(spec=ToolPredicate)
+    tool_filter.return_value = True
+    mcp_toolset = MCPToolset(connection_params=mock_connection_params, tool_filter=tool_filter)
+    mock_tool = MagicMock()
+    assert mcp_toolset._is_selected(mock_tool, None) is True
+    tool_filter.assert_called_once_with(mock_tool, None)
+
+def test_is_selected_tool_list():
+    """Test _is_selected when tool_filter is a list."""
+    mock_connection_params = MagicMock()
+    tool_filter = ["tool1", "tool2"]
+    mcp_toolset = MCPToolset(connection_params=mock_connection_params, tool_filter=tool_filter)
+    mock_tool = MagicMock()
+    mock_tool.name = "tool1"
+    assert mcp_toolset._is_selected(mock_tool, None) is True
+    mock_tool.name = "tool3"
+    assert mcp_toolset._is_selected(mock_tool, None) is False
+
+@pytest.mark.asyncio
+async def test_get_tools_initializes_session():
+    """Test that get_tools initializes the session if it's None."""
+    mock_connection_params = MagicMock(spec=StdioServerParameters)
+    mock_connection_params.command = "test_command"
+    mock_connection_params.args = []  # Add the missing 'args' attribute
+    mcp_toolset = MCPToolset(connection_params=mock_connection_params)
+    mock_session = AsyncMock()
+    mock_mcp_tool = MagicMock()
+    mock_mcp_tool.name = "test_tool"
+    mock_mcp_tool.description = "test_description"
+    mock_session.list_tools = AsyncMock(return_value=MagicMock(tools=[mock_mcp_tool]))
+    mcp_toolset._session_manager = MagicMock()
+    mcp_toolset._session_manager.create_session = AsyncMock(return_value=mock_session)
+
+    mcp_toolset._initialize = AsyncMock()  # Mock with AsyncMock
+    mcp_toolset._initialize.side_effect = lambda: setattr(mcp_toolset, '_session', mock_session) or mock_session # Set session as side effect
+
+    mcp_toolset._session = None
+    tools = await mcp_toolset.get_tools()
+    mcp_toolset._initialize.assert_called_once()