106 lines
3.6 KiB
Python
106 lines
3.6 KiB
Python
import unittest
|
|
|
|
from models import RuntimeMcpServer, RuntimeTool
|
|
from services.runtime_variables import DynamicVariableStore
|
|
from services.tool_executor import ToolExecutionError, ToolExecutor
|
|
from services.tools import McpClientError
|
|
|
|
|
|
def mcp_tool() -> RuntimeTool:
|
|
return RuntimeTool(
|
|
id="tool_mcp_order",
|
|
name="查询订单",
|
|
function_name="query_order",
|
|
type="mcp",
|
|
description="查询订单状态",
|
|
definition={
|
|
"schema_version": 1,
|
|
"type": "mcp",
|
|
"config": {
|
|
"remote_tool_name": "get_order",
|
|
"input_schema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"order_id": {
|
|
"type": "string",
|
|
"description": "订单编号",
|
|
}
|
|
},
|
|
"required": ["order_id"],
|
|
},
|
|
"dynamic_variable_assignments": {
|
|
"order_status": "structuredContent.status"
|
|
},
|
|
},
|
|
},
|
|
mcp_server=RuntimeMcpServer(
|
|
id="mcp_orders",
|
|
name="订单服务",
|
|
url="https://mcp.example.com/{{tenant}}",
|
|
headers={"X-Tenant": "{{tenant}}"},
|
|
secret_headers={"Authorization": "Bearer secret"},
|
|
),
|
|
)
|
|
|
|
|
|
class FakeMcpClient:
|
|
def __init__(self):
|
|
self.calls = []
|
|
self.error: Exception | None = None
|
|
|
|
async def call_tool(self, server, remote_tool_name, arguments):
|
|
self.calls.append((server, remote_tool_name, arguments))
|
|
if self.error:
|
|
raise self.error
|
|
return {
|
|
"text": "订单已发货",
|
|
"structuredContent": {"status": "shipped"},
|
|
"content": [{"type": "text", "text": "订单已发货"}],
|
|
}
|
|
|
|
|
|
class ToolExecutorMcpTests(unittest.IsolatedAsyncioTestCase):
|
|
def test_mcp_schema_is_exposed_to_llm(self):
|
|
properties, required = ToolExecutor.schema_parts(mcp_tool())
|
|
|
|
self.assertEqual(properties["order_id"]["type"], "string")
|
|
self.assertEqual(required, ["order_id"])
|
|
|
|
async def test_mcp_execution_renders_connection_and_assigns_result(self):
|
|
store = DynamicVariableStore(
|
|
{"tenant": "school-a"},
|
|
variable_types={"order_status": "string"},
|
|
)
|
|
client = FakeMcpClient()
|
|
executor = ToolExecutor(store, mcp_client=client)
|
|
|
|
result = await executor.execute(
|
|
mcp_tool(),
|
|
{"order_id": "A-100"},
|
|
)
|
|
|
|
self.assertEqual(result["status"], "ok")
|
|
self.assertEqual(result["updated_variables"], ["order_status"])
|
|
self.assertEqual(store.values["order_status"], "shipped")
|
|
server, remote_name, arguments = client.calls[0]
|
|
self.assertEqual(server.url, "https://mcp.example.com/school-a")
|
|
self.assertEqual(server.headers["X-Tenant"], "school-a")
|
|
self.assertEqual(server.secret_headers["Authorization"], "Bearer secret")
|
|
self.assertEqual(remote_name, "get_order")
|
|
self.assertEqual(arguments, {"order_id": "A-100"})
|
|
|
|
async def test_mcp_client_errors_use_shared_tool_error(self):
|
|
client = FakeMcpClient()
|
|
client.error = McpClientError("远端不可用")
|
|
executor = ToolExecutor(
|
|
DynamicVariableStore({"tenant": "school-a"}),
|
|
mcp_client=client,
|
|
)
|
|
|
|
with self.assertRaisesRegex(ToolExecutionError, "远端不可用"):
|
|
await executor.execute(mcp_tool(), {"order_id": "A-100"})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|