test: fix WS test port race; narrow to single smoke test covering both transport ends (#2267)
This commit is contained in:
Max Isbey
GitHub
parent
7826ade12b
commit
67201a9bbd
3 changed files with 97 additions and 187 deletions
|
|
@ -8,7 +8,7 @@ import anyio
|
|||
import pytest
|
||||
from inline_snapshot import snapshot
|
||||
|
||||
from mcp import types
|
||||
from mcp import MCPError, types
|
||||
from mcp.client._memory import InMemoryTransport
|
||||
from mcp.client.client import Client
|
||||
from mcp.server import Server, ServerRequestContext
|
||||
|
|
@ -175,6 +175,21 @@ async def test_read_resource(app: MCPServer):
|
|||
)
|
||||
|
||||
|
||||
async def test_read_resource_error_propagates():
|
||||
"""MCPError raised by a server handler propagates to the client with its code intact."""
|
||||
|
||||
async def handle_read_resource(
|
||||
ctx: ServerRequestContext, params: types.ReadResourceRequestParams
|
||||
) -> ReadResourceResult:
|
||||
raise MCPError(code=404, message="no resource with that URI was found")
|
||||
|
||||
server = Server("test", on_read_resource=handle_read_resource)
|
||||
async with Client(server) as client:
|
||||
with pytest.raises(MCPError) as exc_info:
|
||||
await client.read_resource("unknown://example")
|
||||
assert exc_info.value.error.code == 404
|
||||
|
||||
|
||||
async def test_get_prompt(app: MCPServer):
|
||||
"""Test getting a prompt."""
|
||||
async with Client(app) as client:
|
||||
|
|
|
|||
|
|
@ -1,210 +1,51 @@
|
|||
import multiprocessing
|
||||
import socket
|
||||
from collections.abc import AsyncGenerator, Generator
|
||||
from urllib.parse import urlparse
|
||||
"""Smoke test for the WebSocket transport.
|
||||
|
||||
Runs the full WS stack end-to-end over a real TCP connection, covering both
|
||||
``src/mcp/client/websocket.py`` and ``src/mcp/server/websocket.py``. MCP
|
||||
semantics (error propagation, timeouts, etc.) are transport-agnostic and are
|
||||
covered in ``tests/client/test_client.py`` and ``tests/issues/test_88_random_error.py``.
|
||||
"""
|
||||
|
||||
from collections.abc import Generator
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
import uvicorn
|
||||
from starlette.applications import Starlette
|
||||
from starlette.routing import WebSocketRoute
|
||||
from starlette.websockets import WebSocket
|
||||
|
||||
from mcp import MCPError
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.client.websocket import websocket_client
|
||||
from mcp.server import Server, ServerRequestContext
|
||||
from mcp.server import Server
|
||||
from mcp.server.websocket import websocket_server
|
||||
from mcp.types import (
|
||||
CallToolRequestParams,
|
||||
CallToolResult,
|
||||
EmptyResult,
|
||||
InitializeResult,
|
||||
ListToolsResult,
|
||||
PaginatedRequestParams,
|
||||
ReadResourceRequestParams,
|
||||
ReadResourceResult,
|
||||
TextContent,
|
||||
TextResourceContents,
|
||||
Tool,
|
||||
)
|
||||
from tests.test_helpers import wait_for_server
|
||||
from mcp.types import EmptyResult, InitializeResult
|
||||
from tests.test_helpers import run_uvicorn_in_thread
|
||||
|
||||
SERVER_NAME = "test_server_for_WS"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def server_port() -> int:
|
||||
with socket.socket() as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
return s.getsockname()[1]
|
||||
def make_server_app() -> Starlette:
|
||||
srv = Server(SERVER_NAME)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def server_url(server_port: int) -> str:
|
||||
return f"ws://127.0.0.1:{server_port}"
|
||||
|
||||
|
||||
async def handle_read_resource( # pragma: no cover
|
||||
ctx: ServerRequestContext, params: ReadResourceRequestParams
|
||||
) -> ReadResourceResult:
|
||||
parsed = urlparse(str(params.uri))
|
||||
if parsed.scheme == "foobar":
|
||||
return ReadResourceResult(
|
||||
contents=[TextResourceContents(uri=str(params.uri), text=f"Read {parsed.netloc}", mime_type="text/plain")]
|
||||
)
|
||||
elif parsed.scheme == "slow":
|
||||
await anyio.sleep(2.0)
|
||||
return ReadResourceResult(
|
||||
contents=[
|
||||
TextResourceContents(
|
||||
uri=str(params.uri), text=f"Slow response from {parsed.netloc}", mime_type="text/plain"
|
||||
)
|
||||
]
|
||||
)
|
||||
raise MCPError(code=404, message="OOPS! no resource with that URI was found")
|
||||
|
||||
|
||||
async def handle_list_tools( # pragma: no cover
|
||||
ctx: ServerRequestContext, params: PaginatedRequestParams | None
|
||||
) -> ListToolsResult:
|
||||
return ListToolsResult(
|
||||
tools=[
|
||||
Tool(
|
||||
name="test_tool",
|
||||
description="A test tool",
|
||||
input_schema={"type": "object", "properties": {}},
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
async def handle_call_tool( # pragma: no cover
|
||||
ctx: ServerRequestContext, params: CallToolRequestParams
|
||||
) -> CallToolResult:
|
||||
return CallToolResult(content=[TextContent(type="text", text=f"Called {params.name}")])
|
||||
|
||||
|
||||
def _create_server() -> Server: # pragma: no cover
|
||||
return Server(
|
||||
SERVER_NAME,
|
||||
on_read_resource=handle_read_resource,
|
||||
on_list_tools=handle_list_tools,
|
||||
on_call_tool=handle_call_tool,
|
||||
)
|
||||
|
||||
|
||||
# Test fixtures
|
||||
def make_server_app() -> Starlette: # pragma: no cover
|
||||
"""Create test Starlette app with WebSocket transport"""
|
||||
server = _create_server()
|
||||
|
||||
async def handle_ws(websocket: WebSocket):
|
||||
async def handle_ws(websocket: WebSocket) -> None:
|
||||
async with websocket_server(websocket.scope, websocket.receive, websocket.send) as streams:
|
||||
await server.run(streams[0], streams[1], server.create_initialization_options())
|
||||
await srv.run(streams[0], streams[1], srv.create_initialization_options())
|
||||
|
||||
app = Starlette(routes=[WebSocketRoute("/ws", endpoint=handle_ws)])
|
||||
return app
|
||||
return Starlette(routes=[WebSocketRoute("/ws", endpoint=handle_ws)])
|
||||
|
||||
|
||||
def run_server(server_port: int) -> None: # pragma: no cover
|
||||
app = make_server_app()
|
||||
server = uvicorn.Server(config=uvicorn.Config(app=app, host="127.0.0.1", port=server_port, log_level="error"))
|
||||
print(f"starting server on {server_port}")
|
||||
server.run()
|
||||
@pytest.fixture
|
||||
def ws_server_url() -> Generator[str, None, None]:
|
||||
with run_uvicorn_in_thread(make_server_app()) as base_url:
|
||||
yield base_url.replace("http://", "ws://") + "/ws"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def server(server_port: int) -> Generator[None, None, None]:
|
||||
proc = multiprocessing.Process(target=run_server, kwargs={"server_port": server_port}, daemon=True)
|
||||
print("starting process")
|
||||
proc.start()
|
||||
|
||||
# Wait for server to be running
|
||||
print("waiting for server to start")
|
||||
wait_for_server(server_port)
|
||||
|
||||
yield
|
||||
|
||||
print("killing server")
|
||||
# Signal the server to stop
|
||||
proc.kill()
|
||||
proc.join(timeout=2)
|
||||
if proc.is_alive(): # pragma: no cover
|
||||
print("server process failed to terminate")
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
async def initialized_ws_client_session(server: None, server_url: str) -> AsyncGenerator[ClientSession, None]:
|
||||
"""Create and initialize a WebSocket client session"""
|
||||
async with websocket_client(server_url + "/ws") as streams:
|
||||
@pytest.mark.anyio
|
||||
async def test_ws_client_basic_connection(ws_server_url: str) -> None:
|
||||
async with websocket_client(ws_server_url) as streams:
|
||||
async with ClientSession(*streams) as session:
|
||||
# Test initialization
|
||||
result = await session.initialize()
|
||||
assert isinstance(result, InitializeResult)
|
||||
assert result.server_info.name == SERVER_NAME
|
||||
|
||||
# Test ping
|
||||
ping_result = await session.send_ping()
|
||||
assert isinstance(ping_result, EmptyResult)
|
||||
|
||||
yield session
|
||||
|
||||
|
||||
# Tests
|
||||
@pytest.mark.anyio
|
||||
async def test_ws_client_basic_connection(server: None, server_url: str) -> None:
|
||||
"""Test the WebSocket connection establishment"""
|
||||
async with websocket_client(server_url + "/ws") as streams:
|
||||
async with ClientSession(*streams) as session:
|
||||
# Test initialization
|
||||
result = await session.initialize()
|
||||
assert isinstance(result, InitializeResult)
|
||||
assert result.server_info.name == SERVER_NAME
|
||||
|
||||
# Test ping
|
||||
ping_result = await session.send_ping()
|
||||
assert isinstance(ping_result, EmptyResult)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_ws_client_happy_request_and_response(
|
||||
initialized_ws_client_session: ClientSession,
|
||||
) -> None:
|
||||
"""Test a successful request and response via WebSocket"""
|
||||
result = await initialized_ws_client_session.read_resource("foobar://example")
|
||||
assert isinstance(result, ReadResourceResult)
|
||||
assert isinstance(result.contents, list)
|
||||
assert len(result.contents) > 0
|
||||
assert isinstance(result.contents[0], TextResourceContents)
|
||||
assert result.contents[0].text == "Read example"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_ws_client_exception_handling(
|
||||
initialized_ws_client_session: ClientSession,
|
||||
) -> None:
|
||||
"""Test exception handling in WebSocket communication"""
|
||||
with pytest.raises(MCPError) as exc_info:
|
||||
await initialized_ws_client_session.read_resource("unknown://example")
|
||||
assert exc_info.value.error.code == 404
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_ws_client_timeout(
|
||||
initialized_ws_client_session: ClientSession,
|
||||
) -> None:
|
||||
"""Test timeout handling in WebSocket communication"""
|
||||
# Set a very short timeout to trigger a timeout exception
|
||||
with pytest.raises(TimeoutError):
|
||||
with anyio.fail_after(0.1): # 100ms timeout
|
||||
await initialized_ws_client_session.read_resource("slow://example")
|
||||
|
||||
# Now test that we can still use the session after a timeout
|
||||
with anyio.fail_after(5): # Longer timeout to allow completion
|
||||
result = await initialized_ws_client_session.read_resource("foobar://example")
|
||||
assert isinstance(result, ReadResourceResult)
|
||||
assert isinstance(result.contents, list)
|
||||
assert len(result.contents) > 0
|
||||
assert isinstance(result.contents[0], TextResourceContents)
|
||||
assert result.contents[0].text == "Read example"
|
||||
|
|
|
|||
|
|
@ -1,7 +1,61 @@
|
|||
"""Common test utilities for MCP server tests."""
|
||||
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Generator
|
||||
from contextlib import contextmanager
|
||||
from typing import Any
|
||||
|
||||
import uvicorn
|
||||
|
||||
_SERVER_SHUTDOWN_TIMEOUT_S = 5.0
|
||||
|
||||
|
||||
@contextmanager
|
||||
def run_uvicorn_in_thread(app: Any, **config_kwargs: Any) -> Generator[str, None, None]:
|
||||
"""Run a uvicorn server in a background thread on an ephemeral port.
|
||||
|
||||
The socket is bound and put into listening state *before* the thread
|
||||
starts, so the port is known immediately with no wait. The kernel's
|
||||
listen queue buffers any connections that arrive before uvicorn's event
|
||||
loop reaches ``accept()``, so callers can connect as soon as this
|
||||
function yields — no polling, no sleeps, no startup race.
|
||||
|
||||
This also avoids the TOCTOU race of the old pick-a-port-then-rebind
|
||||
pattern: the socket passed here is the one uvicorn serves on, with no
|
||||
gap where another pytest-xdist worker could claim it.
|
||||
|
||||
Args:
|
||||
app: ASGI application to serve.
|
||||
**config_kwargs: Additional keyword arguments for :class:`uvicorn.Config`
|
||||
(e.g. ``log_level``). ``host``/``port`` are ignored since the
|
||||
socket is pre-bound.
|
||||
|
||||
Yields:
|
||||
The base URL of the running server, e.g. ``http://127.0.0.1:54321``.
|
||||
"""
|
||||
host = "127.0.0.1"
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
sock.bind((host, 0))
|
||||
sock.listen()
|
||||
port = sock.getsockname()[1]
|
||||
|
||||
config_kwargs.setdefault("log_level", "error")
|
||||
# Uvicorn's interface autodetection calls asyncio.iscoroutinefunction,
|
||||
# which Python 3.14 deprecates. Under filterwarnings=error this crashes
|
||||
# the server thread silently. Starlette is asgi3; skip the autodetect.
|
||||
config_kwargs.setdefault("interface", "asgi3")
|
||||
server = uvicorn.Server(config=uvicorn.Config(app=app, **config_kwargs))
|
||||
|
||||
thread = threading.Thread(target=server.run, kwargs={"sockets": [sock]}, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
yield f"http://{host}:{port}"
|
||||
finally:
|
||||
server.should_exit = True
|
||||
thread.join(timeout=_SERVER_SHUTDOWN_TIMEOUT_S)
|
||||
|
||||
|
||||
def wait_for_server(port: int, timeout: float = 20.0) -> None:
|
||||
|
|
|
|||
Loading…
Reference in a new issue