diff --git a/tests/litellm/experimental_mcp_client/test_tools.py b/tests/litellm/experimental_mcp_client/test_tools.py index 799078f915..7089d83217 100644 --- a/tests/litellm/experimental_mcp_client/test_tools.py +++ b/tests/litellm/experimental_mcp_client/test_tools.py @@ -9,10 +9,16 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -from mcp.types import CallToolRequestParams, CallToolResult, ListToolsResult +from mcp.types import ( + CallToolRequestParams, + CallToolResult, + ListToolsResult, + TextContent, +) from mcp.types import Tool as MCPTool from litellm.experimental_mcp_client.tools import ( + _get_function_arguments, _transform_openai_tool_call_to_mcp_tool_call_request, call_mcp_tool, call_openai_tool, @@ -38,6 +44,27 @@ def mock_session(): return session +@pytest.fixture +def mock_list_tools_result(): + return ListToolsResult( + tools=[ + MCPTool( + name="test_tool", + description="A test tool", + inputSchema={ + "type": "object", + "properties": {"test": {"type": "string"}}, + }, + ) + ] + ) + + +@pytest.fixture +def mock_mcp_tool_call_result(): + return CallToolResult(content=[TextContent(type="text", text="test_output")]) + + def test_transform_mcp_tool_to_openai_tool(mock_mcp_tool): openai_tool = transform_mcp_tool_to_openai_tool(mock_mcp_tool) assert openai_tool["type"] == "function" @@ -58,3 +85,73 @@ def test_transform_openai_tool_call_to_mcp_tool_call_request(mock_mcp_tool): ) assert mcp_tool_call_request.name == "test_tool" assert mcp_tool_call_request.arguments == {"test": "value"} + + +@pytest.mark.asyncio() +async def test_load_mcp_tools_mcp_format(mock_session, mock_list_tools_result): + mock_session.list_tools.return_value = mock_list_tools_result + result = await load_mcp_tools(mock_session, format="mcp") + assert len(result) == 1 + assert isinstance(result[0], MCPTool) + assert result[0].name == "test_tool" + mock_session.list_tools.assert_called_once() + + +@pytest.mark.asyncio() +async def test_load_mcp_tools_openai_format(mock_session, mock_list_tools_result): + mock_session.list_tools.return_value = mock_list_tools_result + result = await load_mcp_tools(mock_session, format="openai") + assert len(result) == 1 + assert result[0]["type"] == "function" + assert result[0]["function"]["name"] == "test_tool" + mock_session.list_tools.assert_called_once() + + +def test_get_function_arguments(): + # Test with string arguments + function = {"arguments": '{"test": "value"}'} + result = _get_function_arguments(function) + assert result == {"test": "value"} + + # Test with dict arguments + function = {"arguments": {"test": "value"}} + result = _get_function_arguments(function) + assert result == {"test": "value"} + + # Test with invalid JSON string + function = {"arguments": "invalid json"} + result = _get_function_arguments(function) + assert result == {} + + # Test with no arguments + function = {} + result = _get_function_arguments(function) + assert result == {} + + +@pytest.mark.asyncio() +async def test_call_openai_tool(mock_session, mock_mcp_tool_call_result): + mock_session.call_tool.return_value = mock_mcp_tool_call_result + openai_tool = { + "function": {"name": "test_tool", "arguments": json.dumps({"test": "value"})} + } + result = await call_openai_tool(mock_session, openai_tool) + print("result of call_openai_tool", result) + assert result.content[0].text == "test_output" + mock_session.call_tool.assert_called_once_with( + name="test_tool", arguments={"test": "value"} + ) + + +@pytest.mark.asyncio() +async def test_call_mcp_tool(mock_session, mock_mcp_tool_call_result): + mock_session.call_tool.return_value = mock_mcp_tool_call_result + request_params = CallToolRequestParams( + name="test_tool", arguments={"test": "value"} + ) + result = await call_mcp_tool(mock_session, request_params) + print("call_mcp_tool result", result) + assert result.content[0].text == "test_output" + mock_session.call_tool.assert_called_once_with( + name="test_tool", arguments={"test": "value"} + )