wehub-resource-sync
77 lines
2.9 KiB
Python
77 lines
2.9 KiB
Python
import unittest
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
from trae_agent.tools.base import ToolCallArguments, ToolExecResult
|
|
from trae_agent.tools.mcp_tool import MCPTool
|
|
|
|
|
|
class TestMCPTool(unittest.IsolatedAsyncioTestCase):
|
|
def setUp(self):
|
|
# simulate a tool schema
|
|
self.mock_tool = MagicMock()
|
|
self.mock_tool.name = "test_tool"
|
|
self.mock_tool.description = "A test tool"
|
|
self.mock_tool.inputSchema = {
|
|
"required": ["param1"],
|
|
"properties": {
|
|
"param1": {"type": "string", "description": "First parameter"},
|
|
"param2": {"type": "integer", "description": "Second parameter"},
|
|
},
|
|
}
|
|
|
|
# simulate client side
|
|
self.mock_client = MagicMock()
|
|
self.tool = MCPTool(self.mock_client, self.mock_tool, model_provider="test_provider")
|
|
|
|
def test_get_name(self):
|
|
self.assertEqual(self.tool.get_name(), "test_tool")
|
|
|
|
def test_get_description(self):
|
|
self.assertEqual(self.tool.get_description(), "A test tool")
|
|
|
|
def test_get_model_provider(self):
|
|
self.assertEqual(self.tool.get_model_provider(), "test_provider")
|
|
|
|
def test_get_parameters(self):
|
|
params = self.tool.get_parameters()
|
|
self.assertEqual(len(params), 2)
|
|
self.assertTrue(any(p.name == "param1" and p.required for p in params))
|
|
self.assertTrue(any(p.name == "param2" and not p.required for p in params))
|
|
|
|
async def test_execute_success(self):
|
|
mock_response = MagicMock()
|
|
mock_response.isError = False
|
|
mock_response.content = [MagicMock(text="Execution successful")]
|
|
self.mock_client.call_tool = AsyncMock(return_value=mock_response)
|
|
|
|
arguments = ToolCallArguments(arguments={"param1": "value", "param2": 123})
|
|
result: ToolExecResult = await self.tool.execute(arguments)
|
|
|
|
self.assertIsNone(result.error)
|
|
self.assertEqual(result.output, "Execution successful")
|
|
|
|
async def test_execute_failure(self):
|
|
mock_response = MagicMock()
|
|
mock_response.isError = True
|
|
mock_response.content = [MagicMock(text="Something went wrong")]
|
|
self.mock_client.call_tool = AsyncMock(return_value=mock_response)
|
|
|
|
arguments = ToolCallArguments(arguments={"param1": "value"})
|
|
result: ToolExecResult = await self.tool.execute(arguments)
|
|
|
|
self.assertIsNone(result.output)
|
|
self.assertEqual(result.error, "Something went wrong")
|
|
|
|
async def test_execute_exception(self):
|
|
self.mock_client.call_tool = AsyncMock(side_effect=RuntimeError("Tool crashed"))
|
|
|
|
arguments = ToolCallArguments(arguments={"param1": "value"})
|
|
result: ToolExecResult = await self.tool.execute(arguments)
|
|
|
|
self.assertIn("Error running mcp tool", result.error)
|
|
self.assertEqual(result.error_code, -1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|