Mocking a WebSocket-Based AI Chat in an End-to-End Test

A WebSocket chat has a client that talks back, a connection that outlives a single answer, and a reconnect path. None of those exist in an SSE test, and all three are where the bugs are.

Why this is not the SSE test again

With server-sent events the request is the unit: one POST, one stream, done. A WebSocket chat multiplexes. One connection carries several conversations’ worth of frames, the client sends frames of its own (a new message, a cancel, a heartbeat), and the connection can drop and come back in the middle of an answer. That last case produces the bug this test exists for: the client reconnects, the server replays from a checkpoint, and the widget appends tokens it has already rendered.

Playwright routes WebSockets with page.routeWebSocket(url, handler). The handler receives a WebSocketRoute, whose documented methods are onMessage(handler), send(message), close(options), connectToServer(), onClose(handler), protocols() and url(). Two of those decide which of the two test styles you are writing.

There is a second structural difference worth naming before you write any code. Over SSE, the framing is fixed by the protocol: text, lines beginning data:, blank line between events. Over a WebSocket the framing is entirely yours, which means your fixture has to encode a private protocol correctly, and a fixture that encodes it slightly wrong produces a test that passes against a mock nothing else agrees with. Generate the fixture frames from the same serialiser the server uses, or at minimum keep one test that validates a recorded real frame against your fixture’s shape.

A fully mocked socket

If you never call connectToServer(), no real connection is made and the route object is the server. Frames the page sends arrive in onMessage; send() pushes a frame to the page. That gives the test complete control of timing, which is the thing the fulfil-based SSE approach could not do.

import { test, expect } from "@playwright/test";

test("streams tokens over a websocket", async ({ page }) => {
  await page.routeWebSocket("**/ws/chat", (ws) => {
    ws.onMessage((raw) => {
      const msg = JSON.parse(String(raw));
      if (msg.type !== "user_message") return;

      ws.send(JSON.stringify({ type: "start", seq: 1, id: "m_1" }));
      ws.send(JSON.stringify({ type: "token", seq: 2, id: "m_1", text: "The invoice " }));
      ws.send(JSON.stringify({ type: "token", seq: 3, id: "m_1", text: "is overdue." }));
      ws.send(JSON.stringify({ type: "done", seq: 4, id: "m_1" }));
    });
  });

  await page.goto("/support");
  await page.getByTestId("composer").fill("where is my invoice");
  await page.getByTestId("send").click();

  await expect(page.getByTestId("message-assistant")).toHaveText(
    "The invoice is overdue.",
  );
  await expect(page.getByTestId("send")).toBeEnabled();
});

Because onMessage sees the client’s frames, the first assertion to add is on the outbound side: that clicking send produced exactly one user_message frame, that it carried the conversation id, and that pressing the stop button produced a cancel frame referencing the message id rather than closing the socket. Collect frames into an array in the handler and assert on it after the interaction.

Passing through to the real server

connectToServer() opens the real connection and returns the server-side route, so you can let almost everything flow and intercept one frame type. This is the right shape when the socket also carries presence, typing indicators or auth handshakes you do not want to reimplement in a stub — reimplementing a protocol in a mock is how a test suite ends up asserting that your mock matches your mock.

await page.routeWebSocket("**/ws/chat", (ws) => {
  const server = ws.connectToServer();

  // Client to server: pass everything through untouched.
  ws.onMessage((m) => server.send(m));

  // Server to client: swap the model's tokens for a fixture,
  // leave every other frame type alone.
  server.onMessage((raw) => {
    const msg = JSON.parse(String(raw));
    if (msg.type === "token") {
      ws.send(JSON.stringify({ ...msg, text: "FIXTURE " }));
    } else {
      ws.send(raw);
    }
  });
});

The trade is coverage against determinism: the handshake and the presence protocol are genuinely exercised, but the test now needs a server running, which makes it slower and couples it to a deployment. Keep the fully mocked version as the fast suite and the pass-through version as a small smoke set — see replaying recorded traffic.

A practical note on choosing between the two: the fully mocked style answers “does my client behave correctly given these frames” and the pass-through style answers “does the server send the frames I think it does”. Those are different questions and a suite usually needs both, but only the first belongs in the fast loop. If you find yourself adding server behaviour to the mock — auth, presence, rate limiting — that is the signal to move that test to pass-through rather than to keep growing the mock.

The reconnect assertion

This is the test worth writing the whole page for. close() shuts one side of the connection, so you can drop the socket mid-answer and let the client’s reconnect logic run for real.

1. Send start and two token frames with sequence numbers 2 and 3, then call ws.close({ code: 1006 }) to simulate an abnormal closure rather than a clean one. Clients frequently treat 1000 and 1006 differently, and only one of them should trigger a reconnect.
2. Assert the partial text is still on screen and marked as interrupted. Losing it on disconnect is the most common failure.
3. The route handler fires again for the new connection. Assert the client’s first frame is a resume carrying the last sequence number it saw, not a fresh subscribe — a client that resubscribes from zero will render the answer twice.
4. Replay sequence 3 deliberately along with 4 and 5. Assert the final text contains the token from frame 3 exactly once. Deduplication by sequence number is the invariant; assert the invariant, not the prose.
5. Assert the reconnect is backed off. Close the socket three times in a row and check the client did not open three sockets inside a few milliseconds.

Without routeWebSocket

Cypress has no first-class WebSocket interception, so the equivalent is to replace window.WebSocket in the onBeforeLoad hook with a fake class that records outbound frames and exposes a method for the test to push inbound ones. It is the same technique as the EventSource stub for a Cypress chat test, with two extra members: a readyState that moves through the documented constants, and an onclose that fires with the code you choose. Getting readyState wrong is the usual reason a fake socket makes a widget hang: clients queue outbound frames until the socket reports open, and a fake that never reports open queues forever.