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
1 changed files with 97 additions and 187 deletions
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Reference in a new issue