159 lines
7.4 KiB
Python
159 lines
7.4 KiB
Python
"""End-to-end: authentication with a clients file, hot reload, admin API, per-client limit,
|
|
cancel on disconnect and the connection limit, against a real server process."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import http.client
|
|
import io
|
|
import json
|
|
import re
|
|
import socket
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from mcp_calculator import admin_cli
|
|
from mcp_calculator.admin import AuditLog, ClientAdmin
|
|
from mcp_calculator.clients import ClientFile
|
|
from tests.support import slow
|
|
from tests.test_transports import _free_port, _http, _start_server
|
|
|
|
ADMIN_TOKEN = "admin-token-0123456789"
|
|
|
|
|
|
def _call(port: int, token: str | None, expression: str = "6*7", *, header: str = "Authorization") -> tuple[int, dict | str]:
|
|
body = {"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "evaluate", "arguments": {"expression": expression}}}
|
|
if header == "Authorization" or token is None:
|
|
status, text = _http("POST", f"http://127.0.0.1:{port}/mcp", body, token=token)
|
|
else: # X-API-Key
|
|
conn = http.client.HTTPConnection("127.0.0.1", port, timeout=30)
|
|
conn.request("POST", "/mcp", json.dumps(body), {"Content-Type": "application/json", "Accept": "application/json, text/event-stream", header: token})
|
|
response = conn.getresponse()
|
|
status, text = response.status, response.read().decode()
|
|
conn.close()
|
|
try:
|
|
return status, json.loads(text)
|
|
except json.JSONDecodeError:
|
|
return status, text
|
|
|
|
|
|
def _metric(port: int, pattern: str) -> float:
|
|
text = _http("GET", f"http://127.0.0.1:{port}/metrics")[1]
|
|
match = re.search(pattern + r" ([0-9.e+]+)", text)
|
|
return float(match.group(1)) if match else 0.0
|
|
|
|
|
|
def _wait_for(predicate, timeout: float = 5.0) -> bool:
|
|
deadline = time.time() + timeout
|
|
while time.time() < deadline:
|
|
if predicate():
|
|
return True
|
|
time.sleep(0.1)
|
|
return False
|
|
|
|
|
|
@slow
|
|
class ClientsFileAuthTest(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls) -> None:
|
|
cls.tmp = tempfile.TemporaryDirectory()
|
|
cls.clients_path = Path(cls.tmp.name) / "clients.json"
|
|
admin = ClientAdmin(ClientFile(cls.clients_path), AuditLog(Path(cls.tmp.name) / "audit.log"))
|
|
_, cls.token = admin.create("team-a", actor="test")
|
|
_, cls.single = admin.create("single", actor="test", max_concurrent=1)
|
|
cls.port, cls.admin_port = _free_port(), _free_port()
|
|
cls.proc = _start_server(
|
|
cls.port,
|
|
MCP_CALC_CLIENTS_FILE=str(cls.clients_path),
|
|
MCP_CALC_CLIENTS_RELOAD_SECONDS="0.2",
|
|
MCP_CALC_ADMIN_TOKENS=f"ops:{ADMIN_TOKEN}",
|
|
MCP_CALC_ADMIN_PORT=str(cls.admin_port),
|
|
MCP_CALC_TIMEOUT_SECONDS="20",
|
|
MCP_CALC_LOG_LEVEL="WARNING",
|
|
)
|
|
|
|
@classmethod
|
|
def tearDownClass(cls) -> None:
|
|
cls.proc.terminate()
|
|
cls.proc.communicate(timeout=60)
|
|
cls.tmp.cleanup()
|
|
|
|
def cli(self, *args: str) -> None:
|
|
with contextlib.redirect_stdout(io.StringIO()):
|
|
self.assertEqual(admin_cli.main(["--file", str(self.clients_path), *args]), 0)
|
|
|
|
def admin(self, method: str, path: str, body: dict | None = None, token: str = ADMIN_TOKEN) -> tuple[int, dict]:
|
|
status, text = _http(method, f"http://127.0.0.1:{self.admin_port}{path}", body, token=token)
|
|
return status, json.loads(text) if text else {}
|
|
|
|
def test_token_header_variants(self):
|
|
status, body = _call(self.port, self.token)
|
|
self.assertEqual(status, 200)
|
|
self.assertIn("42", body["result"]["content"][0]["text"])
|
|
self.assertEqual(_call(self.port, self.token, header="X-API-Key")[0], 200)
|
|
self.assertEqual(_call(self.port, None)[0], 401)
|
|
self.assertEqual(_call(self.port, "mcpc_wrong")[0], 401)
|
|
|
|
def test_disable_via_cli_takes_effect_without_restart(self):
|
|
_, token = ClientAdmin(ClientFile(self.clients_path), AuditLog(None)).create("temp-cli", actor="test")
|
|
self.assertTrue(_wait_for(lambda: _call(self.port, token)[0] == 200))
|
|
self.cli("disable", "temp-cli")
|
|
self.assertTrue(_wait_for(lambda: _call(self.port, token)[0] == 403), "disabled client must be rejected after reload")
|
|
self.cli("enable", "temp-cli")
|
|
self.assertTrue(_wait_for(lambda: _call(self.port, token)[0] == 200))
|
|
|
|
def test_admin_api_over_http(self):
|
|
self.assertEqual(self.admin("GET", "/admin/clients", token="wrong")[0], 401)
|
|
status, body = self.admin("POST", "/admin/clients", {"name": "temp-api", "max_concurrent": 2})
|
|
self.assertEqual(status, 201)
|
|
token = body["token"]
|
|
self.assertEqual(_call(self.port, token)[0], 200, "the replica serving the admin call applies it immediately")
|
|
status, body = self.admin("POST", "/admin/clients/temp-api/rotate")
|
|
self.assertEqual(_call(self.port, token)[0], 401, "old token is invalid after rotation")
|
|
self.assertEqual(_call(self.port, body["token"])[0], 200)
|
|
self.assertEqual(self.admin("DELETE", "/admin/clients/temp-api")[0], 200)
|
|
self.assertEqual(_call(self.port, body["token"])[0], 401)
|
|
|
|
def test_per_client_concurrency_limit(self):
|
|
blocker = threading.Thread(target=_call, args=(self.port, self.single, "9^9^9^9"), daemon=True)
|
|
blocker.start()
|
|
self.assertTrue(_wait_for(lambda: _metric(self.port, r'mcp_calc_client_inflight_computations\{client="single"\}') == 1))
|
|
status, body = _call(self.port, self.single)
|
|
self.assertEqual(status, 200)
|
|
self.assertTrue(body["result"]["isError"])
|
|
self.assertIn("client busy", body["result"]["content"][0]["text"])
|
|
self.assertEqual(_call(self.port, self.token)[0], 200, "other clients are not affected")
|
|
|
|
def test_disconnect_cancels_computation(self):
|
|
before = _metric(self.port, r'mcp_calc_tool_calls_total\{client="team-a",status="cancelled",tool="evaluate"\}')
|
|
body = {"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "evaluate", "arguments": {"expression": "9^9^9^9"}}}
|
|
conn = http.client.HTTPConnection("127.0.0.1", self.port, timeout=0.5)
|
|
conn.request("POST", "/mcp", json.dumps(body), {"Content-Type": "application/json", "Accept": "application/json, text/event-stream", "Authorization": f"Bearer {self.token}"})
|
|
with self.assertRaises((TimeoutError, socket.timeout)):
|
|
conn.getresponse()
|
|
conn.close()
|
|
pattern = r'mcp_calc_tool_calls_total\{client="team-a",status="cancelled",tool="evaluate"\}'
|
|
self.assertTrue(_wait_for(lambda: _metric(self.port, pattern) == before + 1), "computation must be cancelled well before its 20 s timeout")
|
|
|
|
|
|
@slow
|
|
class ConnectionLimitTest(unittest.TestCase):
|
|
def test_connections_beyond_the_limit_get_503(self):
|
|
port = _free_port()
|
|
# uvicorn counts the incoming connection itself: with a limit of 2, one long request
|
|
# is served and the next concurrent connection is rejected.
|
|
proc = _start_server(port, MCP_CALC_MAX_CONNECTIONS="2", MCP_CALC_TIMEOUT_SECONDS="3", MCP_CALC_LOG_LEVEL="WARNING")
|
|
try:
|
|
blocker = threading.Thread(target=_call, args=(port, None, "9^9^9^9"), daemon=True)
|
|
blocker.start()
|
|
time.sleep(0.5)
|
|
self.assertEqual(_call(port, None)[0], 503)
|
|
blocker.join(timeout=10)
|
|
self.assertEqual(_call(port, None)[0], 200)
|
|
finally:
|
|
proc.terminate()
|
|
proc.communicate(timeout=30)
|
|
|