fix(site): add tests for createMockWebSocket (#19172)

Parkreiner · web-flow · commit 5225c5671d40 · 2025-08-07T16:54:50.000-04:00
Needed for #19126 and #18679 ## Changes made - Moved `createWebSocket` to dedicated file and addressed edge cases for making it a reliable mock - Added test cases to validate mock functionality
diff --git a/site/src/testHelpers/websockets.test.ts b/site/src/testHelpers/websockets.test.ts
@@ -0,0 +1,186 @@
+import { createMockWebSocket } from "./websockets";
+
+describe(createMockWebSocket.name, () => {
+	it("Throws if URL does not have ws:// or wss:// protocols", () => {
+		const urls: readonly string[] = [
+			"http://www.dog.ceo/roll-over",
+			"https://www.dog.ceo/roll-over",
+		];
+		for (const url of urls) {
+			expect(() => {
+				void createMockWebSocket(url);
+			}).toThrow("URL must start with ws:// or wss://");
+		}
+	});
+
+	it("Sends events from server to socket", () => {
+		const [socket, server] = createMockWebSocket("wss://www.dog.ceo/shake");
+
+		const onOpen = jest.fn();
+		const onError = jest.fn();
+		const onMessage = jest.fn();
+		const onClose = jest.fn();
+
+		socket.addEventListener("open", onOpen);
+		socket.addEventListener("error", onError);
+		socket.addEventListener("message", onMessage);
+		socket.addEventListener("close", onClose);
+
+		const openEvent = new Event("open");
+		const errorEvent = new Event("error");
+		const messageEvent = new MessageEvent<string>("message");
+		const closeEvent = new CloseEvent("close");
+
+		server.publishOpen(openEvent);
+		server.publishError(errorEvent);
+		server.publishMessage(messageEvent);
+		server.publishClose(closeEvent);
+
+		expect(onOpen).toHaveBeenCalledTimes(1);
+		expect(onOpen).toHaveBeenCalledWith(openEvent);
+
+		expect(onError).toHaveBeenCalledTimes(1);
+		expect(onError).toHaveBeenCalledWith(errorEvent);
+
+		expect(onMessage).toHaveBeenCalledTimes(1);
+		expect(onMessage).toHaveBeenCalledWith(messageEvent);
+
+		expect(onClose).toHaveBeenCalledTimes(1);
+		expect(onClose).toHaveBeenCalledWith(closeEvent);
+	});
+
+	it("Sends JSON data to the socket for message events", () => {
+		const [socket, server] = createMockWebSocket("wss://www.dog.ceo/wag");
+		const onMessage = jest.fn();
+
+		// Could type this as a special JSON type, but unknown is good enough,
+		// since any invalid values will throw in the test case
+		const jsonData: readonly unknown[] = [
+			"blah",
+			42,
+			true,
+			false,
+			null,
+			{},
+			[],
+			[{ value: "blah" }, { value: "guh" }, { value: "huh" }],
+			{
+				name: "Hershel Layton",
+				age: 40,
+				profession: "Puzzle Solver",
+				sadBackstory: true,
+				greatVideoGames: true,
+			},
+		];
+		for (const jd of jsonData) {
+			socket.addEventListener("message", onMessage);
+			server.publishMessage(
+				new MessageEvent("message", { data: JSON.stringify(jd) }),
+			);
+
+			expect(onMessage).toHaveBeenCalledTimes(1);
+			expect(onMessage).toHaveBeenCalledWith(
+				new MessageEvent("message", { data: JSON.stringify(jd) }),
+			);
+
+			socket.removeEventListener("message", onMessage);
+			onMessage.mockClear();
+		}
+	});
+
+	it("Only registers each socket event handler once", () => {
+		const [socket, server] = createMockWebSocket("wss://www.dog.ceo/borf");
+
+		const onOpen = jest.fn();
+		const onError = jest.fn();
+		const onMessage = jest.fn();
+		const onClose = jest.fn();
+
+		// Do it once
+		socket.addEventListener("open", onOpen);
+		socket.addEventListener("error", onError);
+		socket.addEventListener("message", onMessage);
+		socket.addEventListener("close", onClose);
+
+		// Do it again with the exact same functions
+		socket.addEventListener("open", onOpen);
+		socket.addEventListener("error", onError);
+		socket.addEventListener("message", onMessage);
+		socket.addEventListener("close", onClose);
+
+		server.publishOpen(new Event("open"));
+		server.publishError(new Event("error"));
+		server.publishMessage(new MessageEvent<string>("message"));
+		server.publishClose(new CloseEvent("close"));
+
+		expect(onOpen).toHaveBeenCalledTimes(1);
+		expect(onError).toHaveBeenCalledTimes(1);
+		expect(onMessage).toHaveBeenCalledTimes(1);
+		expect(onClose).toHaveBeenCalledTimes(1);
+	});
+
+	it("Lets a socket unsubscribe to event types", () => {
+		const [socket, server] = createMockWebSocket("wss://www.dog.ceo/zoomies");
+
+		const onOpen = jest.fn();
+		const onError = jest.fn();
+		const onMessage = jest.fn();
+		const onClose = jest.fn();
+
+		socket.addEventListener("open", onOpen);
+		socket.addEventListener("error", onError);
+		socket.addEventListener("message", onMessage);
+		socket.addEventListener("close", onClose);
+
+		socket.removeEventListener("open", onOpen);
+		socket.removeEventListener("error", onError);
+		socket.removeEventListener("message", onMessage);
+		socket.removeEventListener("close", onClose);
+
+		server.publishOpen(new Event("open"));
+		server.publishError(new Event("error"));
+		server.publishMessage(new MessageEvent<string>("message"));
+		server.publishClose(new CloseEvent("close"));
+
+		expect(onOpen).not.toHaveBeenCalled();
+		expect(onError).not.toHaveBeenCalled();
+		expect(onMessage).not.toHaveBeenCalled();
+		expect(onClose).not.toHaveBeenCalled();
+	});
+
+	it("Renders socket inert after being closed", () => {
+		const [socket, server] = createMockWebSocket("wss://www.dog.ceo/woof");
+		expect(server.isConnectionOpen).toBe(true);
+
+		const onMessage = jest.fn();
+		socket.addEventListener("message", onMessage);
+
+		socket.close();
+		expect(server.isConnectionOpen).toBe(false);
+
+		server.publishMessage(new MessageEvent<string>("message"));
+		expect(onMessage).not.toHaveBeenCalled();
+	});
+
+	it("Tracks arguments sent by the mock socket", () => {
+		const [socket, server] = createMockWebSocket("wss://www.dog.ceo/wan-wan");
+		const data = JSON.stringify({
+			famousDogs: [
+				"snoopy",
+				"clifford",
+				"lassie",
+				"beethoven",
+				"courage the cowardly dog",
+			],
+		});
+
+		socket.send(data);
+		expect(server.clientSentData).toHaveLength(1);
+		expect(server.clientSentData).toEqual([data]);
+
+		socket.close();
+		socket.send(data);
+		expect(server.clientSentData).toHaveLength(1);
+		expect(server.clientSentData).toEqual([data]);
+	});
+});
diff --git a/site/src/testHelpers/websockets.ts b/site/src/testHelpers/websockets.ts
@@ -0,0 +1,162 @@
+import type { WebSocketEventType } from "utils/OneWayWebSocket";
+
+type SocketSendData = Parameters<WebSocket["send"]>[0];
+
+export type MockWebSocketServer = Readonly<{
+	publishMessage: (event: MessageEvent<string>) => void;
+	publishError: (event: Event) => void;
+	publishClose: (event: CloseEvent) => void;
+	publishOpen: (event: Event) => void;
+
+	readonly isConnectionOpen: boolean;
+	readonly clientSentData: readonly SocketSendData[];
+}>;
+
+type CallbackStore = {
+	[K in keyof WebSocketEventMap]: Set<(event: WebSocketEventMap[K]) => void>;
+};
+
+type MockWebSocket = Omit<WebSocket, "send"> & {
+	/**
+	 * A version of the WebSocket `send` method that has been pre-wrapped inside
+	 * a Jest mock.
+	 *
+	 * The Jest mock functionality should be used at a minimum. Basically:
+	 * 1. If you want to check that the mock socket sent something to the mock
+	 *    server: call the `send` method as a function, and then check the
+	 *    `clientSentData` on `MockWebSocketServer` to see what data got
+	 *    received.
+	 * 2. If you need to make sure that the client-side `send` method got called
+	 *    at all: you can use the Jest mock functionality, but you should
+	 *    probably also be checking `clientSentData` still and making additional
+	 *    assertions with it.
+	 *
+	 * Generally, tests should center around whether socket-to-server
+	 * communication was successful, not whether the client-side method was
+	 * called.
+	 */
+	send: jest.Mock<void, [SocketSendData], unknown>;
+};
+
+export function createMockWebSocket(
+	url: string,
+	protocol?: string | string[] | undefined,
+): readonly [MockWebSocket, MockWebSocketServer] {
+	if (!url.startsWith("ws://") && !url.startsWith("wss://")) {
+		throw new Error("URL must start with ws:// or wss://");
+	}
+
+	const activeProtocol = Array.isArray(protocol)
+		? protocol.join(" ")
+		: (protocol ?? "");
+
+	let isOpen = true;
+	const store: CallbackStore = {
+		message: new Set(),
+		error: new Set(),
+		close: new Set(),
+		open: new Set(),
+	};
+
+	const sentData: SocketSendData[] = [];
+
+	const mockSocket: MockWebSocket = {
+		CONNECTING: 0,
+		OPEN: 1,
+		CLOSING: 2,
+		CLOSED: 3,
+
+		url,
+		protocol: activeProtocol,
+		readyState: 1,
+		binaryType: "blob",
+		bufferedAmount: 0,
+		extensions: "",
+		onclose: null,
+		onerror: null,
+		onmessage: null,
+		onopen: null,
+		dispatchEvent: jest.fn(),
+
+		send: jest.fn((data) => {
+			if (!isOpen) {
+				return;
+			}
+			sentData.push(data);
+		}),
+
+		addEventListener: <E extends WebSocketEventType>(
+			eventType: E,
+			callback: (event: WebSocketEventMap[E]) => void,
+		) => {
+			if (!isOpen) {
+				return;
+			}
+			const subscribers = store[eventType];
+			subscribers.add(callback);
+		},
+
+		removeEventListener: <E extends WebSocketEventType>(
+			eventType: E,
+			callback: (event: WebSocketEventMap[E]) => void,
+		) => {
+			if (!isOpen) {
+				return;
+			}
+			const subscribers = store[eventType];
+			subscribers.delete(callback);
+		},
+
+		close: () => {
+			isOpen = false;
+		},
+	};
+
+	const publisher: MockWebSocketServer = {
+		get isConnectionOpen() {
+			return isOpen;
+		},
+
+		get clientSentData() {
+			return [...sentData];
+		},
+
+		publishOpen: (event) => {
+			if (!isOpen) {
+				return;
+			}
+			for (const sub of store.open) {
+				sub(event);
+			}
+		},
+
+		publishError: (event) => {
+			if (!isOpen) {
+				return;
+			}
+			for (const sub of store.error) {
+				sub(event);
+			}
+		},
+
+		publishMessage: (event) => {
+			if (!isOpen) {
+				return;
+			}
+			for (const sub of store.message) {
+				sub(event);
+			}
+		},
+
+		publishClose: (event) => {
+			if (!isOpen) {
+				return;
+			}
+			for (const sub of store.close) {
+				sub(event);
+			}
+		},
+	};
+
+	return [mockSocket, publisher] as const;
+}
diff --git a/site/src/utils/OneWayWebSocket.test.ts b/site/src/utils/OneWayWebSocket.test.ts
