|
| 1 | +#!/usr/bin/env python3 |
| 2 | +""" |
| 3 | +Simple MCP client smoke test for ZenML MCP server |
| 4 | +""" |
| 5 | + |
| 6 | +import asyncio |
| 7 | +import sys |
| 8 | +from pathlib import Path |
| 9 | +from typing import Any, Dict |
| 10 | + |
| 11 | +from mcp import ClientSession, StdioServerParameters |
| 12 | +from mcp.client.stdio import stdio_client |
| 13 | + |
| 14 | + |
| 15 | +class MCPSmokeTest: |
| 16 | + def __init__(self, server_path: str): |
| 17 | + """Initialize the smoke test with the server path.""" |
| 18 | + self.server_path = Path(server_path) |
| 19 | + self.server_params = StdioServerParameters( |
| 20 | + command="uv", |
| 21 | + args=["run", str(self.server_path)], |
| 22 | + ) |
| 23 | + |
| 24 | + async def run_smoke_test(self) -> Dict[str, Any]: |
| 25 | + """Run a comprehensive smoke test of the MCP server.""" |
| 26 | + results = { |
| 27 | + "connection": False, |
| 28 | + "initialization": False, |
| 29 | + "tools": [], |
| 30 | + "resources": [], |
| 31 | + "prompts": [], |
| 32 | + "tool_test_results": {}, |
| 33 | + "errors": [], |
| 34 | + } |
| 35 | + |
| 36 | + try: |
| 37 | + print(f"🚀 Starting smoke test for MCP server: {self.server_path}") |
| 38 | + |
| 39 | + # Connect to the server |
| 40 | + async with stdio_client(self.server_params) as (read, write): |
| 41 | + print("✅ Connected to MCP server") |
| 42 | + results["connection"] = True |
| 43 | + |
| 44 | + async with ClientSession(read, write) as session: |
| 45 | + # Initialize the session |
| 46 | + print("🔄 Initializing session...") |
| 47 | + await asyncio.wait_for(session.initialize(), timeout=60.0) |
| 48 | + print("✅ Session initialized") |
| 49 | + results["initialization"] = True |
| 50 | + |
| 51 | + # List available tools |
| 52 | + print("🔄 Listing available tools...") |
| 53 | + tools_result = await asyncio.wait_for( |
| 54 | + session.list_tools(), timeout=30.0 |
| 55 | + ) |
| 56 | + print( |
| 57 | + f"🔄 Got tools result: {len(tools_result.tools) if tools_result.tools else 0} tools" |
| 58 | + ) |
| 59 | + if tools_result.tools: |
| 60 | + results["tools"] = [ |
| 61 | + {"name": tool.name, "description": tool.description} |
| 62 | + for tool in tools_result.tools |
| 63 | + ] |
| 64 | + print(f"✅ Found {len(tools_result.tools)} tools:") |
| 65 | + for tool in tools_result.tools: |
| 66 | + print(f" - {tool.name}: {tool.description}") |
| 67 | + |
| 68 | + # List available resources |
| 69 | + print("🔄 Listing available resources...") |
| 70 | + try: |
| 71 | + resources_result = await asyncio.wait_for( |
| 72 | + session.list_resources(), timeout=30.0 |
| 73 | + ) |
| 74 | + print( |
| 75 | + f"🔄 Got resources result: {len(resources_result.resources) if resources_result.resources else 0} resources" |
| 76 | + ) |
| 77 | + if resources_result.resources: |
| 78 | + results["resources"] = [ |
| 79 | + { |
| 80 | + "uri": res.uri, |
| 81 | + "name": res.name, |
| 82 | + "description": res.description, |
| 83 | + } |
| 84 | + for res in resources_result.resources |
| 85 | + ] |
| 86 | + print( |
| 87 | + f"✅ Found {len(resources_result.resources)} resources:" |
| 88 | + ) |
| 89 | + for res in resources_result.resources: |
| 90 | + print(f" - {res.name}: {res.description}") |
| 91 | + except Exception as e: |
| 92 | + print( |
| 93 | + f"ℹ️ No resources available or error listing resources: {e}" |
| 94 | + ) |
| 95 | + |
| 96 | + # List available prompts |
| 97 | + print("🔄 Listing available prompts...") |
| 98 | + try: |
| 99 | + prompts_result = await asyncio.wait_for( |
| 100 | + session.list_prompts(), timeout=30.0 |
| 101 | + ) |
| 102 | + print( |
| 103 | + f"🔄 Got prompts result: {len(prompts_result.prompts) if prompts_result.prompts else 0} prompts" |
| 104 | + ) |
| 105 | + if prompts_result.prompts: |
| 106 | + results["prompts"] = [ |
| 107 | + {"name": prompt.name, "description": prompt.description} |
| 108 | + for prompt in prompts_result.prompts |
| 109 | + ] |
| 110 | + print(f"✅ Found {len(prompts_result.prompts)} prompts:") |
| 111 | + for prompt in prompts_result.prompts: |
| 112 | + print(f" - {prompt.name}: {prompt.description}") |
| 113 | + except Exception as e: |
| 114 | + print(f"ℹ️ No prompts available or error listing prompts: {e}") |
| 115 | + |
| 116 | + # Test a few basic tools (if available) |
| 117 | + print("🔄 Starting tool tests...") |
| 118 | + await self._test_basic_tools(session, results) |
| 119 | + print("✅ Tool tests completed") |
| 120 | + |
| 121 | + except Exception as e: |
| 122 | + error_msg = f"❌ Error during smoke test: {e}" |
| 123 | + print(error_msg) |
| 124 | + results["errors"].append(error_msg) |
| 125 | + |
| 126 | + return results |
| 127 | + |
| 128 | + async def _test_basic_tools(self, session: ClientSession, results: Dict[str, Any]): |
| 129 | + """Test basic tools that are likely to be safe to call.""" |
| 130 | + safe_tools_to_test = [ |
| 131 | + "list_users", |
| 132 | + "list_stacks", |
| 133 | + "list_pipelines", |
| 134 | + "get_server_info", |
| 135 | + ] |
| 136 | + |
| 137 | + available_tools = {tool["name"] for tool in results["tools"]} |
| 138 | + print(f"🔄 Available tools for testing: {available_tools}") |
| 139 | + |
| 140 | + for tool_name in safe_tools_to_test: |
| 141 | + if tool_name in available_tools: |
| 142 | + try: |
| 143 | + print(f"🧪 Testing tool: {tool_name}") |
| 144 | + print(f"🔄 Calling tool {tool_name}...") |
| 145 | + # Add timeout to prevent hanging |
| 146 | + result = await asyncio.wait_for( |
| 147 | + session.call_tool(tool_name, {}), timeout=30.0 |
| 148 | + ) |
| 149 | + print(f"🔄 Tool {tool_name} returned result") |
| 150 | + results["tool_test_results"][tool_name] = { |
| 151 | + "success": True, |
| 152 | + "content_length": len(str(result.content)) |
| 153 | + if result.content |
| 154 | + else 0, |
| 155 | + } |
| 156 | + print(f"✅ Tool {tool_name} executed successfully") |
| 157 | + except Exception as e: |
| 158 | + error_msg = f"Tool {tool_name} failed: {e}" |
| 159 | + print(f"❌ {error_msg}") |
| 160 | + results["tool_test_results"][tool_name] = { |
| 161 | + "success": False, |
| 162 | + "error": str(e), |
| 163 | + } |
| 164 | + else: |
| 165 | + print(f"ℹ️ Tool {tool_name} not available in server") |
| 166 | + |
| 167 | + def print_summary(self, results: Dict[str, Any]): |
| 168 | + """Print a summary of the smoke test results.""" |
| 169 | + print("\n" + "=" * 50) |
| 170 | + print("🔍 SMOKE TEST SUMMARY") |
| 171 | + print("=" * 50) |
| 172 | + |
| 173 | + print(f"Connection: {'✅ PASS' if results['connection'] else '❌ FAIL'}") |
| 174 | + print( |
| 175 | + f"Initialization: {'✅ PASS' if results['initialization'] else '❌ FAIL'}" |
| 176 | + ) |
| 177 | + print(f"Tools found: {len(results['tools'])}") |
| 178 | + print(f"Resources found: {len(results['resources'])}") |
| 179 | + print(f"Prompts found: {len(results['prompts'])}") |
| 180 | + |
| 181 | + if results["tool_test_results"]: |
| 182 | + successful_tests = sum( |
| 183 | + 1 for r in results["tool_test_results"].values() if r["success"] |
| 184 | + ) |
| 185 | + total_tests = len(results["tool_test_results"]) |
| 186 | + print(f"Tool tests: {successful_tests}/{total_tests} passed") |
| 187 | + |
| 188 | + if results["errors"]: |
| 189 | + print(f"Errors: {len(results['errors'])}") |
| 190 | + for error in results["errors"]: |
| 191 | + print(f" - {error}") |
| 192 | + |
| 193 | + overall_status = ( |
| 194 | + results["connection"] |
| 195 | + and results["initialization"] |
| 196 | + and len(results["tools"]) > 0 |
| 197 | + ) |
| 198 | + print(f"\nOverall: {'✅ PASS' if overall_status else '❌ FAIL'}") |
| 199 | + |
| 200 | + |
| 201 | +async def main(): |
| 202 | + """Main entry point for the smoke test.""" |
| 203 | + if len(sys.argv) != 2: |
| 204 | + print("Usage: python test_mcp_server.py <path_to_mcp_server.py>") |
| 205 | + print("Example: python test_mcp_server.py ./zenml_server.py") |
| 206 | + sys.exit(1) |
| 207 | + |
| 208 | + server_path = sys.argv[1] |
| 209 | + |
| 210 | + # Verify server file exists |
| 211 | + if not Path(server_path).exists(): |
| 212 | + print(f"❌ Server file not found: {server_path}") |
| 213 | + sys.exit(1) |
| 214 | + |
| 215 | + smoke_test = MCPSmokeTest(server_path) |
| 216 | + results = await smoke_test.run_smoke_test() |
| 217 | + smoke_test.print_summary(results) |
| 218 | + |
| 219 | + # Exit with appropriate code |
| 220 | + if results["connection"] and results["initialization"]: |
| 221 | + sys.exit(0) |
| 222 | + else: |
| 223 | + sys.exit(1) |
| 224 | + |
| 225 | + |
| 226 | +if __name__ == "__main__": |
| 227 | + asyncio.run(main()) |
0 commit comments