mirror of
https://github.com/hanzoai/agent.git
synced 2026-09-11 21:08:51 +00:00
zeekay
ff7ef7114f
Consolidate all Hanzo agent infrastructure into one repo: - control-plane/: Go control plane server (was agents-new) - sdk/: Go, Python, TypeScript SDKs - deployments/: Docker, Kubernetes, Helm configs - Dockerfile.sdk: Python SDK image (ghcr.io/hanzoai/agent:latest) - CI builds both SDK + control-plane images No more agentfield — everything is hanzoai/agent.
398 lines
13 KiB
Python
398 lines
13 KiB
Python
"""
|
|
Tests for MCP Client
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import AsyncMock, patch
|
|
from aiohttp import ClientSession
|
|
from hanzo_agents.mcp_client import MCPClient
|
|
|
|
|
|
class TestMCPClientInitialization:
|
|
"""Test MCP Client initialization"""
|
|
|
|
def test_init_basic(self):
|
|
"""Test basic initialization"""
|
|
client = MCPClient("http://localhost:8080", "test-server")
|
|
|
|
assert client.base_url == "http://localhost:8080"
|
|
assert client.server_alias == "test-server"
|
|
assert client.dev_mode is False
|
|
assert client.session is None
|
|
assert client._is_stdio_bridge is False
|
|
|
|
def test_init_with_dev_mode(self):
|
|
"""Test initialization with dev mode"""
|
|
client = MCPClient("http://localhost:8080", "test-server", dev_mode=True)
|
|
|
|
assert client.dev_mode is True
|
|
|
|
def test_from_port_legacy(self):
|
|
"""Test legacy from_port constructor"""
|
|
client = MCPClient.from_port("test-server", 8080, dev_mode=False)
|
|
|
|
assert client.base_url == "http://localhost:8080"
|
|
assert client.server_alias == "test-server"
|
|
assert client.dev_mode is False
|
|
|
|
def test_from_port_with_dev_mode(self):
|
|
"""Test from_port with dev mode"""
|
|
client = MCPClient.from_port("test-server", 9000, dev_mode=True)
|
|
|
|
assert client.base_url == "http://localhost:9000"
|
|
assert client.dev_mode is True
|
|
|
|
|
|
class TestMCPClientSession:
|
|
"""Test session management"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ensure_session_creates_new(self):
|
|
"""Test that _ensure_session creates new session"""
|
|
client = MCPClient("http://localhost:8080", "test-server")
|
|
|
|
assert client.session is None
|
|
|
|
await client._ensure_session()
|
|
|
|
assert client.session is not None
|
|
assert isinstance(client.session, ClientSession)
|
|
|
|
# Cleanup
|
|
await client.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ensure_session_reuses_existing(self):
|
|
"""Test that _ensure_session reuses existing session"""
|
|
client = MCPClient("http://localhost:8080", "test-server")
|
|
|
|
await client._ensure_session()
|
|
first_session = client.session
|
|
|
|
await client._ensure_session()
|
|
second_session = client.session
|
|
|
|
assert first_session is second_session
|
|
|
|
# Cleanup
|
|
await client.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_close_session(self):
|
|
"""Test closing session"""
|
|
client = MCPClient("http://localhost:8080", "test-server")
|
|
|
|
await client._ensure_session()
|
|
assert client.session is not None
|
|
assert not client.session.closed
|
|
|
|
await client.close()
|
|
|
|
assert client.session.closed
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_close_without_session(self):
|
|
"""Test closing when no session exists"""
|
|
client = MCPClient("http://localhost:8080", "test-server")
|
|
|
|
# Should not raise error
|
|
await client.close()
|
|
|
|
assert client.session is None
|
|
|
|
|
|
class TestMCPClientHealthCheck:
|
|
"""Test health check functionality"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_health_check_success(self):
|
|
"""Test successful health check"""
|
|
client = MCPClient("http://localhost:8080", "test-server")
|
|
|
|
mock_response = AsyncMock()
|
|
mock_response.status = 200
|
|
mock_response.__aenter__.return_value = mock_response
|
|
mock_response.__aexit__.return_value = None
|
|
|
|
with patch.object(ClientSession, "get", return_value=mock_response):
|
|
result = await client.health_check()
|
|
|
|
assert result is True
|
|
|
|
await client.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_health_check_failure_404(self):
|
|
"""Test health check with 404 status"""
|
|
client = MCPClient("http://localhost:8080", "test-server")
|
|
|
|
mock_response = AsyncMock()
|
|
mock_response.status = 404
|
|
mock_response.__aenter__.return_value = mock_response
|
|
mock_response.__aexit__.return_value = None
|
|
|
|
with patch.object(ClientSession, "get", return_value=mock_response):
|
|
result = await client.health_check()
|
|
|
|
assert result is False
|
|
|
|
await client.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_health_check_network_error(self):
|
|
"""Test health check with network error"""
|
|
client = MCPClient("http://localhost:8080", "test-server", dev_mode=True)
|
|
|
|
with patch.object(
|
|
ClientSession, "get", side_effect=ConnectionError("Connection refused")
|
|
):
|
|
result = await client.health_check()
|
|
|
|
assert result is False
|
|
|
|
await client.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_health_check_timeout(self):
|
|
"""Test health check with timeout"""
|
|
client = MCPClient("http://localhost:8080", "test-server")
|
|
|
|
with patch.object(ClientSession, "get", side_effect=TimeoutError):
|
|
result = await client.health_check()
|
|
|
|
assert result is False
|
|
|
|
await client.close()
|
|
|
|
|
|
class TestMCPClientListTools:
|
|
"""Test list_tools functionality"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_tools_direct_http_success(self):
|
|
"""Test listing tools via direct HTTP (non-stdio)"""
|
|
client = MCPClient("http://localhost:8080", "test-server")
|
|
client._is_stdio_bridge = False
|
|
|
|
mock_response = AsyncMock()
|
|
mock_response.status = 200
|
|
mock_response.json = AsyncMock(
|
|
return_value={
|
|
"result": {
|
|
"tools": [
|
|
{"name": "tool1", "description": "Test tool 1"},
|
|
{"name": "tool2", "description": "Test tool 2"},
|
|
]
|
|
}
|
|
}
|
|
)
|
|
mock_response.__aenter__.return_value = mock_response
|
|
mock_response.__aexit__.return_value = None
|
|
|
|
with patch.object(ClientSession, "post", return_value=mock_response):
|
|
tools = await client.list_tools()
|
|
|
|
assert len(tools) == 2
|
|
assert tools[0]["name"] == "tool1"
|
|
assert tools[1]["name"] == "tool2"
|
|
|
|
await client.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_tools_stdio_bridge_success(self):
|
|
"""Test listing tools via stdio bridge"""
|
|
client = MCPClient("http://localhost:8080", "test-server")
|
|
client._is_stdio_bridge = True
|
|
|
|
mock_response = AsyncMock()
|
|
mock_response.status = 200
|
|
mock_response.json = AsyncMock(
|
|
return_value={
|
|
"tools": [
|
|
{"name": "bridge_tool1", "description": "Bridge tool 1"},
|
|
{"name": "bridge_tool2", "description": "Bridge tool 2"},
|
|
]
|
|
}
|
|
)
|
|
mock_response.__aenter__.return_value = mock_response
|
|
mock_response.__aexit__.return_value = None
|
|
|
|
with patch.object(ClientSession, "post", return_value=mock_response):
|
|
tools = await client.list_tools()
|
|
|
|
assert len(tools) == 2
|
|
assert tools[0]["name"] == "bridge_tool1"
|
|
|
|
await client.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_tools_empty_result(self):
|
|
"""Test listing tools with empty result"""
|
|
client = MCPClient("http://localhost:8080", "test-server")
|
|
|
|
mock_response = AsyncMock()
|
|
mock_response.status = 200
|
|
mock_response.json = AsyncMock(return_value={"result": {"tools": []}})
|
|
mock_response.__aenter__.return_value = mock_response
|
|
mock_response.__aexit__.return_value = None
|
|
|
|
with patch.object(ClientSession, "post", return_value=mock_response):
|
|
tools = await client.list_tools()
|
|
|
|
assert len(tools) == 0
|
|
assert tools == []
|
|
|
|
await client.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_tools_network_error(self):
|
|
"""Test list_tools with network error"""
|
|
client = MCPClient("http://localhost:8080", "test-server", dev_mode=True)
|
|
|
|
with patch.object(
|
|
ClientSession, "post", side_effect=ConnectionError("Connection refused")
|
|
):
|
|
tools = await client.list_tools()
|
|
|
|
assert tools == []
|
|
|
|
await client.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_tools_malformed_response(self):
|
|
"""Test list_tools with malformed JSON response"""
|
|
client = MCPClient("http://localhost:8080", "test-server")
|
|
|
|
mock_response = AsyncMock()
|
|
mock_response.status = 200
|
|
mock_response.json = AsyncMock(
|
|
return_value={"error": "malformed"} # Missing 'result' key
|
|
)
|
|
mock_response.__aenter__.return_value = mock_response
|
|
mock_response.__aexit__.return_value = None
|
|
|
|
with patch.object(ClientSession, "post", return_value=mock_response):
|
|
tools = await client.list_tools()
|
|
|
|
assert tools == []
|
|
|
|
await client.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_tools_http_500(self):
|
|
"""Test list_tools with server error"""
|
|
client = MCPClient("http://localhost:8080", "test-server")
|
|
|
|
mock_response = AsyncMock()
|
|
mock_response.status = 500
|
|
mock_response.__aenter__.return_value = mock_response
|
|
mock_response.__aexit__.return_value = None
|
|
|
|
with patch.object(ClientSession, "post", return_value=mock_response):
|
|
tools = await client.list_tools()
|
|
|
|
assert tools == []
|
|
|
|
await client.close()
|
|
|
|
|
|
class TestMCPClientEdgeCases:
|
|
"""Test edge cases and error handling"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_multiple_operations_on_same_client(self):
|
|
"""Test performing multiple operations on same client"""
|
|
client = MCPClient("http://localhost:8080", "test-server")
|
|
|
|
mock_health_response = AsyncMock()
|
|
mock_health_response.status = 200
|
|
mock_health_response.__aenter__.return_value = mock_health_response
|
|
mock_health_response.__aexit__.return_value = None
|
|
|
|
mock_tools_response = AsyncMock()
|
|
mock_tools_response.status = 200
|
|
mock_tools_response.json = AsyncMock(
|
|
return_value={"result": {"tools": [{"name": "tool1"}]}}
|
|
)
|
|
mock_tools_response.__aenter__.return_value = mock_tools_response
|
|
mock_tools_response.__aexit__.return_value = None
|
|
|
|
with patch.object(ClientSession, "get", return_value=mock_health_response):
|
|
health1 = await client.health_check()
|
|
|
|
assert health1 is True
|
|
|
|
with patch.object(ClientSession, "post", return_value=mock_tools_response):
|
|
tools = await client.list_tools()
|
|
|
|
assert len(tools) == 1
|
|
|
|
with patch.object(ClientSession, "get", return_value=mock_health_response):
|
|
health2 = await client.health_check()
|
|
|
|
assert health2 is True
|
|
|
|
await client.close()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_operations_after_close(self):
|
|
"""Test that operations work after close (should recreate session)"""
|
|
client = MCPClient("http://localhost:8080", "test-server")
|
|
|
|
# First operation
|
|
mock_response = AsyncMock()
|
|
mock_response.status = 200
|
|
mock_response.__aenter__.return_value = mock_response
|
|
mock_response.__aexit__.return_value = None
|
|
|
|
with patch.object(ClientSession, "get", return_value=mock_response):
|
|
result1 = await client.health_check()
|
|
|
|
assert result1 is True
|
|
|
|
# Close session
|
|
await client.close()
|
|
assert client.session.closed
|
|
|
|
# Second operation should work (creates new session)
|
|
with patch.object(ClientSession, "get", return_value=mock_response):
|
|
_ = await client.health_check()
|
|
|
|
# This will fail if _ensure_session doesn't handle closed sessions
|
|
# Current implementation may need fixing for this case
|
|
# Tests document the expected behavior
|
|
|
|
await client.close()
|
|
|
|
def test_client_attributes_immutable(self):
|
|
"""Test that client attributes are properly set"""
|
|
client = MCPClient("http://localhost:8080", "test-server", dev_mode=True)
|
|
|
|
# Attributes should be accessible
|
|
assert client.base_url == "http://localhost:8080"
|
|
assert client.server_alias == "test-server"
|
|
assert client.dev_mode is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_health_checks(self):
|
|
"""Test multiple concurrent health checks"""
|
|
import asyncio
|
|
|
|
client = MCPClient("http://localhost:8080", "test-server")
|
|
|
|
mock_response = AsyncMock()
|
|
mock_response.status = 200
|
|
mock_response.__aenter__.return_value = mock_response
|
|
mock_response.__aexit__.return_value = None
|
|
|
|
with patch.object(ClientSession, "get", return_value=mock_response):
|
|
results = await asyncio.gather(
|
|
client.health_check(),
|
|
client.health_check(),
|
|
client.health_check(),
|
|
)
|
|
|
|
assert all(results)
|
|
|
|
await client.close()
|