206 lines
9 KiB
Python
206 lines
9 KiB
Python
"""End-to-end tests with real subprocesses: worker pool, backpressure, HTTP and stdio transport."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import signal
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import unittest
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
import httpx2
|
|
from mcp import Client, StdioServerParameters
|
|
from mcp.client.streamable_http import streamable_http_client
|
|
|
|
from mcp_calculator.config import load_settings
|
|
from mcp_calculator.executor import ComputationError, Executor
|
|
from tests.support import slow
|
|
|
|
TOKEN = "test-token-0123456789"
|
|
|
|
|
|
def _free_port() -> int:
|
|
with socket.socket() as s:
|
|
s.bind(("127.0.0.1", 0))
|
|
return s.getsockname()[1]
|
|
|
|
|
|
def _http(method: str, url: str, body: dict | None = None, token: str | None = None) -> tuple[int, str]:
|
|
headers = {"Content-Type": "application/json", "Accept": "application/json, text/event-stream"}
|
|
if token:
|
|
headers["Authorization"] = f"Bearer {token}"
|
|
data = json.dumps(body).encode() if body is not None else None
|
|
try:
|
|
with urllib.request.urlopen(urllib.request.Request(url, data=data, headers=headers, method=method), timeout=30) as r:
|
|
return r.status, r.read().decode()
|
|
except urllib.error.HTTPError as exc:
|
|
return exc.code, exc.read().decode()
|
|
|
|
|
|
def _start_server(port: int, **env: str) -> subprocess.Popen:
|
|
full_env = {**os.environ, "MCP_CALC_TRANSPORT": "http", "MCP_CALC_PORT": str(port), "MCP_CALC_WORKERS": "1", **env}
|
|
proc = subprocess.Popen([sys.executable, "-m", "mcp_calculator"], env=full_env, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
|
|
# Poll readiness (503 until the workers are warm). Always pause between polls: every request
|
|
# is access-logged, and a flood of log lines would fill the stderr pipe and block the server.
|
|
deadline = time.time() + 30
|
|
while time.time() < deadline:
|
|
try:
|
|
if _http("GET", f"http://127.0.0.1:{port}/readyz")[0] == 200:
|
|
return proc
|
|
except OSError:
|
|
pass
|
|
time.sleep(0.2)
|
|
proc.kill()
|
|
raise RuntimeError(f"server did not start: {proc.communicate()[1]}")
|
|
|
|
|
|
# --------------------------------------------------------------------------- worker pool
|
|
|
|
|
|
@slow
|
|
class WorkerPoolTest(unittest.IsolatedAsyncioTestCase):
|
|
def setUp(self) -> None:
|
|
self.executor = Executor(load_settings([], env={"MCP_CALC_TIMEOUT_SECONDS": "2", "MCP_CALC_WORKERS": "1"}))
|
|
self.executor.start()
|
|
self.addCleanup(self.executor.stop)
|
|
|
|
async def test_computes(self):
|
|
self.assertEqual((await self.executor.run("algebra.evaluate", expression="2^10"))["exact"], "1024")
|
|
|
|
async def test_timeout_kills_worker_and_recovers(self):
|
|
started = time.perf_counter()
|
|
with self.assertRaisesRegex(ComputationError, "timeout") as ctx:
|
|
await self.executor.run("algebra.evaluate", expression="9^9^9^9")
|
|
self.assertEqual(ctx.exception.kind, "timeout")
|
|
self.assertLess(time.perf_counter() - started, 5)
|
|
# the killed worker is replaced transparently
|
|
self.assertEqual((await self.executor.run("algebra.evaluate", expression="1+1"))["exact"], "2")
|
|
self.assertTrue(self.executor.healthy)
|
|
|
|
async def test_ready_only_after_warm_up(self):
|
|
executor = Executor(load_settings([], env={"MCP_CALC_WORKERS": "1"}))
|
|
executor.start()
|
|
self.addCleanup(executor.stop)
|
|
self.assertFalse(executor.healthy, "workers cannot have loaded SymPy/SciPy yet")
|
|
deadline = time.time() + 30
|
|
while not executor.healthy and time.time() < deadline:
|
|
await asyncio.sleep(0.05)
|
|
self.assertTrue(executor.healthy)
|
|
|
|
async def test_math_error(self):
|
|
with self.assertRaisesRegex(ComputationError, "unsupported character") as ctx:
|
|
await self.executor.run("algebra.evaluate", expression="'x'")
|
|
self.assertEqual(ctx.exception.kind, "error")
|
|
|
|
|
|
@slow
|
|
class BackpressureTest(unittest.IsolatedAsyncioTestCase):
|
|
async def test_requests_beyond_capacity_are_rejected_immediately(self):
|
|
env = {"MCP_CALC_TIMEOUT_SECONDS": "1.5", "MCP_CALC_WORKERS": "1", "MCP_CALC_MAX_QUEUE": "1"}
|
|
executor = Executor(load_settings([], env=env))
|
|
executor.start()
|
|
self.addCleanup(executor.stop)
|
|
self.assertEqual(executor.capacity, 2)
|
|
|
|
slow_tasks = [asyncio.create_task(executor.run("algebra.evaluate", expression="9^9^9^9")) for _ in range(2)]
|
|
await asyncio.sleep(0.1)
|
|
self.assertEqual(executor.inflight, 2)
|
|
|
|
started = time.perf_counter()
|
|
with self.assertRaisesRegex(ComputationError, "busy") as ctx:
|
|
await executor.run("algebra.evaluate", expression="1+1")
|
|
self.assertEqual(ctx.exception.kind, "busy")
|
|
self.assertLess(time.perf_counter() - started, 0.5, "busy rejection must not wait")
|
|
|
|
results = await asyncio.gather(*slow_tasks, return_exceptions=True)
|
|
self.assertTrue(all(isinstance(r, ComputationError) and r.kind == "timeout" for r in results), results)
|
|
self.assertEqual(executor.inflight, 0)
|
|
# capacity is free again
|
|
self.assertEqual((await executor.run("algebra.evaluate", expression="1+1"))["exact"], "2")
|
|
|
|
|
|
# --------------------------------------------------------------------------- HTTP
|
|
|
|
|
|
@slow
|
|
class HttpTransportTest(unittest.IsolatedAsyncioTestCase):
|
|
@classmethod
|
|
def setUpClass(cls) -> None:
|
|
port = _free_port()
|
|
cls.base = f"http://127.0.0.1:{port}"
|
|
cls.proc = _start_server(port, MCP_CALC_AUTH_TOKENS=f"tests:{TOKEN}", MCP_CALC_LOG_LEVEL="WARNING")
|
|
|
|
@classmethod
|
|
def tearDownClass(cls) -> None:
|
|
cls.proc.terminate()
|
|
cls.proc.communicate(timeout=30)
|
|
|
|
def test_health_endpoints_are_public(self):
|
|
self.assertEqual(json.loads(_http("GET", f"{self.base}/healthz")[1])["status"], "ok")
|
|
self.assertEqual(json.loads(_http("GET", f"{self.base}/readyz")[1])["status"], "ready")
|
|
|
|
def test_mcp_endpoint_requires_token(self):
|
|
self.assertEqual(_http("POST", f"{self.base}/mcp", {"jsonrpc": "2.0", "id": 1, "method": "tools/list"})[0], 401)
|
|
self.assertEqual(_http("POST", f"{self.base}/mcp", {}, token="wrong")[0], 401)
|
|
|
|
async def test_mcp_over_http(self):
|
|
async with httpx2.AsyncClient(headers={"Authorization": f"Bearer {TOKEN}"}, timeout=30) as http:
|
|
async with Client(streamable_http_client(f"{self.base}/mcp", http_client=http)) as client:
|
|
self.assertEqual(len((await client.list_tools()).tools), 13)
|
|
result = await client.call_tool("evaluate", {"expression": "sqrt(2)", "precision": 30})
|
|
self.assertEqual(result.structured_content["numeric"], "1.41421356237309504880168872421")
|
|
|
|
def test_metrics_endpoint(self):
|
|
body = {"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "convert_units", "arguments": {"quantity": "1 m", "to_unit": "cm"}}}
|
|
self.assertEqual(_http("POST", f"{self.base}/mcp", body, token=TOKEN)[0], 200)
|
|
status, text = _http("GET", f"{self.base}/metrics") # public, like the probes
|
|
self.assertEqual(status, 200)
|
|
self.assertIn('mcp_calc_tool_calls_total{client="tests",status="ok",tool="convert_units"}', text)
|
|
for name in ("mcp_calc_tool_duration_seconds_bucket", "mcp_calc_inflight_computations", "mcp_calc_capacity 21.0", "mcp_calc_build_info"):
|
|
self.assertIn(name, text)
|
|
|
|
|
|
@slow
|
|
class HttpLifecycleTest(unittest.TestCase):
|
|
def test_metrics_can_be_disabled(self):
|
|
port = _free_port()
|
|
proc = _start_server(port, MCP_CALC_METRICS="false")
|
|
try:
|
|
self.assertEqual(_http("GET", f"http://127.0.0.1:{port}/metrics")[0], 404)
|
|
finally:
|
|
proc.terminate()
|
|
proc.communicate(timeout=30)
|
|
|
|
def test_sigterm_stops_worker_pool_cleanly(self):
|
|
port = _free_port()
|
|
proc = _start_server(port, MCP_CALC_LOG_LEVEL="INFO")
|
|
body = {"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "evaluate", "arguments": {"expression": "1+1"}}}
|
|
self.assertEqual(_http("POST", f"http://127.0.0.1:{port}/mcp", body)[0], 200)
|
|
proc.send_signal(signal.SIGTERM)
|
|
# communicate() also waits for the multiprocessing resource tracker, which shares stderr
|
|
# and reports leaked semaphores only when it exits.
|
|
_, stderr = proc.communicate(timeout=30)
|
|
self.assertIn("worker pool stopped", stderr)
|
|
self.assertNotIn("leaked", stderr)
|
|
|
|
|
|
# --------------------------------------------------------------------------- stdio
|
|
|
|
|
|
@slow
|
|
class StdioTransportTest(unittest.IsolatedAsyncioTestCase):
|
|
async def test_mcp_over_stdio(self):
|
|
params = StdioServerParameters(
|
|
command=sys.executable,
|
|
args=["-m", "mcp_calculator"],
|
|
env={**os.environ, "MCP_CALC_WORKERS": "1", "MCP_CALC_LOG_LEVEL": "WARNING"},
|
|
)
|
|
async with Client(params) as client:
|
|
result = await client.call_tool("solve", {"equations": ["x + y = 3", "x - y = 1"]})
|
|
self.assertEqual(result.structured_content["solutions"], [{"x": "2", "y": "1"}])
|