test: rewrite MCP integration tests from user perspective

tests/tools/test_mcp_integration.py

@@ -1,75 +1,136 @@
 
 import pytest
-from swarms.tools.mcp_integration import MCPServerSseParams
+import asyncio
-from swarms import Agent
+from swarms.tools.mcp_integration import (
-from swarms.prompts.finance_agent_sys_prompt import FINANCIAL_AGENT_SYS_PROMPT
+    MCPServer,
+    MCPServerStdio,
+    MCPServerSse,
+    mcp_flow,
+    mcp_flow_get_tool_schema,
+    batch_mcp_flow
+)
 
-def test_interactive_multi_agent_mcp():
+# Test basic server connectivity
-    # Configure two MCP servers
+def test_server_connection():
-    server_one = MCPServerSseParams(
+    """
-        url="http://0.0.0.0:6274",
+    Test that a user can connect to the MCP server successfully
-        headers={"Content-Type": "application/json"}
+    """
-    )
+    params = {"url": "http://localhost:8000"}
+    server = MCPServerSse(params, cache_tools_list=True)
     
-    server_two = MCPServerSseParams(
+    # Connect should work
-        url="http://0.0.0.0:6275", 
+    asyncio.run(server.connect())
-        headers={"Content-Type": "application/json"}
+    assert server.session is not None
-    )
+    
-
+    # Cleanup should work
-    # Create two agents with different roles
+    asyncio.run(server.cleanup())
-    finance_agent = Agent(
+    assert server.session is None
-        agent_name="Finance-Agent",
-        agent_description="Financial analysis expert",
-        system_prompt=FINANCIAL_AGENT_SYS_PROMPT,
-        max_loops=1,
-        mcp_servers=[server_one],
-        interactive=True,
-        streaming_on=True
-    )
 
-    research_agent = Agent(
+# Test tool listing functionality
-        agent_name="Research-Agent", 
+def test_list_tools():
-        agent_description="Market research specialist",
+    """
-        system_prompt="You are a market research specialist. Analyze market trends and provide insights.",
+    Test that a user can retrieve available tools from the server
-        max_loops=1,
+    """
-        mcp_servers=[server_two],
+    params = {"url": "http://localhost:8000"}
-        interactive=True,
+    server = MCPServerSse(params)
-        streaming_on=True
+    
-    )
+    asyncio.run(server.connect())
+    tools = asyncio.run(server.list_tools())
+    
+    assert isinstance(tools, list)
+    assert len(tools) > 0
+    
+    asyncio.run(server.cleanup())
 
-    try:
+# Test tool execution
-        # Interactive loop
+def test_tool_execution():
-        while True:
+    """
-            # Get user input for which agent to use
+    Test that a user can execute a tool successfully
-            print("\nWhich agent would you like to interact with?")
+    """
-            print("1. Finance Agent")
+    params = {"url": "http://localhost:8000"}
-            print("2. Research Agent")
+    function_call = {
-            print("3. Exit")
+        "tool_name": "add",
-            
+        "arguments": {"a": 5, "b": 3}
-            choice = input("Enter your choice (1-3): ")
+    }
-            
+    
-            if choice == "3":
+    result = mcp_flow(params, function_call)
-                break
+    assert result is not None
-                
-            # Get the task from user
-            task = input("\nEnter your task for the agent: ")
-            
-            # Route to appropriate agent
-            if choice == "1":
-                response = finance_agent.run(task)
-                print(f"\nFinance Agent Response:\n{response}")
-            elif choice == "2":
-                response = research_agent.run(task)
-                print(f"\nResearch Agent Response:\n{response}")
-            else:
-                print("Invalid choice, please try again")
 
-    except Exception as e:
+# Test batch operations
-        pytest.fail(f"Interactive multi-agent test failed: {e}")
+def test_batch_execution():
+    """
+    Test that a user can execute multiple tools in batch
+    """
+    params_list = [
+        {"url": "http://localhost:8000"},
+        {"url": "http://localhost:8000"}
+    ]
+    function_calls = [
+        {"tool_name": "add", "arguments": {"a": 1, "b": 2}},
+        {"tool_name": "subtract", "arguments": {"a": 5, "b": 3}}
+    ]
+    
+    results = batch_mcp_flow(params_list, function_calls)
+    assert len(results) == 2
+    assert all(result is not None for result in results)
 
-def test_mcp_invalid_params():
+# Test error handling
+def test_error_handling():
+    """
+    Test that users receive proper error messages for invalid operations
+    """
+    params = {"url": "http://localhost:8000"}
+    invalid_function = {
+        "tool_name": "nonexistent_tool",
+        "arguments": {}
+    }
+    
     with pytest.raises(Exception):
-        mcp_flow(None, {})
+        mcp_flow(params, invalid_function)
 
-if __name__ == "__main__":
+# Test tool schema retrieval
-    test_interactive_multi_agent_mcp()
+def test_get_tool_schema():
+    """
+    Test that users can retrieve tool schemas correctly
+    """
+    params = {"url": "http://localhost:8000"}
+    schema = mcp_flow_get_tool_schema(params)
+    
+    assert isinstance(schema, dict)
+    assert "tools" in schema or "functions" in schema
+
+# Test server reconnection
+def test_server_reconnection():
+    """
+    Test that users can reconnect to the server after disconnection
+    """
+    params = {"url": "http://localhost:8000"}
+    server = MCPServerSse(params)
+    
+    # First connection
+    asyncio.run(server.connect())
+    asyncio.run(server.cleanup())
+    
+    # Second connection should work
+    asyncio.run(server.connect())
+    assert server.session is not None
+    asyncio.run(server.cleanup())
+
+# Test cache functionality
+def test_cache_behavior():
+    """
+    Test that tool caching works as expected for users
+    """
+    params = {"url": "http://localhost:8000"}
+    server = MCPServerSse(params, cache_tools_list=True)
+    
+    asyncio.run(server.connect())
+    
+    # First call should cache
+    tools1 = asyncio.run(server.list_tools())
+    # Second call should use cache
+    tools2 = asyncio.run(server.list_tools())
+    
+    assert tools1 == tools2
+    
+    asyncio.run(server.cleanup())