X Tutup
Skip to content

test: convert WebSocket tests to in-memory transport#2267

Draft
maxisbey wants to merge 1 commit intomainfrom
maxisbey/max-157-flaky-convert-websocket-tests-to-in-memory-transport-remove
Draft

test: convert WebSocket tests to in-memory transport#2267
maxisbey wants to merge 1 commit intomainfrom
maxisbey/max-157-flaky-convert-websocket-tests-to-in-memory-transport-remove

Conversation

@maxisbey
Copy link
Contributor

Summary

tests/shared/test_ws.py used the subprocess + TCP port pattern that races under pytest-xdist: a worker allocates a port with socket.bind(0), releases it, then spawns a uvicorn subprocess hoping to rebind. Between release and rebind, another worker can claim the port — observed in CI as the WS client connecting to another worker's HTTP/SSE server and getting a 403 Forbidden on the upgrade (job 63414508256).

What changed

Three tests → in-memory Client(server) transport (no network, no subprocess, no race):

  • test_ws_client_happy_request_and_response
  • test_ws_client_exception_handling
  • test_ws_client_timeout

These verify transport-agnostic MCP semantics (read_resource happy path, MCPError propagation, session recovery after client-side timeout). None of them were testing WebSocket framing.

One test kept as real-TCP smoke testtest_ws_client_basic_connection

Uses a new run_uvicorn_in_thread helper that binds port=0 atomically and reads the actual port back from the server's bound socket — the OS holds the port from bind to shutdown, eliminating the race window. This test alone provides 100% coverage of src/mcp/client/websocket.py.

Cleanup

  • Removed handle_list_tools / handle_call_tool (never exercised by any test)
  • Removed # pragma: no cover from handle_read_resource (now runs in-process)
  • Replaced anyio.sleep(2.0) with anyio.sleep_forever() for the slow-resource handler — semantically clearer that it blocks until cancellation

Overlap with MAX-158

MAX-158 adds the same run_uvicorn_in_thread helper to fix the port race across many test files. Whichever PR lands first, the other will have a trivial rebase (the helper implementation is identical). This PR keeps wait_for_server in test_helpers.py since other test files on main still use it.

Test Plan

  • ./scripts/test — full suite + 100% coverage ✓
  • pytest tests/shared/test_ws.py -n 4 --flake-finder --flake-runs=5 — 20 parallel runs, 0 flakes ✓
  • pyright + ruff

AI Disclaimer

test_ws.py used the subprocess + TCP port pattern that races under
pytest-xdist: a worker allocates a port with socket.bind(0), releases
it, then spawns a uvicorn subprocess hoping to rebind. Between release
and rebind, another worker can claim the port, causing the WS client to
connect to an unrelated server (observed: HTTP 403 Forbidden on the
WebSocket upgrade).

Three of the four tests here verify transport-agnostic MCP semantics
(read_resource happy path, MCPError propagation, session recovery after
client-side timeout). These now use the in-memory Client transport — no
network, no subprocess, no race.

The fourth test (test_ws_client_basic_connection) is kept as a smoke
test running the real WS stack end-to-end. It uses a new
run_uvicorn_in_thread helper that binds port=0 atomically and reads the
actual port back from the server's socket — the OS holds the port from
bind to shutdown, eliminating the race window entirely. This test alone
provides 100% coverage of src/mcp/client/websocket.py.

Also removed dead handler code (list_tools/call_tool were never
exercised) and the no-longer-needed pragma: no cover annotations on the
read_resource handler (it now runs in-process).
@maxisbey maxisbey marked this pull request as draft March 10, 2026 14:42
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant

X Tutup