wehub-resource-sync
424 lines
15 KiB
Python
424 lines
15 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Tests for MCP tool executor (omlx/mcp/executor.py).
|
|
"""
|
|
|
|
import asyncio
|
|
from typing import Any, Dict
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from omlx.mcp.executor import ToolExecutor, execute_single_tool
|
|
from omlx.mcp.manager import MCPClientManager
|
|
from omlx.mcp.types import (
|
|
MCPConfig,
|
|
MCPServerConfig,
|
|
MCPServerState,
|
|
MCPTool,
|
|
MCPToolResult,
|
|
MCPTransport,
|
|
)
|
|
|
|
|
|
class TestToolExecutorInit:
|
|
"""Tests for ToolExecutor initialization."""
|
|
|
|
def test_init_defaults(self):
|
|
"""Test executor initialization with defaults."""
|
|
config = MCPConfig(default_timeout=30.0)
|
|
manager = MCPClientManager(config)
|
|
executor = ToolExecutor(manager)
|
|
|
|
assert executor.manager is manager
|
|
assert executor.max_parallel == 5
|
|
assert executor.default_timeout == 30.0
|
|
|
|
def test_init_custom_values(self):
|
|
"""Test executor initialization with custom values."""
|
|
config = MCPConfig(default_timeout=60.0)
|
|
manager = MCPClientManager(config)
|
|
executor = ToolExecutor(manager, max_parallel=10, default_timeout=120.0)
|
|
|
|
assert executor.max_parallel == 10
|
|
assert executor.default_timeout == 120.0
|
|
|
|
|
|
class TestToolExecutorExecuteToolCalls:
|
|
"""Tests for ToolExecutor.execute_tool_calls()."""
|
|
|
|
@pytest.fixture
|
|
def executor_with_manager(self) -> ToolExecutor:
|
|
"""Create executor with mock manager."""
|
|
config = MCPConfig.from_dict({
|
|
"servers": {
|
|
"test": {
|
|
"transport": "stdio",
|
|
"command": "python",
|
|
},
|
|
},
|
|
"default_timeout": 10.0,
|
|
})
|
|
manager = MCPClientManager(config)
|
|
manager._clients["test"]._state = MCPServerState.CONNECTED
|
|
manager._clients["test"]._tools = [
|
|
MCPTool(server_name="test", name="tool1", description=""),
|
|
MCPTool(server_name="test", name="tool2", description=""),
|
|
]
|
|
return ToolExecutor(manager)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_empty_list(self, executor_with_manager: ToolExecutor):
|
|
"""Test executing empty tool calls list."""
|
|
results = await executor_with_manager.execute_tool_calls([])
|
|
assert results == []
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_parallel(self, executor_with_manager: ToolExecutor):
|
|
"""Test parallel execution of tool calls."""
|
|
executor_with_manager.manager.execute_tool_call = AsyncMock(
|
|
side_effect=[
|
|
MCPToolResult(tool_name="tool1", content="Result 1"),
|
|
MCPToolResult(tool_name="tool2", content="Result 2"),
|
|
]
|
|
)
|
|
|
|
tool_calls = [
|
|
{"id": "call_1", "function": {"name": "test__tool1", "arguments": "{}"}},
|
|
{"id": "call_2", "function": {"name": "test__tool2", "arguments": "{}"}},
|
|
]
|
|
|
|
results = await executor_with_manager.execute_tool_calls(tool_calls, parallel=True)
|
|
|
|
assert len(results) == 2
|
|
assert results[0][0].content == "Result 1"
|
|
assert results[0][1] == "call_1"
|
|
assert results[1][0].content == "Result 2"
|
|
assert results[1][1] == "call_2"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_sequential(self, executor_with_manager: ToolExecutor):
|
|
"""Test sequential execution of tool calls."""
|
|
call_order = []
|
|
|
|
async def track_call(tc, **kwargs):
|
|
name = tc["function"]["name"]
|
|
call_order.append(name)
|
|
return MCPToolResult(tool_name=name, content=f"Result for {name}")
|
|
|
|
executor_with_manager.manager.execute_tool_call = AsyncMock(side_effect=track_call)
|
|
|
|
tool_calls = [
|
|
{"id": "call_1", "function": {"name": "tool1", "arguments": "{}"}},
|
|
{"id": "call_2", "function": {"name": "tool2", "arguments": "{}"}},
|
|
]
|
|
|
|
results = await executor_with_manager.execute_tool_calls(tool_calls, parallel=False)
|
|
|
|
assert len(results) == 2
|
|
# Verify sequential order
|
|
assert call_order == ["tool1", "tool2"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_parallel_with_exception(self, executor_with_manager: ToolExecutor):
|
|
"""Test parallel execution handles exceptions."""
|
|
executor_with_manager.manager.execute_tool_call = AsyncMock(
|
|
side_effect=[
|
|
MCPToolResult(tool_name="tool1", content="Success"),
|
|
RuntimeError("Tool failed"),
|
|
]
|
|
)
|
|
|
|
tool_calls = [
|
|
{"id": "call_1", "function": {"name": "tool1", "arguments": "{}"}},
|
|
{"id": "call_2", "function": {"name": "tool2", "arguments": "{}"}},
|
|
]
|
|
|
|
results = await executor_with_manager.execute_tool_calls(tool_calls, parallel=True)
|
|
|
|
assert len(results) == 2
|
|
assert results[0][0].is_error is False
|
|
assert results[1][0].is_error is True
|
|
assert "Tool failed" in results[1][0].error_message
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_sequential_with_exception(self, executor_with_manager: ToolExecutor):
|
|
"""Test sequential execution handles exceptions."""
|
|
executor_with_manager.manager.execute_tool_call = AsyncMock(
|
|
side_effect=[
|
|
RuntimeError("First failed"),
|
|
MCPToolResult(tool_name="tool2", content="Success"),
|
|
]
|
|
)
|
|
|
|
tool_calls = [
|
|
{"id": "call_1", "function": {"name": "tool1", "arguments": "{}"}},
|
|
{"id": "call_2", "function": {"name": "tool2", "arguments": "{}"}},
|
|
]
|
|
|
|
results = await executor_with_manager.execute_tool_calls(tool_calls, parallel=False)
|
|
|
|
assert len(results) == 2
|
|
assert results[0][0].is_error is True
|
|
assert results[1][0].is_error is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_parallel_respects_semaphore(self, executor_with_manager: ToolExecutor):
|
|
"""Test parallel execution respects max_parallel limit."""
|
|
executor_with_manager.max_parallel = 2
|
|
concurrent_count = []
|
|
current_concurrent = 0
|
|
lock = asyncio.Lock()
|
|
|
|
async def track_concurrency(tc, **kwargs):
|
|
nonlocal current_concurrent
|
|
async with lock:
|
|
current_concurrent += 1
|
|
concurrent_count.append(current_concurrent)
|
|
await asyncio.sleep(0.05) # Simulate work
|
|
async with lock:
|
|
current_concurrent -= 1
|
|
return MCPToolResult(tool_name="tool", content="ok")
|
|
|
|
executor_with_manager.manager.execute_tool_call = AsyncMock(side_effect=track_concurrency)
|
|
|
|
tool_calls = [
|
|
{"id": f"call_{i}", "function": {"name": f"tool{i}", "arguments": "{}"}}
|
|
for i in range(5)
|
|
]
|
|
|
|
await executor_with_manager.execute_tool_calls(tool_calls, parallel=True)
|
|
|
|
# Max concurrent should never exceed 2
|
|
assert max(concurrent_count) <= 2
|
|
|
|
|
|
class TestToolExecutorExecuteAndFormat:
|
|
"""Tests for ToolExecutor.execute_and_format()."""
|
|
|
|
@pytest.fixture
|
|
def executor(self) -> ToolExecutor:
|
|
"""Create executor with mock manager."""
|
|
config = MCPConfig(default_timeout=10.0)
|
|
manager = MCPClientManager(config)
|
|
return ToolExecutor(manager)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_and_format(self, executor: ToolExecutor):
|
|
"""Test execute_and_format returns formatted messages."""
|
|
executor.manager.execute_tool_call = AsyncMock(
|
|
return_value=MCPToolResult(tool_name="test", content="Result data")
|
|
)
|
|
|
|
tool_calls = [
|
|
{"id": "call_123", "function": {"name": "test__tool", "arguments": "{}"}},
|
|
]
|
|
|
|
messages = await executor.execute_and_format(tool_calls)
|
|
|
|
assert len(messages) == 1
|
|
assert messages[0]["role"] == "tool"
|
|
assert messages[0]["tool_call_id"] == "call_123"
|
|
assert messages[0]["content"] == "Result data"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_and_format_error(self, executor: ToolExecutor):
|
|
"""Test execute_and_format formats error messages."""
|
|
executor.manager.execute_tool_call = AsyncMock(
|
|
return_value=MCPToolResult(
|
|
tool_name="test",
|
|
content=None,
|
|
is_error=True,
|
|
error_message="Tool failed",
|
|
)
|
|
)
|
|
|
|
tool_calls = [
|
|
{"id": "call_456", "function": {"name": "test__tool", "arguments": "{}"}},
|
|
]
|
|
|
|
messages = await executor.execute_and_format(tool_calls)
|
|
|
|
assert len(messages) == 1
|
|
assert "Error: Tool failed" in messages[0]["content"]
|
|
|
|
|
|
class TestToolExecutorExtractAndValidate:
|
|
"""Tests for ToolExecutor.extract_and_validate()."""
|
|
|
|
@pytest.fixture
|
|
def executor_with_tools(self) -> ToolExecutor:
|
|
"""Create executor with tools."""
|
|
config = MCPConfig.from_dict({
|
|
"servers": {
|
|
"test": {"transport": "stdio", "command": "python"},
|
|
}
|
|
})
|
|
manager = MCPClientManager(config)
|
|
manager._clients["test"]._state = MCPServerState.CONNECTED
|
|
manager._clients["test"]._tools = [
|
|
MCPTool(server_name="test", name="valid_tool", description=""),
|
|
]
|
|
return ToolExecutor(manager)
|
|
|
|
def test_extract_valid_tools(self, executor_with_tools: ToolExecutor):
|
|
"""Test extracting and validating valid tool calls."""
|
|
response = {
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_1",
|
|
"function": {
|
|
"name": "test__valid_tool",
|
|
"arguments": "{}",
|
|
},
|
|
}
|
|
]
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
calls, all_valid = executor_with_tools.extract_and_validate(response)
|
|
|
|
assert len(calls) == 1
|
|
assert all_valid is True
|
|
|
|
def test_extract_invalid_tool(self, executor_with_tools: ToolExecutor):
|
|
"""Test extracting and validating invalid tool."""
|
|
response = {
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_1",
|
|
"function": {
|
|
"name": "nonexistent_tool",
|
|
"arguments": "{}",
|
|
},
|
|
}
|
|
]
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
calls, all_valid = executor_with_tools.extract_and_validate(response)
|
|
|
|
assert len(calls) == 1
|
|
assert all_valid is False
|
|
|
|
def test_extract_no_tool_calls(self, executor_with_tools: ToolExecutor):
|
|
"""Test extracting from response without tool calls."""
|
|
response = {
|
|
"choices": [{"message": {"content": "No tools"}}]
|
|
}
|
|
|
|
calls, all_valid = executor_with_tools.extract_and_validate(response)
|
|
|
|
assert calls == []
|
|
assert all_valid is True
|
|
|
|
def test_extract_mixed_valid_invalid(self, executor_with_tools: ToolExecutor):
|
|
"""Test extracting mix of valid and invalid tools."""
|
|
response = {
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_1",
|
|
"function": {"name": "test__valid_tool", "arguments": "{}"},
|
|
},
|
|
{
|
|
"id": "call_2",
|
|
"function": {"name": "invalid_tool", "arguments": "{}"},
|
|
},
|
|
]
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
calls, all_valid = executor_with_tools.extract_and_validate(response)
|
|
|
|
assert len(calls) == 2
|
|
assert all_valid is False
|
|
|
|
|
|
class TestToolExecutorToolExists:
|
|
"""Tests for ToolExecutor._tool_exists()."""
|
|
|
|
@pytest.fixture
|
|
def executor(self) -> ToolExecutor:
|
|
"""Create executor with tools."""
|
|
config = MCPConfig.from_dict({
|
|
"servers": {
|
|
"server1": {"transport": "stdio", "command": "python"},
|
|
"server2": {"transport": "stdio", "command": "node"},
|
|
}
|
|
})
|
|
manager = MCPClientManager(config)
|
|
manager._clients["server1"]._state = MCPServerState.CONNECTED
|
|
manager._clients["server1"]._tools = [
|
|
MCPTool(server_name="server1", name="tool_a", description=""),
|
|
]
|
|
manager._clients["server2"]._state = MCPServerState.CONNECTED
|
|
manager._clients["server2"]._tools = [
|
|
MCPTool(server_name="server2", name="tool_b", description=""),
|
|
]
|
|
return ToolExecutor(manager)
|
|
|
|
def test_tool_exists_full_name(self, executor: ToolExecutor):
|
|
"""Test finding tool by full name."""
|
|
assert executor._tool_exists("server1__tool_a") is True
|
|
assert executor._tool_exists("server2__tool_b") is True
|
|
assert executor._tool_exists("server1__nonexistent") is False
|
|
|
|
def test_tool_exists_short_name(self, executor: ToolExecutor):
|
|
"""Test finding tool by short name."""
|
|
assert executor._tool_exists("tool_a") is True
|
|
assert executor._tool_exists("tool_b") is True
|
|
assert executor._tool_exists("nonexistent") is False
|
|
|
|
|
|
class TestExecuteSingleTool:
|
|
"""Tests for execute_single_tool convenience function."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_single_tool(self):
|
|
"""Test execute_single_tool calls manager."""
|
|
config = MCPConfig()
|
|
manager = MCPClientManager(config)
|
|
manager.execute_tool = AsyncMock(
|
|
return_value=MCPToolResult(tool_name="tool", content="Result")
|
|
)
|
|
|
|
result = await execute_single_tool(
|
|
manager,
|
|
"server__tool",
|
|
{"arg": "value"},
|
|
timeout=30.0,
|
|
)
|
|
|
|
assert result.content == "Result"
|
|
manager.execute_tool.assert_called_once_with(
|
|
"server__tool", {"arg": "value"}, 30.0
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_single_tool_default_timeout(self):
|
|
"""Test execute_single_tool with default timeout."""
|
|
config = MCPConfig()
|
|
manager = MCPClientManager(config)
|
|
manager.execute_tool = AsyncMock(
|
|
return_value=MCPToolResult(tool_name="tool", content="Result")
|
|
)
|
|
|
|
await execute_single_tool(manager, "tool", {})
|
|
|
|
manager.execute_tool.assert_called_once_with("tool", {}, None)
|