test: fix WS test port race; narrow to single smoke test covering both transport ends (#2267)

This commit is contained in:
Max Isbey 2026-03-18 15:48:30 +00:00 • committed by GitHub
commit 67201a9bbd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 97 additions and 187 deletions

View file

@ -10,7 +10,7 @@ from mcp import types
from mcp.shared.message import SessionMessage from mcp.shared.message import SessionMessage
@asynccontextmanager # pragma: no cover @asynccontextmanager
async def websocket_server(scope: Scope, receive: Receive, send: Send): async def websocket_server(scope: Scope, receive: Receive, send: Send):
"""WebSocket server transport for MCP. This is an ASGI application, suitable for use """WebSocket server transport for MCP. This is an ASGI application, suitable for use
with a framework like Starlette and a server like Hypercorn. with a framework like Starlette and a server like Hypercorn.
@ -34,13 +34,13 @@ async def websocket_server(scope: Scope, receive: Receive, send: Send):
async for msg in websocket.iter_text(): async for msg in websocket.iter_text():
try: try:
client_message = types.jsonrpc_message_adapter.validate_json(msg, by_name=False) client_message = types.jsonrpc_message_adapter.validate_json(msg, by_name=False)
except ValidationError as exc: except ValidationError as exc: # pragma: no cover
await read_stream_writer.send(exc) await read_stream_writer.send(exc)
continue continue
session_message = SessionMessage(client_message) session_message = SessionMessage(client_message)
await read_stream_writer.send(session_message) await read_stream_writer.send(session_message)
except anyio.ClosedResourceError: except anyio.ClosedResourceError: # pragma: no cover
await websocket.close() await websocket.close()
async def ws_writer(): async def ws_writer():
@ -49,7 +49,7 @@ async def websocket_server(scope: Scope, receive: Receive, send: Send):
async for session_message in write_stream_reader: async for session_message in write_stream_reader:
obj = session_message.message.model_dump_json(by_alias=True, exclude_unset=True) obj = session_message.message.model_dump_json(by_alias=True, exclude_unset=True)
await websocket.send_text(obj) await websocket.send_text(obj)
except anyio.ClosedResourceError: except anyio.ClosedResourceError: # pragma: no cover
await websocket.close() await websocket.close()
async with anyio.create_task_group() as tg: async with anyio.create_task_group() as tg:

View file

@ -8,7 +8,7 @@ import anyio
import pytest import pytest
from inline_snapshot import snapshot from inline_snapshot import snapshot
from mcp import types from mcp import MCPError, types
from mcp.client._memory import InMemoryTransport from mcp.client._memory import InMemoryTransport
from mcp.client.client import Client from mcp.client.client import Client
from mcp.server import Server, ServerRequestContext 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): async def test_get_prompt(app: MCPServer):
"""Test getting a prompt.""" """Test getting a prompt."""
async with Client(app) as client: async with Client(app) as client:

View file

@ -1,210 +1,51 @@
import multiprocessing """Smoke test for the WebSocket transport.
import socket
from collections.abc import AsyncGenerator, Generator Runs the full WS stack end-to-end over a real TCP connection, covering both
from urllib.parse import urlparse ``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 pytest
import uvicorn
from starlette.applications import Starlette from starlette.applications import Starlette
from starlette.routing import WebSocketRoute from starlette.routing import WebSocketRoute
from starlette.websockets import WebSocket from starlette.websockets import WebSocket
from mcp import MCPError
from mcp.client.session import ClientSession from mcp.client.session import ClientSession
from mcp.client.websocket import websocket_client 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.server.websocket import websocket_server
from mcp.types import ( from mcp.types import EmptyResult, InitializeResult
CallToolRequestParams, from tests.test_helpers import run_uvicorn_in_thread
CallToolResult,
EmptyResult,
InitializeResult,
ListToolsResult,
PaginatedRequestParams,
ReadResourceRequestParams,
ReadResourceResult,
TextContent,
TextResourceContents,
Tool,
)
from tests.test_helpers import wait_for_server
SERVER_NAME = "test_server_for_WS" SERVER_NAME = "test_server_for_WS"
@pytest.fixture def make_server_app() -> Starlette:
def server_port() -> int: srv = Server(SERVER_NAME)
with socket.socket() as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
async def handle_ws(websocket: WebSocket) -> None:
@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 with websocket_server(websocket.scope, websocket.receive, websocket.send) as streams: 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 Starlette(routes=[WebSocketRoute("/ws", endpoint=handle_ws)])
return app
def run_server(server_port: int) -> None: # pragma: no cover @pytest.fixture
app = make_server_app() def ws_server_url() -> Generator[str, None, None]:
server = uvicorn.Server(config=uvicorn.Config(app=app, host="127.0.0.1", port=server_port, log_level="error")) with run_uvicorn_in_thread(make_server_app()) as base_url:
print(f"starting server on {server_port}") yield base_url.replace("http://", "ws://") + "/ws"
server.run()
@pytest.fixture() @pytest.mark.anyio
def server(server_port: int) -> Generator[None, None, None]: async def test_ws_client_basic_connection(ws_server_url: str) -> None:
proc = multiprocessing.Process(target=run_server, kwargs={"server_port": server_port}, daemon=True) async with websocket_client(ws_server_url) as streams:
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:
async with ClientSession(*streams) as session: async with ClientSession(*streams) as session:
# Test initialization
result = await session.initialize() result = await session.initialize()
assert isinstance(result, InitializeResult) assert isinstance(result, InitializeResult)
assert result.server_info.name == SERVER_NAME assert result.server_info.name == SERVER_NAME
# Test ping
ping_result = await session.send_ping() ping_result = await session.send_ping()
assert isinstance(ping_result, EmptyResult) 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"

View file

@ -1,7 +1,61 @@
"""Common test utilities for MCP server tests.""" """Common test utilities for MCP server tests."""
import socket import socket
import threading
import time 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: def wait_for_server(port: int, timeout: float = 20.0) -> None: