diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 17abd1d6a5..cf64d16029 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -7,21 +7,24 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -from litellm.proxy._experimental.mcp_server.mcp_client_manager import ( +from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, MCPSSEServer, ) -MCP_SERVERS = [ - MCPSSEServer(name="zapier_mcp_server", url=os.environ.get("ZAPIER_MCP_SERVER_URL")), -] - -mcp_server_manager = MCPServerManager(mcp_servers=MCP_SERVERS) +mcp_server_manager = MCPServerManager() @pytest.mark.asyncio async def test_mcp_server_manager(): + mcp_server_manager.load_servers_from_config( + { + "zapier_mcp_server": { + "url": os.environ.get("ZAPIER_MCP_SERVER_URL"), + } + } + ) tools = await mcp_server_manager.list_tools() print("TOOLS FROM MCP SERVER MANAGER== ", tools) @@ -29,9 +32,3 @@ async def test_mcp_server_manager(): name="gmail_send_email", arguments={"body": "Test"} ) print("RESULT FROM CALLING TOOL FROM MCP SERVER MANAGER== ", result) - - -""" -TODO test with multiple MCP servers and calling a specific - -"""