import asyncio
import sys
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from mcp.client.sse import sse_client
async def test_stdio():
print("🧪 Testing Stdio Transport...")
server_params = StdioServerParameters(
command="python",
args=["src/main.py"],
env=None # Inherit env from current process (assuming .env is loaded or env vars set)
)
async with stdio_client(server_params) as (read, write):
async with ClientSession(read, write) as session:
await run_tests(session)
async def test_sse():
print("🧪 Testing SSE Transport (http://localhost:8000/sse)...")
async with sse_client("http://localhost:8000/sse") as (read, write):
async with ClientSession(read, write) as session:
await run_tests(session)
async def run_tests(session: ClientSession):
print("✅ Connected to MCP Server!")
# List Tools
print("📋 Listing Tools...")
tools = await session.list_tools()
print(f" Found {len(tools.tools)} tools: {[t.name for t in tools.tools]}")
# Test: Add Memory
print("💾 Testing add_memory...")
result = await session.call_tool("add_memory", arguments={
"content": "The user is testing the Mem0 server integration.",
"user_id": "test_user_123",
"metadata": {"source": "test_script"}
})
print(f" Result: {result.content[0].text[:100]}...") # Truncate for brevity
# Test: Search Memory
print("🔍 Testing search_memories...")
search_res = await session.call_tool("search_memories", arguments={
"query": "testing server",
"user_id": "test_user_123"
})
print(f" Result: {search_res.content[0].text[:100]}...")
print("🎉 All Tests Passed!")
if __name__ == "__main__":
mode = sys.argv[1] if len(sys.argv) > 1 else "stdio"
try:
if mode == "sse":
asyncio.run(test_sse())
else:
asyncio.run(test_stdio())
except Exception as e:
print(f"❌ Test Failed: {e}")
sys.exit(1)
