add more tests (#69)

vbarda · web-flow · commit 072bf41a773a · 2025-04-03T21:33:22.000-04:00
diff --git a/tests/servers/__init__.py b/tests/servers/__init__.py
diff --git a/tests/servers/math_server.py b/tests/servers/math_server.py
@@ -0,0 +1,29 @@
+from mcp.server.fastmcp import FastMCP
+
+mcp = FastMCP("Math")
+
+
+@mcp.tool()
+def add(a: int, b: int) -> int:
+    """Add two numbers"""
+    return a + b
+
+
+@mcp.tool()
+def multiply(a: int, b: int) -> int:
+    """Multiply two numbers"""
+    return a * b
+
+
+@mcp.prompt()
+def configure_assistant(skills: str) -> str:
+    return [
+        {
+            "role": "assistant",
+            "content": f"You are a helpful assistant. You have the following skills: {skills}. Always use only one tool at a time.",
+        }
+    ]
+
+
+if __name__ == "__main__":
+    mcp.run(transport="stdio")
diff --git a/tests/servers/weather_server.py b/tests/servers/weather_server.py
@@ -0,0 +1,13 @@
+from mcp.server.fastmcp import FastMCP
+
+mcp = FastMCP("Weather")
+
+
+@mcp.tool()
+async def get_weather(location: str) -> str:
+    """Get weather for location."""
+    return f"It's always sunny in {location}"
+
+
+if __name__ == "__main__":
+    mcp.run(transport="stdio")
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -0,0 +1,132 @@
+import os
+from pathlib import Path
+
+import pytest
+from langchain_core.messages import AIMessage
+from langchain_core.tools import BaseTool
+
+from langchain_mcp_adapters.client import MultiServerMCPClient
+
+
+@pytest.mark.asyncio
+async def test_multi_server_mcp_client():
+    """Test that the MultiServerMCPClient can connect to multiple servers and load tools."""
+
+    # Get the absolute path to the server scripts
+    current_dir = Path(__file__).parent
+    math_server_path = os.path.join(current_dir, "servers/math_server.py")
+    weather_server_path = os.path.join(current_dir, "servers/weather_server.py")
+
+    async with MultiServerMCPClient(
+        {
+            "math": {
+                "command": "python",
+                "args": [math_server_path],
+                "transport": "stdio",
+            },
+            "weather": {
+                "command": "python",
+                "args": [weather_server_path],
+                "transport": "stdio",
+            },
+        }
+    ) as client:
+        # Check that we have tools from both servers
+        all_tools = client.get_tools()
+
+        # Should have 3 tools (add, multiply, get_weather)
+        assert len(all_tools) == 3
+
+        # Check that tools are BaseTool instances
+        for tool in all_tools:
+            assert isinstance(tool, BaseTool)
+
+        # Verify tool names
+        tool_names = {tool.name for tool in all_tools}
+        assert tool_names == {"add", "multiply", "get_weather"}
+
+        # Check math server tools
+        math_tools = client.server_name_to_tools["math"]
+        assert len(math_tools) == 2
+        math_tool_names = {tool.name for tool in math_tools}
+        assert math_tool_names == {"add", "multiply"}
+
+        # Check weather server tools
+        weather_tools = client.server_name_to_tools["weather"]
+        assert len(weather_tools) == 1
+        assert weather_tools[0].name == "get_weather"
+
+        # Test that we can call a math tool
+        add_tool = next(tool for tool in all_tools if tool.name == "add")
+        result = await add_tool.ainvoke({"a": 2, "b": 3})
+        assert result == "5"
+
+        # Test that we can call a weather tool
+        weather_tool = next(tool for tool in all_tools if tool.name == "get_weather")
+        result = await weather_tool.ainvoke({"location": "London"})
+        assert result == "It's always sunny in London"
+
+        # Test the multiply tool
+        multiply_tool = next(tool for tool in all_tools if tool.name == "multiply")
+        result = await multiply_tool.ainvoke({"a": 4, "b": 5})
+        assert result == "20"
+
+
+@pytest.mark.asyncio
+async def test_multi_server_connect_methods():
+    """Test the different connect methods for MultiServerMCPClient."""
+
+    # Get the absolute path to the server scripts
+    current_dir = Path(__file__).parent
+    math_server_path = os.path.join(current_dir, "servers/math_server.py")
+    weather_server_path = os.path.join(current_dir, "servers/weather_server.py")
+
+    # Initialize client without initial connections
+    client = MultiServerMCPClient()
+    async with client:
+        # Connect to math server using connect_to_server
+        await client.connect_to_server(
+            "math", transport="stdio", command="python", args=[math_server_path]
+        )
+
+        # Connect to weather server using connect_to_server_via_stdio
+        await client.connect_to_server_via_stdio(
+            "weather", command="python", args=[weather_server_path]
+        )
+
+        # Check that we have tools from both servers
+        all_tools = client.get_tools()
+        assert len(all_tools) == 3
+
+        # Verify tool names
+        tool_names = {tool.name for tool in all_tools}
+        assert tool_names == {"add", "multiply", "get_weather"}
+
+
+@pytest.mark.asyncio
+async def test_get_prompt():
+    """Test retrieving prompts from MCP servers."""
+
+    # Get the absolute path to the server scripts
+    current_dir = Path(__file__).parent
+    math_server_path = os.path.join(current_dir, "servers/math_server.py")
+
+    async with MultiServerMCPClient(
+        {
+            "math": {
+                "command": "python",
+                "args": [math_server_path],
+                "transport": "stdio",
+            }
+        }
+    ) as client:
+        # Test getting a prompt from the math server
+        messages = await client.get_prompt(
+            "math", "configure_assistant", {"skills": "math, addition, multiplication"}
+        )
+
+        # Check that we got an AIMessage back
+        assert len(messages) == 1
+        assert isinstance(messages[0], AIMessage)
+        assert "You are a helpful assistant" in messages[0].content
+        assert "math, addition, multiplication" in messages[0].content
diff --git a/tests/test_tools.py b/tests/test_tools.py
@@ -0,0 +1,196 @@
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import BaseTool, ToolException
+from mcp.types import (
+    CallToolResult,
+    EmbeddedResource,
+    ImageContent,
+    TextContent,
+    TextResourceContents,
+)
+from mcp.types import Tool as MCPTool
+
+from langchain_mcp_adapters.tools import (
+    _convert_call_tool_result,
+    convert_mcp_tool_to_langchain_tool,
+    load_mcp_tools,
+)
+
+
+def test_convert_single_text_content():
+    # Test with a single text content
+    result = CallToolResult(
+        content=[TextContent(type="text", text="test result")],
+        isError=False,
+    )
+
+    text_content, non_text_content = _convert_call_tool_result(result)
+
+    assert text_content == "test result"
+    assert non_text_content is None
+
+
+def test_convert_multiple_text_contents():
+    # Test with multiple text contents
+    result = CallToolResult(
+        content=[
+            TextContent(type="text", text="result 1"),
+            TextContent(type="text", text="result 2"),
+        ],
+        isError=False,
+    )
+
+    text_content, non_text_content = _convert_call_tool_result(result)
+
+    assert text_content == ["result 1", "result 2"]
+    assert non_text_content is None
+
+
+def test_convert_with_non_text_content():
+    # Test with non-text content
+    image_content = ImageContent(type="image", mimeType="image/png", data="base64data")
+    resource_content = EmbeddedResource(
+        type="resource",
+        resource=TextResourceContents(uri="resource://test", mimeType="text/plain", text="hi"),
+    )
+
+    result = CallToolResult(
+        content=[
+            TextContent(type="text", text="text result"),
+            image_content,
+            resource_content,
+        ],
+        isError=False,
+    )
+
+    text_content, non_text_content = _convert_call_tool_result(result)
+
+    assert text_content == "text result"
+    assert non_text_content == [image_content, resource_content]
+
+
+def test_convert_with_error():
+    # Test with error
+    result = CallToolResult(
+        content=[TextContent(type="text", text="error message")],
+        isError=True,
+    )
+
+    with pytest.raises(ToolException) as exc_info:
+        _convert_call_tool_result(result)
+
+    assert str(exc_info.value) == "error message"
+
+
+@pytest.mark.asyncio
+async def test_convert_mcp_tool_to_langchain_tool():
+    tool_input_schema = {
+        "properties": {
+            "param1": {"title": "Param1", "type": "string"},
+            "param2": {"title": "Param2", "type": "integer"},
+        },
+        "required": ["param1", "param2"],
+        "title": "ToolSchema",
+        "type": "object",
+    }
+    # Mock session and MCP tool
+    session = AsyncMock()
+    session.call_tool.return_value = CallToolResult(
+        content=[TextContent(type="text", text="tool result")],
+        isError=False,
+    )
+
+    mcp_tool = MCPTool(
+        name="test_tool",
+        description="Test tool description",
+        inputSchema=tool_input_schema,
+    )
+
+    # Convert MCP tool to LangChain tool
+    lc_tool = convert_mcp_tool_to_langchain_tool(session, mcp_tool)
+
+    # Verify the converted tool
+    assert lc_tool.name == "test_tool"
+    assert lc_tool.description == "Test tool description"
+    assert lc_tool.args_schema == tool_input_schema
+
+    # Test calling the tool
+    result = await lc_tool.ainvoke(
+        {"args": {"param1": "test", "param2": 42}, "id": "1", "type": "tool_call"}
+    )
+
+    # Verify session.call_tool was called with correct arguments
+    session.call_tool.assert_called_once_with("test_tool", {"param1": "test", "param2": 42})
+
+    # Verify result
+    assert result == ToolMessage(content="tool result", name="test_tool", tool_call_id="1")
+
+
+@pytest.mark.asyncio
+async def test_load_mcp_tools():
+    tool_input_schema = {
+        "properties": {
+            "param1": {"title": "Param1", "type": "string"},
+            "param2": {"title": "Param2", "type": "integer"},
+        },
+        "required": ["param1", "param2"],
+        "title": "ToolSchema",
+        "type": "object",
+    }
+    # Mock session and list_tools response
+    session = AsyncMock()
+    mcp_tools = [
+        MCPTool(
+            name="tool1",
+            description="Tool 1 description",
+            inputSchema=tool_input_schema,
+        ),
+        MCPTool(
+            name="tool2",
+            description="Tool 2 description",
+            inputSchema=tool_input_schema,
+        ),
+    ]
+    session.list_tools.return_value = MagicMock(tools=mcp_tools)
+
+    # Mock call_tool to return different results for different tools
+    async def mock_call_tool(tool_name, arguments):
+        if tool_name == "tool1":
+            return CallToolResult(
+                content=[TextContent(type="text", text=f"tool1 result with {arguments}")],
+                isError=False,
+            )
+        else:
+            return CallToolResult(
+                content=[TextContent(type="text", text=f"tool2 result with {arguments}")],
+                isError=False,
+            )
+
+    session.call_tool.side_effect = mock_call_tool
+
+    # Load MCP tools
+    tools = await load_mcp_tools(session)
+
+    # Verify the tools
+    assert len(tools) == 2
+    assert all(isinstance(tool, BaseTool) for tool in tools)
+    assert tools[0].name == "tool1"
+    assert tools[1].name == "tool2"
+
+    # Test calling the first tool
+    result1 = await tools[0].ainvoke(
+        {"args": {"param1": "test1", "param2": 1}, "id": "1", "type": "tool_call"}
+    )
+    assert result1 == ToolMessage(
+        content="tool1 result with {'param1': 'test1', 'param2': 1}", name="tool1", tool_call_id="1"
+    )
+
+    # Test calling the second tool
+    result2 = await tools[1].ainvoke(
+        {"args": {"param1": "test2", "param2": 2}, "id": "2", "type": "tool_call"}
+    )
+    assert result2 == ToolMessage(
+        content="tool2 result with {'param1': 'test2', 'param2': 2}", name="tool2", tool_call_id="2"
+    )
