test_server.py

#!/usr/bin/env python3
"""
FORUS Cortex Agent MCP Server - Test Suite
"""

import asyncio
import json
import sys
import time
import requests
from typing import Dict, Any

# Test configuration
SERVER_URL = "http://localhost:8083"
MCP_ENDPOINT = f"{SERVER_URL}/mcp"

class MCPTester:
    """Test suite for FORUS MCP Server"""
    
    def __init__(self):
        self.passed = 0
        self.failed = 0
        self.request_id = 0
    
    def get_next_id(self):
        self.request_id += 1
        return self.request_id
    
    def test_result(self, test_name: str, success: bool, message: str = ""):
        """Record test result"""
        if success:
            print(f"✅ {test_name}")
            self.passed += 1
        else:
            print(f"❌ {test_name}: {message}")
            self.failed += 1
    
    def make_mcp_request(self, method: str, params: Dict = None) -> Dict[str, Any]:
        """Make MCP JSON-RPC request"""
        request = {
            "jsonrpc": "2.0",
            "id": self.get_next_id(),
            "method": method
        }
        
        if params:
            request["params"] = params
        
        try:
            response = requests.post(MCP_ENDPOINT, json=request, timeout=10)
            return response.json()
        except Exception as e:
            return {"error": {"code": -32603, "message": str(e)}}
    
    def test_server_info(self):
        """Test server info endpoint"""
        try:
            response = requests.get(SERVER_URL, timeout=5)
            if response.status_code == 200:
                data = response.json()
                if "server" in data and "FORUS" in data["server"]:
                    self.test_result("Server Info Endpoint", True)
                    return True
                else:
                    self.test_result("Server Info Endpoint", False, "Invalid response format")
                    return False
            else:
                self.test_result("Server Info Endpoint", False, f"HTTP {response.status_code}")
                return False
        except Exception as e:
            self.test_result("Server Info Endpoint", False, str(e))
            return False
    
    def test_tools_list(self):
        """Test tools/list method"""
        response = self.make_mcp_request("tools/list")
        
        if "error" in response:
            self.test_result("Tools List", False, response["error"]["message"])
            return False
        
        if "result" in response and "tools" in response["result"]:
            tools = response["result"]["tools"]
            if len(tools) >= 6:  # Should have at least 6 tools
                self.test_result("Tools List", True)
                return True
            else:
                self.test_result("Tools List", False, f"Only {len(tools)} tools found")
                return False
        else:
            self.test_result("Tools List", False, "Invalid response format")
            return False
    
    def test_initialize(self):
        """Test initialize method"""
        response = self.make_mcp_request("initialize", {
            "protocolVersion": "2024-11-05",
            "capabilities": {},
            "clientInfo": {"name": "test-client", "version": "1.0.0"}
        })
        
        if "error" in response:
            self.test_result("Initialize", False, response["error"]["message"])
            return False
        
        if "result" in response:
            result = response["result"]
            if "protocolVersion" in result and "capabilities" in result:
                self.test_result("Initialize", True)
                return True
            else:
                self.test_result("Initialize", False, "Invalid response format")
                return False
        else:
            self.test_result("Initialize", False, "No result in response")
            return False
    
    def test_memory_search(self):
        """Test memory_search tool"""
        response = self.make_mcp_request("tools/call", {
            "name": "memory_search",
            "arguments": {
                "query": "test query",
                "limit": 5
            }
        })
        
        if "error" in response:
            self.test_result("Memory Search Tool", False, response["error"]["message"])
            return False
        
        if "result" in response and "content" in response["result"]:
            self.test_result("Memory Search Tool", True)
            return True
        else:
            self.test_result("Memory Search Tool", False, "Invalid response format")
            return False
    
    def test_python_execution(self):
        """Test execute_python tool"""
        code = "print('Hello from test!'); result = 2 + 2; print(f'2 + 2 = {result}')"
        response = self.make_mcp_request("tools/call", {
            "name": "execute_python",
            "arguments": {
                "code": code,
                "timeout": 10
            }
        })
        
        if "error" in response:
            self.test_result("Python Execution Tool", False, response["error"]["message"])
            return False
        
        if "result" in response and "content" in response["result"]:
            # Check if execution was successful
            content = response["result"]["content"][0]["text"]
            try:
                result_data = json.loads(content)
                if result_data.get("success") and "Hello from test!" in result_data.get("stdout", ""):
                    self.test_result("Python Execution Tool", True)
                    return True
                else:
                    self.test_result("Python Execution Tool", False, "Execution failed or unexpected output")
                    return False
            except json.JSONDecodeError:
                self.test_result("Python Execution Tool", False, "Invalid JSON response")
                return False
        else:
            self.test_result("Python Execution Tool", False, "Invalid response format")
            return False
    
    def test_shell_execution(self):
        """Test execute_shell tool"""
        response = self.make_mcp_request("tools/call", {
            "name": "execute_shell",
            "arguments": {
                "command": "echo 'Test successful'",
                "timeout": 5
            }
        })
        
        if "error" in response:
            self.test_result("Shell Execution Tool", False, response["error"]["message"])
            return False
        
        if "result" in response and "content" in response["result"]:
            content = response["result"]["content"][0]["text"]
            try:
                result_data = json.loads(content)
                if result_data.get("success") and "Test successful" in result_data.get("stdout", ""):
                    self.test_result("Shell Execution Tool", True)
                    return True
                else:
                    self.test_result("Shell Execution Tool", False, "Command failed or unexpected output")
                    return False
            except json.JSONDecodeError:
                self.test_result("Shell Execution Tool", False, "Invalid JSON response")
                return False
        else:
            self.test_result("Shell Execution Tool", False, "Invalid response format")
            return False
    
    def test_system_status(self):
        """Test get_system_status tool"""
        response = self.make_mcp_request("tools/call", {
            "name": "get_system_status",
            "arguments": {}
        })
        
        if "error" in response:
            self.test_result("System Status Tool", False, response["error"]["message"])
            return False
        
        if "result" in response and "content" in response["result"]:
            content = response["result"]["content"][0]["text"]
            try:
                result_data = json.loads(content)
                if result_data.get("success") and "server" in result_data:
                    self.test_result("System Status Tool", True)
                    return True
                else:
                    self.test_result("System Status Tool", False, "Invalid status data")
                    return False
            except json.JSONDecodeError:
                self.test_result("System Status Tool", False, "Invalid JSON response")
                return False
        else:
            self.test_result("System Status Tool", False, "Invalid response format")
            return False
    
    def test_invalid_method(self):
        """Test invalid method handling"""
        response = self.make_mcp_request("invalid/method")
        
        if "error" in response and response["error"]["code"] == -32601:
            self.test_result("Invalid Method Handling", True)
            return True
        else:
            self.test_result("Invalid Method Handling", False, "Should return method not found error")
            return False
    
    def test_invalid_tool(self):
        """Test invalid tool handling"""
        response = self.make_mcp_request("tools/call", {
            "name": "non_existent_tool",
            "arguments": {}
        })
        
        if "error" in response and response["error"]["code"] == -32601:
            self.test_result("Invalid Tool Handling", True)
            return True
        else:
            self.test_result("Invalid Tool Handling", False, "Should return tool not found error")
            return False
    
    def run_performance_test(self):
        """Test server performance"""
        print(f"\n🚀 Running performance tests...")
        
        # Test response time
        start_time = time.time()
        response = self.make_mcp_request("tools/list")
        end_time = time.time()
        
        response_time = (end_time - start_time) * 1000  # Convert to milliseconds
        
        if response_time < 500:  # Should respond within 500ms
            self.test_result(f"Response Time ({response_time:.1f}ms)", True)
        else:
            self.test_result(f"Response Time ({response_time:.1f}ms)", False, "Too slow")
        
        # Test concurrent requests
        import threading
        import queue
        
        results_queue = queue.Queue()
        
        def make_concurrent_request():
            try:
                response = self.make_mcp_request("get_system_status")
                results_queue.put("error" not in response)
            except:
                results_queue.put(False)
        
        threads = []
        for _ in range(10):  # 10 concurrent requests
            thread = threading.Thread(target=make_concurrent_request)
            threads.append(thread)
            thread.start()
        
        for thread in threads:
            thread.join()
        
        successful_requests = sum(results_queue.get() for _ in range(10))
        
        if successful_requests >= 8:  # At least 80% should succeed
            self.test_result(f"Concurrent Requests ({successful_requests}/10)", True)
        else:
            self.test_result(f"Concurrent Requests ({successful_requests}/10)", False, "Too many failed")
    
    def run_all_tests(self):
        """Run all tests"""
        print("🧪 FORUS Cortex Agent MCP Server Test Suite")
        print("=" * 50)
        
        # Basic connectivity tests
        print("\n📡 Testing Basic Connectivity...")
        self.test_server_info()
        
        # MCP protocol tests
        print("\n🔧 Testing MCP Protocol...")
        self.test_initialize()
        self.test_tools_list()
        
        # Tool functionality tests  
        print("\n🛠️ Testing Tool Functions...")
        self.test_memory_search()
        self.test_python_execution()
        self.test_shell_execution()
        self.test_system_status()
        
        # Error handling tests
        print("\n⚠️ Testing Error Handling...")
        self.test_invalid_method()
        self.test_invalid_tool()
        
        # Performance tests
        self.run_performance_test()
        
        # Summary
        print("\n" + "=" * 50)
        print("📊 Test Summary")
        print("=" * 50)
        total_tests = self.passed + self.failed
        success_rate = (self.passed / total_tests) * 100 if total_tests > 0 else 0
        
        print(f"✅ Passed: {self.passed}")
        print(f"❌ Failed: {self.failed}")
        print(f"📈 Success Rate: {success_rate:.1f}%")
        
        if self.failed == 0:
            print("\n🎉 All tests passed! MCP Server is working correctly.")
            return True
        else:
            print(f"\n⚠️ {self.failed} test(s) failed. Check server configuration.")
            return False

def main():
    """Main test runner"""
    print("Starting FORUS MCP Server tests...")
    print(f"Server URL: {SERVER_URL}")
    print(f"MCP Endpoint: {MCP_ENDPOINT}")
    
    tester = MCPTester()
    success = tester.run_all_tests()
    
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()