initial commit after one day coding agent session
This commit is contained in:
Hans Aschauer
commit
372af75b90
88 changed files with 22694 additions and 0 deletions
359
tests/mcp/test_bridge.py
Normal file
359
tests/mcp/test_bridge.py
Normal file
|
|
@ -0,0 +1,359 @@
|
|||
"""Tests for Tool Bridge Server."""
|
||||
|
||||
import pytest
|
||||
import socket
|
||||
import json
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
from mcp_forge.mcp.bridge import ToolBridgeServer
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_socket_path():
|
||||
"""Create temporary socket path."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
yield Path(tmpdir) / "test_bridge.sock"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_client_manager():
|
||||
"""Create mock MCP client manager."""
|
||||
manager = AsyncMock()
|
||||
manager.call_tool = AsyncMock(return_value={"result": "success"})
|
||||
return manager
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_audit_logger():
|
||||
"""Create mock audit logger."""
|
||||
logger = Mock()
|
||||
logger.log = Mock()
|
||||
return logger
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bridge_server_starts_and_stops(temp_socket_path, mock_client_manager, mock_audit_logger):
|
||||
"""Test that bridge server starts and stops cleanly."""
|
||||
bridge = ToolBridgeServer(
|
||||
socket_path=temp_socket_path,
|
||||
client_manager=mock_client_manager,
|
||||
audit_logger=mock_audit_logger
|
||||
)
|
||||
|
||||
# Start server
|
||||
bridge.start()
|
||||
|
||||
# Verify socket was created
|
||||
assert temp_socket_path.exists()
|
||||
|
||||
# Stop server
|
||||
bridge.stop()
|
||||
|
||||
# Verify socket was removed
|
||||
assert not temp_socket_path.exists()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_receive_and_forward_tool_call(temp_socket_path, mock_client_manager, mock_audit_logger):
|
||||
"""Test receiving tool call request and forwarding to client."""
|
||||
bridge = ToolBridgeServer(
|
||||
socket_path=temp_socket_path,
|
||||
client_manager=mock_client_manager,
|
||||
audit_logger=mock_audit_logger
|
||||
)
|
||||
|
||||
bridge.start()
|
||||
|
||||
try:
|
||||
# Connect as client and send tool call request
|
||||
client_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
client_sock.connect(str(temp_socket_path))
|
||||
|
||||
request = {
|
||||
"tool": "test_tool",
|
||||
"params": {"arg1": "value1", "arg2": 42}
|
||||
}
|
||||
|
||||
client_sock.sendall(json.dumps(request).encode('utf-8'))
|
||||
client_sock.shutdown(socket.SHUT_WR)
|
||||
|
||||
# Receive response
|
||||
response_data = b''
|
||||
while True:
|
||||
chunk = client_sock.recv(4096)
|
||||
if not chunk:
|
||||
break
|
||||
response_data += chunk
|
||||
|
||||
response = json.loads(response_data.decode('utf-8'))
|
||||
|
||||
# Verify response
|
||||
assert response["success"] is True
|
||||
assert response["result"] == {"result": "success"}
|
||||
|
||||
# Verify tool was called with correct arguments
|
||||
mock_client_manager.call_tool.assert_called_once_with(
|
||||
"test_tool",
|
||||
{"arg1": "value1", "arg2": 42}
|
||||
)
|
||||
|
||||
# Verify audit log was called
|
||||
mock_audit_logger.log.assert_called()
|
||||
|
||||
client_sock.close()
|
||||
finally:
|
||||
bridge.stop()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_tool_call_error(temp_socket_path, mock_client_manager, mock_audit_logger):
|
||||
"""Test handling tool call errors."""
|
||||
# Make client manager raise error
|
||||
mock_client_manager.call_tool = AsyncMock(side_effect=RuntimeError("Tool failed"))
|
||||
|
||||
bridge = ToolBridgeServer(
|
||||
socket_path=temp_socket_path,
|
||||
client_manager=mock_client_manager,
|
||||
audit_logger=mock_audit_logger
|
||||
)
|
||||
|
||||
bridge.start()
|
||||
|
||||
try:
|
||||
# Connect and send request
|
||||
client_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
client_sock.connect(str(temp_socket_path))
|
||||
|
||||
request = {"tool": "failing_tool", "params": {}}
|
||||
client_sock.sendall(json.dumps(request).encode('utf-8'))
|
||||
client_sock.shutdown(socket.SHUT_WR)
|
||||
|
||||
# Receive response
|
||||
response_data = b''
|
||||
while True:
|
||||
chunk = client_sock.recv(4096)
|
||||
if not chunk:
|
||||
break
|
||||
response_data += chunk
|
||||
|
||||
response = json.loads(response_data.decode('utf-8'))
|
||||
|
||||
# Verify error response
|
||||
assert response["success"] is False
|
||||
assert "Tool failed" in response["error"]
|
||||
|
||||
client_sock.close()
|
||||
finally:
|
||||
bridge.stop()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_invalid_json(temp_socket_path, mock_client_manager, mock_audit_logger):
|
||||
"""Test handling invalid JSON in request."""
|
||||
bridge = ToolBridgeServer(
|
||||
socket_path=temp_socket_path,
|
||||
client_manager=mock_client_manager,
|
||||
audit_logger=mock_audit_logger
|
||||
)
|
||||
|
||||
bridge.start()
|
||||
|
||||
try:
|
||||
# Connect and send invalid JSON
|
||||
client_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
client_sock.connect(str(temp_socket_path))
|
||||
|
||||
client_sock.sendall(b"not valid json")
|
||||
client_sock.shutdown(socket.SHUT_WR)
|
||||
|
||||
# Receive response
|
||||
response_data = b''
|
||||
while True:
|
||||
chunk = client_sock.recv(4096)
|
||||
if not chunk:
|
||||
break
|
||||
response_data += chunk
|
||||
|
||||
response = json.loads(response_data.decode('utf-8'))
|
||||
|
||||
# Verify error response
|
||||
assert response["success"] is False
|
||||
assert "Invalid JSON" in response["error"]
|
||||
|
||||
client_sock.close()
|
||||
finally:
|
||||
bridge.stop()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_missing_tool_field(temp_socket_path, mock_client_manager, mock_audit_logger):
|
||||
"""Test handling request missing 'tool' field."""
|
||||
bridge = ToolBridgeServer(
|
||||
socket_path=temp_socket_path,
|
||||
client_manager=mock_client_manager,
|
||||
audit_logger=mock_audit_logger
|
||||
)
|
||||
|
||||
bridge.start()
|
||||
|
||||
try:
|
||||
# Connect and send request without 'tool' field
|
||||
client_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
client_sock.connect(str(temp_socket_path))
|
||||
|
||||
request = {"params": {"arg1": "value1"}} # Missing 'tool'
|
||||
client_sock.sendall(json.dumps(request).encode('utf-8'))
|
||||
client_sock.shutdown(socket.SHUT_WR)
|
||||
|
||||
# Receive response
|
||||
response_data = b''
|
||||
while True:
|
||||
chunk = client_sock.recv(4096)
|
||||
if not chunk:
|
||||
break
|
||||
response_data += chunk
|
||||
|
||||
response = json.loads(response_data.decode('utf-8'))
|
||||
|
||||
# Verify error response
|
||||
assert response["success"] is False
|
||||
assert "tool" in response["error"].lower()
|
||||
|
||||
client_sock.close()
|
||||
finally:
|
||||
bridge.stop()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_requests(temp_socket_path, mock_client_manager, mock_audit_logger):
|
||||
"""Test handling multiple concurrent requests."""
|
||||
import threading
|
||||
|
||||
# Track call counts
|
||||
call_count = {"count": 0}
|
||||
|
||||
async def mock_call_tool(tool_name, arguments):
|
||||
call_count["count"] += 1
|
||||
return {"result": f"success_{call_count['count']}"}
|
||||
|
||||
mock_client_manager.call_tool = mock_call_tool
|
||||
|
||||
bridge = ToolBridgeServer(
|
||||
socket_path=temp_socket_path,
|
||||
client_manager=mock_client_manager,
|
||||
audit_logger=mock_audit_logger
|
||||
)
|
||||
|
||||
bridge.start()
|
||||
|
||||
try:
|
||||
results = []
|
||||
|
||||
def make_request(tool_name):
|
||||
client_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
client_sock.connect(str(temp_socket_path))
|
||||
|
||||
request = {"tool": tool_name, "params": {}}
|
||||
client_sock.sendall(json.dumps(request).encode('utf-8'))
|
||||
client_sock.shutdown(socket.SHUT_WR)
|
||||
|
||||
response_data = b''
|
||||
while True:
|
||||
chunk = client_sock.recv(4096)
|
||||
if not chunk:
|
||||
break
|
||||
response_data += chunk
|
||||
|
||||
response = json.loads(response_data.decode('utf-8'))
|
||||
results.append(response)
|
||||
client_sock.close()
|
||||
|
||||
# Make 3 concurrent requests
|
||||
threads = []
|
||||
for i in range(3):
|
||||
thread = threading.Thread(target=make_request, args=(f"tool_{i}",))
|
||||
threads.append(thread)
|
||||
thread.start()
|
||||
|
||||
# Wait for all threads
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
|
||||
# Verify all requests succeeded
|
||||
assert len(results) == 3
|
||||
for result in results:
|
||||
assert result["success"] is True
|
||||
|
||||
# Verify all were processed
|
||||
assert call_count["count"] == 3
|
||||
finally:
|
||||
bridge.stop()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audit_logging_tool_name_only(temp_socket_path, mock_client_manager, mock_audit_logger):
|
||||
"""Test that audit log only logs tool name, not parameters."""
|
||||
bridge = ToolBridgeServer(
|
||||
socket_path=temp_socket_path,
|
||||
client_manager=mock_client_manager,
|
||||
audit_logger=mock_audit_logger
|
||||
)
|
||||
|
||||
bridge.start()
|
||||
|
||||
try:
|
||||
# Send request with sensitive parameters
|
||||
client_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
client_sock.connect(str(temp_socket_path))
|
||||
|
||||
request = {
|
||||
"tool": "sensitive_tool",
|
||||
"params": {"password": "secret123", "token": "abc123"}
|
||||
}
|
||||
|
||||
client_sock.sendall(json.dumps(request).encode('utf-8'))
|
||||
client_sock.shutdown(socket.SHUT_WR)
|
||||
|
||||
# Receive response
|
||||
response_data = b''
|
||||
while True:
|
||||
chunk = client_sock.recv(4096)
|
||||
if not chunk:
|
||||
break
|
||||
response_data += chunk
|
||||
|
||||
client_sock.close()
|
||||
|
||||
# Verify audit log was called
|
||||
mock_audit_logger.log.assert_called()
|
||||
|
||||
# Get the log call arguments
|
||||
log_call = mock_audit_logger.log.call_args
|
||||
|
||||
# Verify tool name is in log
|
||||
log_str = str(log_call)
|
||||
assert "sensitive_tool" in log_str
|
||||
|
||||
# Verify sensitive parameters are NOT in log
|
||||
assert "secret123" not in log_str
|
||||
assert "abc123" not in log_str
|
||||
finally:
|
||||
bridge.stop()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_socket_cleanup_on_error(temp_socket_path, mock_client_manager, mock_audit_logger):
|
||||
"""Test that socket is cleaned up even if server encounters error."""
|
||||
bridge = ToolBridgeServer(
|
||||
socket_path=temp_socket_path,
|
||||
client_manager=mock_client_manager,
|
||||
audit_logger=mock_audit_logger
|
||||
)
|
||||
|
||||
bridge.start()
|
||||
assert temp_socket_path.exists()
|
||||
|
||||
# Stop should cleanup
|
||||
bridge.stop()
|
||||
assert not temp_socket_path.exists()
|
||||
249
tests/mcp/test_client.py
Normal file
249
tests/mcp/test_client.py
Normal file
|
|
@ -0,0 +1,249 @@
|
|||
"""Tests for MCP Client Wrapper."""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import Mock, AsyncMock, patch
|
||||
from mcp_forge.mcp.client import MCPClientWrapper
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_fastmcp_client():
|
||||
"""Mock fastmcp Client."""
|
||||
client = AsyncMock()
|
||||
|
||||
# Create tool mocks with actual string names (not Mock.name)
|
||||
tool1 = Mock()
|
||||
tool1.name = "tool1"
|
||||
tool1.description = "Tool 1"
|
||||
tool1.inputSchema = {"type": "object", "properties": {}}
|
||||
|
||||
tool2 = Mock()
|
||||
tool2.name = "tool2"
|
||||
tool2.description = "Tool 2"
|
||||
tool2.inputSchema = {"type": "object", "properties": {}}
|
||||
|
||||
# Mock list_tools to return tool objects
|
||||
client.list_tools = AsyncMock(return_value=Mock(tools=[tool1, tool2]))
|
||||
|
||||
# Mock call_tool to return result with data
|
||||
client.call_tool = AsyncMock(return_value=Mock(
|
||||
data="result",
|
||||
content=[Mock(text="result")],
|
||||
is_error=False
|
||||
))
|
||||
|
||||
# Mock context manager
|
||||
client.__aenter__ = AsyncMock(return_value=client)
|
||||
client.__aexit__ = AsyncMock(return_value=None)
|
||||
|
||||
return client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connect_success(mock_fastmcp_client):
|
||||
"""Test successful connection to MCP server."""
|
||||
with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
|
||||
client = MCPClientWrapper(
|
||||
name="test-server",
|
||||
command="python",
|
||||
args=["-m", "test_server"]
|
||||
)
|
||||
|
||||
assert not client.is_connected()
|
||||
|
||||
await client.connect()
|
||||
|
||||
assert client.is_connected()
|
||||
mock_fastmcp_client.__aenter__.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connect_with_env(mock_fastmcp_client):
|
||||
"""Test connection with environment variables."""
|
||||
with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
|
||||
env = {"API_KEY": "test123", "DEBUG": "true"}
|
||||
client = MCPClientWrapper(
|
||||
name="test-server",
|
||||
command="python",
|
||||
args=["-m", "test_server"],
|
||||
env=env
|
||||
)
|
||||
|
||||
await client.connect()
|
||||
|
||||
assert client.is_connected()
|
||||
assert client.env == env
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disconnect(mock_fastmcp_client):
|
||||
"""Test disconnect from MCP server."""
|
||||
with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
|
||||
client = MCPClientWrapper(
|
||||
name="test-server",
|
||||
command="python",
|
||||
args=["-m", "test_server"]
|
||||
)
|
||||
|
||||
await client.connect()
|
||||
assert client.is_connected()
|
||||
|
||||
await client.disconnect()
|
||||
|
||||
assert not client.is_connected()
|
||||
mock_fastmcp_client.__aexit__.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_tools(mock_fastmcp_client):
|
||||
"""Test listing available tools."""
|
||||
with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
|
||||
client = MCPClientWrapper(
|
||||
name="test-server",
|
||||
command="python",
|
||||
args=["-m", "test_server"]
|
||||
)
|
||||
|
||||
await client.connect()
|
||||
tools = await client.list_tools()
|
||||
|
||||
assert tools == ["tool1", "tool2"]
|
||||
mock_fastmcp_client.list_tools.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_tools_not_connected():
|
||||
"""Test list_tools raises error when not connected."""
|
||||
client = MCPClientWrapper(
|
||||
name="test-server",
|
||||
command="python",
|
||||
args=["-m", "test_server"]
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="not connected"):
|
||||
await client.list_tools()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_tool_schema(mock_fastmcp_client):
|
||||
"""Test getting tool schema."""
|
||||
with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
|
||||
client = MCPClientWrapper(
|
||||
name="test-server",
|
||||
command="python",
|
||||
args=["-m", "test_server"]
|
||||
)
|
||||
|
||||
await client.connect()
|
||||
schema = await client.get_tool_schema("tool1")
|
||||
|
||||
assert schema == {"type": "object", "properties": {}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_tool_schema_not_found(mock_fastmcp_client):
|
||||
"""Test get_tool_schema raises error for unknown tool."""
|
||||
with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
|
||||
client = MCPClientWrapper(
|
||||
name="test-server",
|
||||
command="python",
|
||||
args=["-m", "test_server"]
|
||||
)
|
||||
|
||||
await client.connect()
|
||||
|
||||
with pytest.raises(KeyError, match="Tool 'unknown' not found"):
|
||||
await client.get_tool_schema("unknown")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_tool_success(mock_fastmcp_client):
|
||||
"""Test successful tool call."""
|
||||
with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
|
||||
client = MCPClientWrapper(
|
||||
name="test-server",
|
||||
command="python",
|
||||
args=["-m", "test_server"]
|
||||
)
|
||||
|
||||
await client.connect()
|
||||
result = await client.call_tool("tool1", {"param": "value"})
|
||||
|
||||
assert result == "result"
|
||||
mock_fastmcp_client.call_tool.assert_called_once_with("tool1", {"param": "value"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_tool_not_connected():
|
||||
"""Test call_tool raises error when not connected."""
|
||||
client = MCPClientWrapper(
|
||||
name="test-server",
|
||||
command="python",
|
||||
args=["-m", "test_server"]
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="not connected"):
|
||||
await client.call_tool("tool1", {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_tool_failure(mock_fastmcp_client):
|
||||
"""Test tool call failure handling."""
|
||||
mock_fastmcp_client.call_tool = AsyncMock(side_effect=Exception("Tool failed"))
|
||||
|
||||
with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
|
||||
client = MCPClientWrapper(
|
||||
name="test-server",
|
||||
command="python",
|
||||
args=["-m", "test_server"]
|
||||
)
|
||||
|
||||
await client.connect()
|
||||
|
||||
with pytest.raises(RuntimeError, match="Tool call failed.*Tool failed"):
|
||||
await client.call_tool("tool1", {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connection_failure():
|
||||
"""Test connection failure handling."""
|
||||
failing_client = AsyncMock()
|
||||
failing_client.__aenter__ = AsyncMock(side_effect=Exception("Connection failed"))
|
||||
|
||||
with patch('mcp_forge.mcp.client.Client', return_value=failing_client):
|
||||
client = MCPClientWrapper(
|
||||
name="test-server",
|
||||
command="python",
|
||||
args=["-m", "test_server"]
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="Failed to connect.*Connection failed"):
|
||||
await client.connect()
|
||||
|
||||
assert not client.is_connected()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reconnect(mock_fastmcp_client):
|
||||
"""Test reconnection after disconnect."""
|
||||
with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
|
||||
client = MCPClientWrapper(
|
||||
name="test-server",
|
||||
command="python",
|
||||
args=["-m", "test_server"]
|
||||
)
|
||||
|
||||
# First connection
|
||||
await client.connect()
|
||||
assert client.is_connected()
|
||||
|
||||
# Disconnect
|
||||
await client.disconnect()
|
||||
assert not client.is_connected()
|
||||
|
||||
# Reconnect
|
||||
await client.connect()
|
||||
assert client.is_connected()
|
||||
|
||||
# Should be able to use tools
|
||||
tools = await client.list_tools()
|
||||
assert tools == ["tool1", "tool2"]
|
||||
114
tests/mcp/test_http_transport.py
Normal file
114
tests/mcp/test_http_transport.py
Normal file
|
|
@ -0,0 +1,114 @@
|
|||
"""Tests for HTTP and SSE transport support in MCP client."""
|
||||
|
||||
import pytest
|
||||
from mcp_forge.mcp.client import MCPClientWrapper
|
||||
|
||||
|
||||
def test_stdio_transport_creation():
|
||||
"""Test creating a client with stdio transport."""
|
||||
client = MCPClientWrapper(
|
||||
name="test_stdio",
|
||||
transport_type="stdio",
|
||||
command="python",
|
||||
args=["-m", "server"],
|
||||
env={"KEY": "value"}
|
||||
)
|
||||
|
||||
assert client.name == "test_stdio"
|
||||
assert client.transport_type == "stdio"
|
||||
assert client.command == "python"
|
||||
assert client.args == ["-m", "server"]
|
||||
assert client.env == {"KEY": "value"}
|
||||
assert client.transport is not None
|
||||
|
||||
|
||||
def test_http_transport_creation():
|
||||
"""Test creating a client with HTTP transport."""
|
||||
client = MCPClientWrapper(
|
||||
name="test_http",
|
||||
transport_type="http",
|
||||
url="http://localhost:8006/mcp",
|
||||
headers={"Authorization": "Bearer token123"}
|
||||
)
|
||||
|
||||
assert client.name == "test_http"
|
||||
assert client.transport_type == "http"
|
||||
assert client.url == "http://localhost:8006/mcp"
|
||||
assert client.headers == {"Authorization": "Bearer token123"}
|
||||
assert client.transport is not None
|
||||
|
||||
|
||||
def test_sse_transport_creation():
|
||||
"""Test creating a client with SSE transport."""
|
||||
client = MCPClientWrapper(
|
||||
name="test_sse",
|
||||
transport_type="sse",
|
||||
url="http://localhost:9000/events",
|
||||
headers={"X-Custom": "value"}
|
||||
)
|
||||
|
||||
assert client.name == "test_sse"
|
||||
assert client.transport_type == "sse"
|
||||
assert client.url == "http://localhost:9000/events"
|
||||
assert client.headers == {"X-Custom": "value"}
|
||||
assert client.transport is not None
|
||||
|
||||
|
||||
def test_stdio_without_command_raises_error():
|
||||
"""Test that stdio transport requires a command."""
|
||||
with pytest.raises(ValueError, match="command required for stdio transport"):
|
||||
MCPClientWrapper(
|
||||
name="test_stdio",
|
||||
transport_type="stdio"
|
||||
)
|
||||
|
||||
|
||||
def test_http_without_url_raises_error():
|
||||
"""Test that HTTP transport requires a URL."""
|
||||
with pytest.raises(ValueError, match="url required for http transport"):
|
||||
MCPClientWrapper(
|
||||
name="test_http",
|
||||
transport_type="http"
|
||||
)
|
||||
|
||||
|
||||
def test_sse_without_url_raises_error():
|
||||
"""Test that SSE transport requires a URL."""
|
||||
with pytest.raises(ValueError, match="url required for sse transport"):
|
||||
MCPClientWrapper(
|
||||
name="test_sse",
|
||||
transport_type="sse"
|
||||
)
|
||||
|
||||
|
||||
def test_invalid_transport_type_raises_error():
|
||||
"""Test that an invalid transport type raises an error."""
|
||||
with pytest.raises(ValueError, match="Unknown transport type"):
|
||||
MCPClientWrapper(
|
||||
name="test_invalid",
|
||||
transport_type="invalid"
|
||||
)
|
||||
|
||||
|
||||
def test_http_transport_with_empty_headers():
|
||||
"""Test HTTP transport with empty headers dict."""
|
||||
client = MCPClientWrapper(
|
||||
name="test_http",
|
||||
transport_type="http",
|
||||
url="http://localhost:8006/mcp"
|
||||
)
|
||||
|
||||
assert client.headers == {}
|
||||
assert client.transport is not None
|
||||
|
||||
|
||||
def test_default_stdio_transport():
|
||||
"""Test that stdio is the default transport type."""
|
||||
client = MCPClientWrapper(
|
||||
name="test_default",
|
||||
command="python",
|
||||
args=["-m", "server"]
|
||||
)
|
||||
|
||||
assert client.transport_type == "stdio"
|
||||
assert client.command == "python"
|
||||
227
tests/mcp/test_injection.py
Normal file
227
tests/mcp/test_injection.py
Normal file
|
|
@ -0,0 +1,227 @@
|
|||
"""Tests for Tool Injection Generator."""
|
||||
|
||||
import pytest
|
||||
import ast
|
||||
from unittest.mock import AsyncMock
|
||||
from mcp_forge.mcp.injection import ToolInjectionGenerator
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_client_manager():
|
||||
"""Create mock MCP client manager with tools."""
|
||||
manager = AsyncMock()
|
||||
|
||||
# Tool schemas
|
||||
manager.get_tool_schema = AsyncMock(side_effect=lambda tool_name: {
|
||||
"read_file": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "File path to read"},
|
||||
"encoding": {"type": "string", "description": "File encoding"}
|
||||
},
|
||||
"required": ["path"]
|
||||
},
|
||||
"write_file": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "File path to write"},
|
||||
"content": {"type": "string", "description": "Content to write"},
|
||||
"mode": {"type": "string", "description": "Write mode"}
|
||||
},
|
||||
"required": ["path", "content"]
|
||||
},
|
||||
"calculate": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"expression": {"type": "string", "description": "Math expression"},
|
||||
"precision": {"type": "integer", "description": "Decimal precision"}
|
||||
},
|
||||
"required": ["expression"]
|
||||
}
|
||||
}[tool_name])
|
||||
|
||||
return manager
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_injection_code_is_valid_python(mock_client_manager):
|
||||
"""Test that generated code is valid Python."""
|
||||
generator = ToolInjectionGenerator(mock_client_manager)
|
||||
|
||||
code = await generator.generate_injection_code(
|
||||
tool_names=["read_file", "write_file"],
|
||||
bridge_socket_path="/tmp/bridge.sock"
|
||||
)
|
||||
|
||||
# Try to parse the generated code
|
||||
try:
|
||||
ast.parse(code)
|
||||
except SyntaxError as e:
|
||||
pytest.fail(f"Generated code has syntax error: {e}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generated_code_includes_bridge_client(mock_client_manager):
|
||||
"""Test that generated code includes bridge client."""
|
||||
generator = ToolInjectionGenerator(mock_client_manager)
|
||||
|
||||
code = await generator.generate_injection_code(
|
||||
tool_names=["read_file"],
|
||||
bridge_socket_path="/tmp/bridge.sock"
|
||||
)
|
||||
|
||||
# Verify bridge client function is included
|
||||
assert "_mcp_call" in code
|
||||
assert "socket.socket" in code
|
||||
assert "socket.AF_UNIX" in code
|
||||
assert "/tmp/bridge.sock" in code
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generated_code_includes_tool_functions(mock_client_manager):
|
||||
"""Test that generated code includes wrapper functions for each tool."""
|
||||
generator = ToolInjectionGenerator(mock_client_manager)
|
||||
|
||||
code = await generator.generate_injection_code(
|
||||
tool_names=["read_file", "write_file", "calculate"],
|
||||
bridge_socket_path="/tmp/bridge.sock"
|
||||
)
|
||||
|
||||
# Verify tool functions are defined
|
||||
assert "def read_file(" in code
|
||||
assert "def write_file(" in code
|
||||
assert "def calculate(" in code
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_function_signatures_match_schemas(mock_client_manager):
|
||||
"""Test that function signatures match tool schemas."""
|
||||
generator = ToolInjectionGenerator(mock_client_manager)
|
||||
|
||||
code = await generator.generate_injection_code(
|
||||
tool_names=["read_file", "write_file"],
|
||||
bridge_socket_path="/tmp/bridge.sock"
|
||||
)
|
||||
|
||||
# read_file has path (required) and encoding (optional)
|
||||
assert "def read_file(path: str, encoding: str = None)" in code
|
||||
|
||||
# write_file has path, content (required) and mode (optional)
|
||||
assert "def write_file(path: str, content: str, mode: str = None)" in code
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generated_functions_have_docstrings(mock_client_manager):
|
||||
"""Test that generated functions have docstrings."""
|
||||
generator = ToolInjectionGenerator(mock_client_manager)
|
||||
|
||||
code = await generator.generate_injection_code(
|
||||
tool_names=["read_file"],
|
||||
bridge_socket_path="/tmp/bridge.sock"
|
||||
)
|
||||
|
||||
# Verify docstring is present (contains parameter descriptions)
|
||||
assert '"""' in code
|
||||
assert "File path to read" in code or "path:" in code.lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generated_functions_call_bridge(mock_client_manager):
|
||||
"""Test that generated functions call _mcp_call."""
|
||||
generator = ToolInjectionGenerator(mock_client_manager)
|
||||
|
||||
code = await generator.generate_injection_code(
|
||||
tool_names=["read_file"],
|
||||
bridge_socket_path="/tmp/bridge.sock"
|
||||
)
|
||||
|
||||
# Verify function calls _mcp_call with tool name
|
||||
lines = code.split('\n')
|
||||
in_read_file = False
|
||||
found_call = False
|
||||
|
||||
for line in lines:
|
||||
if "def read_file(" in line:
|
||||
in_read_file = True
|
||||
if in_read_file and "_mcp_call" in line and "read_file" in line:
|
||||
found_call = True
|
||||
break
|
||||
|
||||
assert found_call, "Generated function should call _mcp_call with tool name"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_type_hints_from_schema(mock_client_manager):
|
||||
"""Test that type hints are generated from schema types."""
|
||||
generator = ToolInjectionGenerator(mock_client_manager)
|
||||
|
||||
code = await generator.generate_injection_code(
|
||||
tool_names=["calculate"],
|
||||
bridge_socket_path="/tmp/bridge.sock"
|
||||
)
|
||||
|
||||
# calculate has string expression and integer precision
|
||||
assert "expression: str" in code
|
||||
assert "precision: int" in code or "precision: " in code
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_required_vs_optional_parameters(mock_client_manager):
|
||||
"""Test that required and optional parameters are handled correctly."""
|
||||
generator = ToolInjectionGenerator(mock_client_manager)
|
||||
|
||||
code = await generator.generate_injection_code(
|
||||
tool_names=["read_file"],
|
||||
bridge_socket_path="/tmp/bridge.sock"
|
||||
)
|
||||
|
||||
# path is required (no default), encoding is optional (has default)
|
||||
assert "def read_file(path: str, encoding: str = None)" in code or \
|
||||
"def read_file(path: str, encoding: " in code
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generated_code_has_imports(mock_client_manager):
|
||||
"""Test that generated code includes necessary imports."""
|
||||
generator = ToolInjectionGenerator(mock_client_manager)
|
||||
|
||||
code = await generator.generate_injection_code(
|
||||
tool_names=["read_file"],
|
||||
bridge_socket_path="/tmp/bridge.sock"
|
||||
)
|
||||
|
||||
# Verify imports
|
||||
assert "import socket" in code
|
||||
assert "import json" in code
|
||||
assert "from typing import Any" in code or "typing" in code
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_tool_list(mock_client_manager):
|
||||
"""Test handling of empty tool list."""
|
||||
generator = ToolInjectionGenerator(mock_client_manager)
|
||||
|
||||
code = await generator.generate_injection_code(
|
||||
tool_names=[],
|
||||
bridge_socket_path="/tmp/bridge.sock"
|
||||
)
|
||||
|
||||
# Should still include bridge client
|
||||
assert "_mcp_call" in code
|
||||
# But no tool functions
|
||||
assert "def read_file(" not in code
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_socket_path(mock_client_manager):
|
||||
"""Test that custom socket path is used correctly."""
|
||||
generator = ToolInjectionGenerator(mock_client_manager)
|
||||
|
||||
custom_path = "/custom/path/to/socket.sock"
|
||||
code = await generator.generate_injection_code(
|
||||
tool_names=["read_file"],
|
||||
bridge_socket_path=custom_path
|
||||
)
|
||||
|
||||
# Verify custom path is in generated code
|
||||
assert custom_path in code
|
||||
245
tests/mcp/test_manager.py
Normal file
245
tests/mcp/test_manager.py
Normal file
|
|
@ -0,0 +1,245 @@
|
|||
"""Tests for MCP Client Manager."""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, patch
|
||||
from mcp_forge.mcp.manager import MCPClientManager
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_client1():
|
||||
"""Create mock MCP client 1."""
|
||||
client = AsyncMock()
|
||||
client.name = "client1"
|
||||
client.is_connected.return_value = False
|
||||
client.connect = AsyncMock()
|
||||
client.disconnect = AsyncMock()
|
||||
client.list_tools = AsyncMock(return_value=["tool1", "tool2"])
|
||||
client.get_tool_schema = AsyncMock(return_value={
|
||||
"type": "object",
|
||||
"properties": {"param1": {"type": "string"}}
|
||||
})
|
||||
client.call_tool = AsyncMock(return_value={"result": "success"})
|
||||
return client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_client2():
|
||||
"""Create mock MCP client 2."""
|
||||
client = AsyncMock()
|
||||
client.name = "client2"
|
||||
client.is_connected.return_value = False
|
||||
client.connect = AsyncMock()
|
||||
client.disconnect = AsyncMock()
|
||||
client.list_tools = AsyncMock(return_value=["tool3", "tool4"])
|
||||
client.get_tool_schema = AsyncMock(return_value={
|
||||
"type": "object",
|
||||
"properties": {"param2": {"type": "number"}}
|
||||
})
|
||||
client.call_tool = AsyncMock(return_value={"result": "success2"})
|
||||
return client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client_config():
|
||||
"""Create test client configuration."""
|
||||
return {
|
||||
"client1": {
|
||||
"command": "python",
|
||||
"args": ["server1.py"],
|
||||
"env": {"KEY1": "value1"}
|
||||
},
|
||||
"client2": {
|
||||
"command": "python",
|
||||
"args": ["server2.py"],
|
||||
"env": {"KEY2": "value2"}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_clients_from_config(mock_client1, mock_client2, client_config):
|
||||
"""Test initializing clients from configuration."""
|
||||
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
||||
# Setup mock to return different clients for different configs
|
||||
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
||||
|
||||
manager = MCPClientManager(client_config)
|
||||
await manager.initialize()
|
||||
|
||||
# Verify clients were created with correct config
|
||||
assert mock_wrapper.call_count == 2
|
||||
|
||||
# Verify first client created with correct params
|
||||
call1 = mock_wrapper.call_args_list[0]
|
||||
assert call1[1]["name"] == "client1"
|
||||
assert call1[1]["command"] == "python"
|
||||
assert call1[1]["args"] == ["server1.py"]
|
||||
assert call1[1]["env"] == {"KEY1": "value1"}
|
||||
|
||||
# Verify second client created with correct params
|
||||
call2 = mock_wrapper.call_args_list[1]
|
||||
assert call2[1]["name"] == "client2"
|
||||
assert call2[1]["command"] == "python"
|
||||
assert call2[1]["args"] == ["server2.py"]
|
||||
assert call2[1]["env"] == {"KEY2": "value2"}
|
||||
|
||||
# Verify clients were connected
|
||||
mock_client1.connect.assert_called_once()
|
||||
mock_client2.connect.assert_called_once()
|
||||
|
||||
# Verify tools were listed
|
||||
mock_client1.list_tools.assert_called_once()
|
||||
mock_client2.list_tools.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_client_for_tool(mock_client1, mock_client2, client_config):
|
||||
"""Test getting client that provides a specific tool."""
|
||||
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
||||
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
||||
|
||||
manager = MCPClientManager(client_config)
|
||||
await manager.initialize()
|
||||
|
||||
# Get client for tool1 (from client1)
|
||||
client = await manager.get_client_for_tool("tool1")
|
||||
assert client == mock_client1
|
||||
|
||||
# Get client for tool3 (from client2)
|
||||
client = await manager.get_client_for_tool("tool3")
|
||||
assert client == mock_client2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_client_for_unknown_tool(mock_client1, mock_client2, client_config):
|
||||
"""Test getting client for tool that doesn't exist."""
|
||||
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
||||
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
||||
|
||||
manager = MCPClientManager(client_config)
|
||||
await manager.initialize()
|
||||
|
||||
# Try to get client for non-existent tool
|
||||
with pytest.raises(KeyError, match="Tool 'unknown_tool' not found"):
|
||||
await manager.get_client_for_tool("unknown_tool")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_all_tools(mock_client1, mock_client2, client_config):
|
||||
"""Test listing all tools across all clients."""
|
||||
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
||||
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
||||
|
||||
manager = MCPClientManager(client_config)
|
||||
await manager.initialize()
|
||||
|
||||
tools = await manager.list_all_tools()
|
||||
|
||||
# Should have all tools from both clients
|
||||
assert set(tools) == {"tool1", "tool2", "tool3", "tool4"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_detect_tool_name_collision():
|
||||
"""Test detection of tool name collisions across clients."""
|
||||
# Create clients with overlapping tool names
|
||||
client1 = AsyncMock()
|
||||
client1.name = "client1"
|
||||
client1.connect = AsyncMock()
|
||||
client1.list_tools = AsyncMock(return_value=["tool1", "tool2"])
|
||||
|
||||
client2 = AsyncMock()
|
||||
client2.name = "client2"
|
||||
client2.connect = AsyncMock()
|
||||
client2.list_tools = AsyncMock(return_value=["tool2", "tool3"]) # tool2 collision!
|
||||
|
||||
config = {
|
||||
"client1": {"command": "python", "args": ["server1.py"]},
|
||||
"client2": {"command": "python", "args": ["server2.py"]}
|
||||
}
|
||||
|
||||
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
||||
mock_wrapper.side_effect = [client1, client2]
|
||||
|
||||
manager = MCPClientManager(config)
|
||||
|
||||
# Should raise ValueError about collision during initialization
|
||||
with pytest.raises(ValueError, match="Tool name collision.*tool2.*client1.*client2"):
|
||||
await manager.initialize()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_tool_schema(mock_client1, mock_client2, client_config):
|
||||
"""Test getting tool schema via manager."""
|
||||
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
||||
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
||||
|
||||
manager = MCPClientManager(client_config)
|
||||
await manager.initialize()
|
||||
|
||||
# Get schema for tool1 (from client1)
|
||||
schema = await manager.get_tool_schema("tool1")
|
||||
assert schema == {"type": "object", "properties": {"param1": {"type": "string"}}}
|
||||
mock_client1.get_tool_schema.assert_called_once_with("tool1")
|
||||
|
||||
# Get schema for tool3 (from client2)
|
||||
schema = await manager.get_tool_schema("tool3")
|
||||
assert schema == {"type": "object", "properties": {"param2": {"type": "number"}}}
|
||||
mock_client2.get_tool_schema.assert_called_once_with("tool3")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_tool(mock_client1, mock_client2, client_config):
|
||||
"""Test calling tool via manager."""
|
||||
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
||||
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
||||
|
||||
manager = MCPClientManager(client_config)
|
||||
await manager.initialize()
|
||||
|
||||
# Call tool1 (from client1)
|
||||
result = await manager.call_tool("tool1", {"param1": "value"})
|
||||
assert result == {"result": "success"}
|
||||
mock_client1.call_tool.assert_called_once_with("tool1", {"param1": "value"})
|
||||
|
||||
# Call tool3 (from client2)
|
||||
result = await manager.call_tool("tool3", {"param2": 42})
|
||||
assert result == {"result": "success2"}
|
||||
mock_client2.call_tool.assert_called_once_with("tool3", {"param2": 42})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shutdown_all_clients(mock_client1, mock_client2, client_config):
|
||||
"""Test shutting down all clients."""
|
||||
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
||||
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
||||
|
||||
manager = MCPClientManager(client_config)
|
||||
await manager.initialize()
|
||||
|
||||
# Shutdown all clients
|
||||
await manager.shutdown()
|
||||
|
||||
# Verify both clients were disconnected
|
||||
mock_client1.disconnect.assert_called_once()
|
||||
mock_client2.disconnect.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_manager_before_initialization():
|
||||
"""Test that manager methods fail before initialization."""
|
||||
config = {"client1": {"command": "python", "args": ["server.py"]}}
|
||||
manager = MCPClientManager(config)
|
||||
|
||||
# Should raise RuntimeError if not initialized
|
||||
with pytest.raises(RuntimeError, match="Manager not initialized"):
|
||||
await manager.list_all_tools()
|
||||
|
||||
with pytest.raises(RuntimeError, match="Manager not initialized"):
|
||||
await manager.get_client_for_tool("tool1")
|
||||
|
||||
with pytest.raises(RuntimeError, match="Manager not initialized"):
|
||||
await manager.get_tool_schema("tool1")
|
||||
|
||||
with pytest.raises(RuntimeError, match="Manager not initialized"):
|
||||
await manager.call_tool("tool1", {})
|
||||
Loading…
Add table
Add a link
Reference in a new issue