From a87e93e3f12338d458439537cd0bfe509ced813d Mon Sep 17 00:00:00 2001
From: Michael Smith <throwawayclover@gmail.com>
Date: Mon, 4 Aug 2025 23:37:18 +0000
Subject: [PATCH 01/12] wip: commit progress on tests

---
 site/src/testHelpers/websockets.test.ts | 135 +++++++++++++++++++++++
 site/src/testHelpers/websockets.ts      | 138 ++++++++++++++++++++++++
 2 files changed, 273 insertions(+)
 create mode 100644 site/src/testHelpers/websockets.test.ts
 create mode 100644 site/src/testHelpers/websockets.ts

diff --git a/site/src/testHelpers/websockets.test.ts b/site/src/testHelpers/websockets.test.ts
new file mode 100644
index 0000000000000..a9ca5751680fc
--- /dev/null
+++ b/site/src/testHelpers/websockets.test.ts
@@ -0,0 +1,135 @@
+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 publisher to socket", () => {
+		const [socket, publisher] = 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");
+
+		publisher.publishOpen(openEvent);
+		publisher.publishError(errorEvent);
+		publisher.publishMessage(messageEvent);
+		publisher.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, publisher] = 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);
+			publisher.publishMessage(new MessageEvent<string>("message", {
+				"data": JSON.stringify(jd)
+			}));
+
+			expect(onMessage).toHaveBeenCalledTimes(1);
+			expect(onMessage).toHaveBeenCalledWith(new MessageEvent<string>("message", {
+				data: JSON.stringify(jd)
+			}));
+
+			socket.removeEventListener("message", onMessage);
+			onMessage.mockClear();
+		}
+	})
+
+	it("Only registers each socket event handler once", () => {
+		const [socket, publisher] = 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);
+
+		publisher.publishOpen(new Event("open"));
+		publisher.publishError(new Event("error"));
+		publisher.publishMessage(new MessageEvent<string>("message"));
+		publisher.publishClose(new CloseEvent("close"));
+
+		expect(onOpen).toHaveBeenCalledTimes(1);
+		expect(onError).toHaveBeenCalledTimes(1);
+		expect(onMessage).toHaveBeenCalledTimes(1);
+		expect(onClose).toHaveBeenCalledTimes(1);
+	});
+
+	it("Renders socket inert after being closed", () => {
+		const [socket, publisher] = createMockWebSocket("wss://www.dog.ceo/woof");
+		expect(publisher.isConnectionOpen).toBe(true);
+
+		const onMessage = jest.fn();
+		socket.addEventListener("message", onMessage);
+
+		socket.close();
+		expect(publisher.isConnectionOpen).toBe(false);
+
+		publisher.publishMessage(new MessageEvent<string>("message"));
+		expect(onMessage).not.toHaveBeenCalled();
+	});
+});
diff --git a/site/src/testHelpers/websockets.ts b/site/src/testHelpers/websockets.ts
new file mode 100644
index 0000000000000..418326f20a6fd
--- /dev/null
+++ b/site/src/testHelpers/websockets.ts
@@ -0,0 +1,138 @@
+import type { WebSocketEventType } from "utils/OneWayWebSocket";
+
+export type MockWebSocketPublisher = Readonly<{
+	publishMessage: (event: MessageEvent<string>) => void;
+	publishError: (event: Event) => void;
+	publishClose: (event: CloseEvent) => void;
+	publishOpen: (event: Event) => void;
+	readonly isConnectionOpen: boolean;
+}>;
+
+export function createMockWebSocket(
+	url: string,
+	protocols?: string | string[],
+): readonly [WebSocket, MockWebSocketPublisher] {
+	type EventMap = {
+		message: MessageEvent<string>;
+		error: Event;
+		close: CloseEvent;
+		open: Event;
+	};
+	type CallbackStore = {
+		[K in keyof EventMap]: ((event: EventMap[K]) => void)[];
+	};
+
+	let activeProtocol: string;
+	if (Array.isArray(protocols)) {
+		activeProtocol = protocols[0] ?? "";
+	} else if (typeof protocols === "string") {
+		activeProtocol = protocols;
+	} else {
+		activeProtocol = "";
+	}
+
+	let isOpen = true;
+	const store: CallbackStore = {
+		message: [],
+		error: [],
+		close: [],
+		open: [],
+	};
+
+	const mockSocket: WebSocket = {
+		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,
+		send: jest.fn(),
+		dispatchEvent: jest.fn(),
+
+		addEventListener: <E extends WebSocketEventType>(
+			eventType: E,
+			callback: (event: WebSocketEventMap[E]) => void,
+		) => {
+			if (!isOpen) {
+				return;
+			}
+
+			const subscribers = store[eventType];
+			if (!subscribers.includes(callback)) {
+				subscribers.push(callback);
+			}
+		},
+
+		removeEventListener: <E extends WebSocketEventType>(
+			eventType: E,
+			callback: (event: WebSocketEventMap[E]) => void,
+		) => {
+			if (!isOpen) {
+				return;
+			}
+
+			const subscribers = store[eventType];
+			if (subscribers.includes(callback)) {
+				const updated = store[eventType].filter((c) => c !== callback);
+				store[eventType] = updated as CallbackStore[E];
+			}
+		},
+
+		close: () => {
+			isOpen = false;
+		},
+	};
+
+	const publisher: MockWebSocketPublisher = {
+		get isConnectionOpen() {
+			return isOpen;
+		},
+
+		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;
+}

From 1c48afcf1f98f8323d1f2a4df22ad387b917fcc5 Mon Sep 17 00:00:00 2001
From: Michael Smith <throwawayclover@gmail.com>
Date: Tue, 5 Aug 2025 00:24:34 +0000
Subject: [PATCH 02/12] chore: get another test working

---
 site/src/testHelpers/websockets.test.ts | 33 +++++++++++++++++++++++++
 site/src/testHelpers/websockets.ts      | 11 ++++++---
 2 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/site/src/testHelpers/websockets.test.ts b/site/src/testHelpers/websockets.test.ts
index a9ca5751680fc..d937091e4d600 100644
--- a/site/src/testHelpers/websockets.test.ts
+++ b/site/src/testHelpers/websockets.test.ts
@@ -119,6 +119,35 @@ describe(createMockWebSocket.name, () => {
 		expect(onClose).toHaveBeenCalledTimes(1);
 	});
 
+	it("Lets a socket unsubscribe to event types", () => {
+		const [socket, publisher] = 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);
+
+		publisher.publishOpen(new Event("open"));
+		publisher.publishError(new Event("error"));
+		publisher.publishMessage(new MessageEvent<string>("message"));
+		publisher.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, publisher] = createMockWebSocket("wss://www.dog.ceo/woof");
 		expect(publisher.isConnectionOpen).toBe(true);
@@ -132,4 +161,8 @@ describe(createMockWebSocket.name, () => {
 		publisher.publishMessage(new MessageEvent<string>("message"));
 		expect(onMessage).not.toHaveBeenCalled();
 	});
+
+	it("Lets a socket send data to the mock server", () => {
+		expect.hasAssertions();
+	});
 });
diff --git a/site/src/testHelpers/websockets.ts b/site/src/testHelpers/websockets.ts
index 418326f20a6fd..d0288014a1c61 100644
--- a/site/src/testHelpers/websockets.ts
+++ b/site/src/testHelpers/websockets.ts
@@ -1,17 +1,18 @@
 import type { WebSocketEventType } from "utils/OneWayWebSocket";
 
-export type MockWebSocketPublisher = Readonly<{
+export type MockWebSocketServer = Readonly<{
 	publishMessage: (event: MessageEvent<string>) => void;
 	publishError: (event: Event) => void;
 	publishClose: (event: CloseEvent) => void;
 	publishOpen: (event: Event) => void;
+
 	readonly isConnectionOpen: boolean;
 }>;
 
 export function createMockWebSocket(
 	url: string,
 	protocols?: string | string[],
-): readonly [WebSocket, MockWebSocketPublisher] {
+): readonly [WebSocket, MockWebSocketServer] {
 	type EventMap = {
 		message: MessageEvent<string>;
 		error: Event;
@@ -22,6 +23,10 @@ export function createMockWebSocket(
 		[K in keyof EventMap]: ((event: EventMap[K]) => void)[];
 	};
 
+	if (!url.startsWith("ws://") && !url.startsWith("wss://")) {
+		throw new Error("URL must start with ws:// or wss://");
+	}
+
 	let activeProtocol: string;
 	if (Array.isArray(protocols)) {
 		activeProtocol = protocols[0] ?? "";
@@ -92,7 +97,7 @@ export function createMockWebSocket(
 		},
 	};
 
-	const publisher: MockWebSocketPublisher = {
+	const publisher: MockWebSocketServer = {
 		get isConnectionOpen() {
 			return isOpen;
 		},

From 5eff4dfe314e658465530ede9c64c4960924f213 Mon Sep 17 00:00:00 2001
From: Michael Smith <throwawayclover@gmail.com>
Date: Tue, 5 Aug 2025 00:50:06 +0000
Subject: [PATCH 03/12] chore: get all initial tests passing

---
 site/src/testHelpers/websockets.test.ts | 67 ++++++++++++++++---------
 site/src/testHelpers/websockets.ts      | 17 ++++++-
 2 files changed, 59 insertions(+), 25 deletions(-)

diff --git a/site/src/testHelpers/websockets.test.ts b/site/src/testHelpers/websockets.test.ts
index d937091e4d600..2b5cae97fee67 100644
--- a/site/src/testHelpers/websockets.test.ts
+++ b/site/src/testHelpers/websockets.test.ts
@@ -1,3 +1,4 @@
+import { subscribe } from "node:diagnostics_channel";
 import { createMockWebSocket } from "./websockets";
 
 describe(createMockWebSocket.name, () => {
@@ -13,8 +14,8 @@ describe(createMockWebSocket.name, () => {
 		}
 	})
 
-	it("Sends events from publisher to socket", () => {
-		const [socket, publisher] = createMockWebSocket("wss://www.dog.ceo/shake");
+	it("Sends events from server to socket", () => {
+		const [socket, server] = createMockWebSocket("wss://www.dog.ceo/shake");
 
 		const onOpen = jest.fn();
 		const onError = jest.fn();
@@ -31,10 +32,10 @@ describe(createMockWebSocket.name, () => {
 		const messageEvent = new MessageEvent<string>("message")
 		const closeEvent = new CloseEvent("close");
 
-		publisher.publishOpen(openEvent);
-		publisher.publishError(errorEvent);
-		publisher.publishMessage(messageEvent);
-		publisher.publishClose(closeEvent);
+		server.publishOpen(openEvent);
+		server.publishError(errorEvent);
+		server.publishMessage(messageEvent);
+		server.publishClose(closeEvent);
 
 		expect(onOpen).toHaveBeenCalledTimes(1);
 		expect(onOpen).toHaveBeenCalledWith(openEvent);
@@ -50,7 +51,7 @@ describe(createMockWebSocket.name, () => {
 	});
 
 	it("Sends JSON data to the socket for message events", () => {
-		const [socket, publisher] = createMockWebSocket("wss://www.dog.ceo/wag");
+		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,
@@ -74,7 +75,7 @@ describe(createMockWebSocket.name, () => {
 		];
 		for (const jd of jsonData) {
 			socket.addEventListener("message", onMessage);
-			publisher.publishMessage(new MessageEvent<string>("message", {
+			server.publishMessage(new MessageEvent<string>("message", {
 				"data": JSON.stringify(jd)
 			}));
 
@@ -89,7 +90,7 @@ describe(createMockWebSocket.name, () => {
 	})
 
 	it("Only registers each socket event handler once", () => {
-		const [socket, publisher] = createMockWebSocket("wss://www.dog.ceo/borf");
+		const [socket, server] = createMockWebSocket("wss://www.dog.ceo/borf");
 
 		const onOpen = jest.fn();
 		const onError = jest.fn();
@@ -108,10 +109,10 @@ describe(createMockWebSocket.name, () => {
 		socket.addEventListener("message", onMessage);
 		socket.addEventListener("close", onClose);
 
-		publisher.publishOpen(new Event("open"));
-		publisher.publishError(new Event("error"));
-		publisher.publishMessage(new MessageEvent<string>("message"));
-		publisher.publishClose(new CloseEvent("close"));
+		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);
@@ -120,7 +121,7 @@ describe(createMockWebSocket.name, () => {
 	});
 
 	it("Lets a socket unsubscribe to event types", () => {
-		const [socket, publisher] = createMockWebSocket("wss://www.dog.ceo/zoomies");
+		const [socket, server] = createMockWebSocket("wss://www.dog.ceo/zoomies");
 
 		const onOpen = jest.fn();
 		const onError = jest.fn();
@@ -137,10 +138,10 @@ describe(createMockWebSocket.name, () => {
 		socket.removeEventListener("message", onMessage);
 		socket.removeEventListener("close", onClose);
 
-		publisher.publishOpen(new Event("open"));
-		publisher.publishError(new Event("error"));
-		publisher.publishMessage(new MessageEvent<string>("message"));
-		publisher.publishClose(new CloseEvent("close"));
+		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();
@@ -149,20 +150,38 @@ describe(createMockWebSocket.name, () => {
 	});
 
 	it("Renders socket inert after being closed", () => {
-		const [socket, publisher] = createMockWebSocket("wss://www.dog.ceo/woof");
-		expect(publisher.isConnectionOpen).toBe(true);
+		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(publisher.isConnectionOpen).toBe(false);
+		expect(server.isConnectionOpen).toBe(false);
 
-		publisher.publishMessage(new MessageEvent<string>("message"));
+		server.publishMessage(new MessageEvent<string>("message"));
 		expect(onMessage).not.toHaveBeenCalled();
 	});
 
-	it("Lets a socket send data to the mock server", () => {
-		expect.hasAssertions();
+	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.socketSendArguments).toHaveLength(1);
+		expect(server.socketSendArguments).toEqual([data]);
+
+		socket.close();
+		socket.send(data);
+		expect(server.socketSendArguments).toHaveLength(1);
+		expect(server.socketSendArguments).toEqual([data]);
 	});
 });
diff --git a/site/src/testHelpers/websockets.ts b/site/src/testHelpers/websockets.ts
index d0288014a1c61..adde2606a7b2a 100644
--- a/site/src/testHelpers/websockets.ts
+++ b/site/src/testHelpers/websockets.ts
@@ -1,5 +1,7 @@
 import type { WebSocketEventType } from "utils/OneWayWebSocket";
 
+type SocketSendData = string | ArrayBufferLike | Blob | ArrayBufferView;
+
 export type MockWebSocketServer = Readonly<{
 	publishMessage: (event: MessageEvent<string>) => void;
 	publishError: (event: Event) => void;
@@ -7,6 +9,7 @@ export type MockWebSocketServer = Readonly<{
 	publishOpen: (event: Event) => void;
 
 	readonly isConnectionOpen: boolean;
+	readonly socketSendArguments: readonly SocketSendData[];
 }>;
 
 export function createMockWebSocket(
@@ -44,6 +47,8 @@ export function createMockWebSocket(
 		open: [],
 	};
 
+	let sendData: SocketSendData[] = [];
+
 	const mockSocket: WebSocket = {
 		CONNECTING: 0,
 		OPEN: 1,
@@ -60,9 +65,15 @@ export function createMockWebSocket(
 		onerror: null,
 		onmessage: null,
 		onopen: null,
-		send: jest.fn(),
 		dispatchEvent: jest.fn(),
 
+		send: (data) => {
+			if (!isOpen) {
+				return;
+			}
+			sendData.push(data);
+		},
+
 		addEventListener: <E extends WebSocketEventType>(
 			eventType: E,
 			callback: (event: WebSocketEventMap[E]) => void,
@@ -102,6 +113,10 @@ export function createMockWebSocket(
 			return isOpen;
 		},
 
+		get socketSendArguments() {
+			return [...sendData];
+		},
+
 		publishOpen: (event) => {
 			if (!isOpen) {
 				return;

From c2d07fbe2ffcf4d9352440861dfd95831f4b10c7 Mon Sep 17 00:00:00 2001
From: Michael Smith <throwawayclover@gmail.com>
Date: Tue, 5 Aug 2025 01:20:39 +0000
Subject: [PATCH 04/12] refactor: centralize socket implementation

---
 site/src/utils/OneWayWebSocket.test.ts | 153 ++-----------------------
 1 file changed, 8 insertions(+), 145 deletions(-)

diff --git a/site/src/utils/OneWayWebSocket.test.ts b/site/src/utils/OneWayWebSocket.test.ts
index c6b00b593111f..bdb8f9202ea6c 100644
--- a/site/src/utils/OneWayWebSocket.test.ts
+++ b/site/src/utils/OneWayWebSocket.test.ts
@@ -7,145 +7,8 @@
  * WebSocket class doesn't have any bugs, and can safely be mocked out.
  */
 
-import {
-	type OneWayMessageEvent,
-	OneWayWebSocket,
-	type WebSocketEventType,
-} from "./OneWayWebSocket";
-
-type MockPublisher = Readonly<{
-	publishMessage: (event: MessageEvent<string>) => void;
-	publishError: (event: ErrorEvent) => void;
-	publishClose: (event: CloseEvent) => void;
-	publishOpen: (event: Event) => void;
-}>;
-
-function createMockWebSocket(
-	url: string,
-	protocols?: string | string[],
-): readonly [WebSocket, MockPublisher] {
-	type EventMap = {
-		message: MessageEvent<string>;
-		error: ErrorEvent;
-		close: CloseEvent;
-		open: Event;
-	};
-	type CallbackStore = {
-		[K in keyof EventMap]: ((event: EventMap[K]) => void)[];
-	};
-
-	let activeProtocol: string;
-	if (Array.isArray(protocols)) {
-		activeProtocol = protocols[0] ?? "";
-	} else if (typeof protocols === "string") {
-		activeProtocol = protocols;
-	} else {
-		activeProtocol = "";
-	}
-
-	let closed = false;
-	const store: CallbackStore = {
-		message: [],
-		error: [],
-		close: [],
-		open: [],
-	};
-
-	const mockSocket: WebSocket = {
-		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,
-		send: jest.fn(),
-		dispatchEvent: jest.fn(),
-
-		addEventListener: <E extends WebSocketEventType>(
-			eventType: E,
-			callback: WebSocketEventMap[E],
-		) => {
-			if (closed) {
-				return;
-			}
-
-			const subscribers = store[eventType];
-			const cb = callback as unknown as CallbackStore[E][0];
-			if (!subscribers.includes(cb)) {
-				subscribers.push(cb);
-			}
-		},
-
-		removeEventListener: <E extends WebSocketEventType>(
-			eventType: E,
-			callback: WebSocketEventMap[E],
-		) => {
-			if (closed) {
-				return;
-			}
-
-			const subscribers = store[eventType];
-			const cb = callback as unknown as CallbackStore[E][0];
-			if (subscribers.includes(cb)) {
-				const updated = store[eventType].filter((c) => c !== cb);
-				store[eventType] = updated as unknown as CallbackStore[E];
-			}
-		},
-
-		close: () => {
-			closed = true;
-		},
-	};
-
-	const publisher: MockPublisher = {
-		publishOpen: (event) => {
-			if (closed) {
-				return;
-			}
-			for (const sub of store.open) {
-				sub(event);
-			}
-		},
-
-		publishError: (event) => {
-			if (closed) {
-				return;
-			}
-			for (const sub of store.error) {
-				sub(event);
-			}
-		},
-
-		publishMessage: (event) => {
-			if (closed) {
-				return;
-			}
-			for (const sub of store.message) {
-				sub(event);
-			}
-		},
-
-		publishClose: (event) => {
-			if (closed) {
-				return;
-			}
-			for (const sub of store.close) {
-				sub(event);
-			}
-		},
-	};
-
-	return [mockSocket, publisher] as const;
-}
+import { createMockWebSocket, MockWebSocketServer } from "testHelpers/websockets";
+import { type OneWayMessageEvent, OneWayWebSocket } from "./OneWayWebSocket";
 
 describe(OneWayWebSocket.name, () => {
 	const dummyRoute = "/api/v2/blah";
@@ -167,7 +30,7 @@ describe(OneWayWebSocket.name, () => {
 	});
 
 	it("Lets a consumer add an event listener of each type", () => {
-		let publisher!: MockPublisher;
+		let publisher!: MockWebSocketServer;
 		const oneWay = new OneWayWebSocket({
 			apiRoute: dummyRoute,
 			websocketInit: (url, protocols) => {
@@ -207,7 +70,7 @@ describe(OneWayWebSocket.name, () => {
 	});
 
 	it("Lets a consumer remove an event listener of each type", () => {
-		let publisher!: MockPublisher;
+		let publisher!: MockWebSocketServer;
 		const oneWay = new OneWayWebSocket({
 			apiRoute: dummyRoute,
 			websocketInit: (url, protocols) => {
@@ -252,7 +115,7 @@ describe(OneWayWebSocket.name, () => {
 	});
 
 	it("Only calls each callback once if callback is added multiple times", () => {
-		let publisher!: MockPublisher;
+		let publisher!: MockWebSocketServer;
 		const oneWay = new OneWayWebSocket({
 			apiRoute: dummyRoute,
 			websocketInit: (url, protocols) => {
@@ -294,7 +157,7 @@ describe(OneWayWebSocket.name, () => {
 	});
 
 	it("Lets consumers register multiple callbacks for each event type", () => {
-		let publisher!: MockPublisher;
+		let publisher!: MockWebSocketServer;
 		const oneWay = new OneWayWebSocket({
 			apiRoute: dummyRoute,
 			websocketInit: (url, protocols) => {
@@ -375,7 +238,7 @@ describe(OneWayWebSocket.name, () => {
 	});
 
 	it("Gives consumers pre-parsed versions of message events", () => {
-		let publisher!: MockPublisher;
+		let publisher!: MockWebSocketServer;
 		const oneWay = new OneWayWebSocket({
 			apiRoute: dummyRoute,
 			websocketInit: (url, protocols) => {
@@ -405,7 +268,7 @@ describe(OneWayWebSocket.name, () => {
 	});
 
 	it("Exposes parsing error if message payload could not be parsed as JSON", () => {
-		let publisher!: MockPublisher;
+		let publisher!: MockWebSocketServer;
 		const oneWay = new OneWayWebSocket({
 			apiRoute: dummyRoute,
 			websocketInit: (url, protocols) => {

From 21c9eb8c6791472c468bc85801d660b027c8941e Mon Sep 17 00:00:00 2001
From: Michael Smith <throwawayclover@gmail.com>
Date: Tue, 5 Aug 2025 01:22:44 +0000
Subject: [PATCH 05/12] chore: format

---
 site/src/testHelpers/websockets.test.ts | 29 ++++++++++++++-----------
 site/src/testHelpers/websockets.ts      |  2 +-
 site/src/utils/OneWayWebSocket.test.ts  |  5 ++++-
 3 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/site/src/testHelpers/websockets.test.ts b/site/src/testHelpers/websockets.test.ts
index 2b5cae97fee67..f0288f090bd0a 100644
--- a/site/src/testHelpers/websockets.test.ts
+++ b/site/src/testHelpers/websockets.test.ts
@@ -1,18 +1,17 @@
-import { subscribe } from "node:diagnostics_channel";
 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"
+			"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");
@@ -29,7 +28,7 @@ describe(createMockWebSocket.name, () => {
 
 		const openEvent = new Event("open");
 		const errorEvent = new Event("error");
-		const messageEvent = new MessageEvent<string>("message")
+		const messageEvent = new MessageEvent<string>("message");
 		const closeEvent = new CloseEvent("close");
 
 		server.publishOpen(openEvent);
@@ -56,7 +55,7 @@ describe(createMockWebSocket.name, () => {
 
 		// 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[] =[
+		const jsonData: readonly unknown[] = [
 			"blah",
 			42,
 			true,
@@ -71,23 +70,27 @@ describe(createMockWebSocket.name, () => {
 				profession: "Puzzle Solver",
 				sadBackstory: true,
 				greatVideoGames: true,
-			}
+			},
 		];
 		for (const jd of jsonData) {
 			socket.addEventListener("message", onMessage);
-			server.publishMessage(new MessageEvent<string>("message", {
-				"data": JSON.stringify(jd)
-			}));
+			server.publishMessage(
+				new MessageEvent<string>("message", {
+					data: JSON.stringify(jd),
+				}),
+			);
 
 			expect(onMessage).toHaveBeenCalledTimes(1);
-			expect(onMessage).toHaveBeenCalledWith(new MessageEvent<string>("message", {
-				data: JSON.stringify(jd)
-			}));
+			expect(onMessage).toHaveBeenCalledWith(
+				new MessageEvent<string>("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");
diff --git a/site/src/testHelpers/websockets.ts b/site/src/testHelpers/websockets.ts
index adde2606a7b2a..c3dcd5fdd9766 100644
--- a/site/src/testHelpers/websockets.ts
+++ b/site/src/testHelpers/websockets.ts
@@ -47,7 +47,7 @@ export function createMockWebSocket(
 		open: [],
 	};
 
-	let sendData: SocketSendData[] = [];
+	const sendData: SocketSendData[] = [];
 
 	const mockSocket: WebSocket = {
 		CONNECTING: 0,
diff --git a/site/src/utils/OneWayWebSocket.test.ts b/site/src/utils/OneWayWebSocket.test.ts
index bdb8f9202ea6c..3b674cc194fac 100644
--- a/site/src/utils/OneWayWebSocket.test.ts
+++ b/site/src/utils/OneWayWebSocket.test.ts
@@ -7,7 +7,10 @@
  * WebSocket class doesn't have any bugs, and can safely be mocked out.
  */
 
-import { createMockWebSocket, MockWebSocketServer } from "testHelpers/websockets";
+import {
+	type MockWebSocketServer,
+	createMockWebSocket,
+} from "testHelpers/websockets";
 import { type OneWayMessageEvent, OneWayWebSocket } from "./OneWayWebSocket";
 
 describe(OneWayWebSocket.name, () => {

From 1834c5908d9d82ba6800cce7f5009b6e8230a153 Mon Sep 17 00:00:00 2001
From: Michael Smith <throwawayclover@gmail.com>
Date: Tue, 5 Aug 2025 01:27:38 +0000
Subject: [PATCH 06/12] refactor: rename variable

---
 site/src/utils/OneWayWebSocket.test.ts | 72 +++++++++++++-------------
 1 file changed, 36 insertions(+), 36 deletions(-)

diff --git a/site/src/utils/OneWayWebSocket.test.ts b/site/src/utils/OneWayWebSocket.test.ts
index 3b674cc194fac..a0492dab9b439 100644
--- a/site/src/utils/OneWayWebSocket.test.ts
+++ b/site/src/utils/OneWayWebSocket.test.ts
@@ -33,12 +33,12 @@ describe(OneWayWebSocket.name, () => {
 	});
 
 	it("Lets a consumer add an event listener of each type", () => {
-		let publisher!: MockWebSocketServer;
+		let mockServer!: MockWebSocketServer;
 		const oneWay = new OneWayWebSocket({
 			apiRoute: dummyRoute,
 			websocketInit: (url, protocols) => {
-				const [socket, pub] = createMockWebSocket(url, protocols);
-				publisher = pub;
+				const [socket, server] = createMockWebSocket(url, protocols);
+				mockServer = server;
 				return socket;
 			},
 		});
@@ -53,14 +53,14 @@ describe(OneWayWebSocket.name, () => {
 		oneWay.addEventListener("error", onError);
 		oneWay.addEventListener("message", onMessage);
 
-		publisher.publishOpen(new Event("open"));
-		publisher.publishClose(new CloseEvent("close"));
-		publisher.publishError(
+		mockServer.publishOpen(new Event("open"));
+		mockServer.publishClose(new CloseEvent("close"));
+		mockServer.publishError(
 			new ErrorEvent("error", {
 				error: new Error("Whoops - connection broke"),
 			}),
 		);
-		publisher.publishMessage(
+		mockServer.publishMessage(
 			new MessageEvent("message", {
 				data: "null",
 			}),
@@ -73,12 +73,12 @@ describe(OneWayWebSocket.name, () => {
 	});
 
 	it("Lets a consumer remove an event listener of each type", () => {
-		let publisher!: MockWebSocketServer;
+		let mockServer!: MockWebSocketServer;
 		const oneWay = new OneWayWebSocket({
 			apiRoute: dummyRoute,
 			websocketInit: (url, protocols) => {
-				const [socket, pub] = createMockWebSocket(url, protocols);
-				publisher = pub;
+				const [socket, server] = createMockWebSocket(url, protocols);
+				mockServer = server;
 				return socket;
 			},
 		});
@@ -98,14 +98,14 @@ describe(OneWayWebSocket.name, () => {
 		oneWay.removeEventListener("error", onError);
 		oneWay.removeEventListener("message", onMessage);
 
-		publisher.publishOpen(new Event("open"));
-		publisher.publishClose(new CloseEvent("close"));
-		publisher.publishError(
+		mockServer.publishOpen(new Event("open"));
+		mockServer.publishClose(new CloseEvent("close"));
+		mockServer.publishError(
 			new ErrorEvent("error", {
 				error: new Error("Whoops - connection broke"),
 			}),
 		);
-		publisher.publishMessage(
+		mockServer.publishMessage(
 			new MessageEvent("message", {
 				data: "null",
 			}),
@@ -118,12 +118,12 @@ describe(OneWayWebSocket.name, () => {
 	});
 
 	it("Only calls each callback once if callback is added multiple times", () => {
-		let publisher!: MockWebSocketServer;
+		let mockServer!: MockWebSocketServer;
 		const oneWay = new OneWayWebSocket({
 			apiRoute: dummyRoute,
 			websocketInit: (url, protocols) => {
-				const [socket, pub] = createMockWebSocket(url, protocols);
-				publisher = pub;
+				const [socket, server] = createMockWebSocket(url, protocols);
+				mockServer = server;
 				return socket;
 			},
 		});
@@ -140,14 +140,14 @@ describe(OneWayWebSocket.name, () => {
 			oneWay.addEventListener("message", onMessage);
 		}
 
-		publisher.publishOpen(new Event("open"));
-		publisher.publishClose(new CloseEvent("close"));
-		publisher.publishError(
+		mockServer.publishOpen(new Event("open"));
+		mockServer.publishClose(new CloseEvent("close"));
+		mockServer.publishError(
 			new ErrorEvent("error", {
 				error: new Error("Whoops - connection broke"),
 			}),
 		);
-		publisher.publishMessage(
+		mockServer.publishMessage(
 			new MessageEvent("message", {
 				data: "null",
 			}),
@@ -160,12 +160,12 @@ describe(OneWayWebSocket.name, () => {
 	});
 
 	it("Lets consumers register multiple callbacks for each event type", () => {
-		let publisher!: MockWebSocketServer;
+		let mockServer!: MockWebSocketServer;
 		const oneWay = new OneWayWebSocket({
 			apiRoute: dummyRoute,
 			websocketInit: (url, protocols) => {
-				const [socket, pub] = createMockWebSocket(url, protocols);
-				publisher = pub;
+				const [socket, server] = createMockWebSocket(url, protocols);
+				mockServer = server;
 				return socket;
 			},
 		});
@@ -188,14 +188,14 @@ describe(OneWayWebSocket.name, () => {
 		oneWay.addEventListener("error", onError2);
 		oneWay.addEventListener("message", onMessage2);
 
-		publisher.publishOpen(new Event("open"));
-		publisher.publishClose(new CloseEvent("close"));
-		publisher.publishError(
+		mockServer.publishOpen(new Event("open"));
+		mockServer.publishClose(new CloseEvent("close"));
+		mockServer.publishError(
 			new ErrorEvent("error", {
 				error: new Error("Whoops - connection broke"),
 			}),
 		);
-		publisher.publishMessage(
+		mockServer.publishMessage(
 			new MessageEvent("message", {
 				data: "null",
 			}),
@@ -241,12 +241,12 @@ describe(OneWayWebSocket.name, () => {
 	});
 
 	it("Gives consumers pre-parsed versions of message events", () => {
-		let publisher!: MockWebSocketServer;
+		let mockServer!: MockWebSocketServer;
 		const oneWay = new OneWayWebSocket({
 			apiRoute: dummyRoute,
 			websocketInit: (url, protocols) => {
-				const [socket, pub] = createMockWebSocket(url, protocols);
-				publisher = pub;
+				const [socket, server] = createMockWebSocket(url, protocols);
+				mockServer = server;
 				return socket;
 			},
 		});
@@ -262,7 +262,7 @@ describe(OneWayWebSocket.name, () => {
 			data: JSON.stringify(payload),
 		});
 
-		publisher.publishMessage(event);
+		mockServer.publishMessage(event);
 		expect(onMessage).toHaveBeenCalledWith({
 			sourceEvent: event,
 			parsedMessage: payload,
@@ -271,12 +271,12 @@ describe(OneWayWebSocket.name, () => {
 	});
 
 	it("Exposes parsing error if message payload could not be parsed as JSON", () => {
-		let publisher!: MockWebSocketServer;
+		let mockServer!: MockWebSocketServer;
 		const oneWay = new OneWayWebSocket({
 			apiRoute: dummyRoute,
 			websocketInit: (url, protocols) => {
-				const [socket, pub] = createMockWebSocket(url, protocols);
-				publisher = pub;
+				const [socket, server] = createMockWebSocket(url, protocols);
+				mockServer = server;
 				return socket;
 			},
 		});
@@ -288,7 +288,7 @@ describe(OneWayWebSocket.name, () => {
 		const event = new MessageEvent("message", {
 			data: payload,
 		});
-		publisher.publishMessage(event);
+		mockServer.publishMessage(event);
 
 		const arg: OneWayMessageEvent<never> = onMessage.mock.lastCall[0];
 		expect(arg.sourceEvent).toEqual(event);

From 6ee9afc035ae4a0cfbda4091f64dab84fd4ecb61 Mon Sep 17 00:00:00 2001
From: Michael Smith <throwawayclover@gmail.com>
Date: Tue, 5 Aug 2025 01:46:00 +0000
Subject: [PATCH 07/12] refactor: rename property

---
 site/src/testHelpers/websockets.test.ts |  8 ++++----
 site/src/testHelpers/websockets.ts      | 10 +++++-----
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/site/src/testHelpers/websockets.test.ts b/site/src/testHelpers/websockets.test.ts
index f0288f090bd0a..eceefbc00b638 100644
--- a/site/src/testHelpers/websockets.test.ts
+++ b/site/src/testHelpers/websockets.test.ts
@@ -179,12 +179,12 @@ describe(createMockWebSocket.name, () => {
 		});
 
 		socket.send(data);
-		expect(server.socketSendArguments).toHaveLength(1);
-		expect(server.socketSendArguments).toEqual([data]);
+		expect(server.clientSentData).toHaveLength(1);
+		expect(server.clientSentData).toEqual([data]);
 
 		socket.close();
 		socket.send(data);
-		expect(server.socketSendArguments).toHaveLength(1);
-		expect(server.socketSendArguments).toEqual([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
index c3dcd5fdd9766..5792fa945b52b 100644
--- a/site/src/testHelpers/websockets.ts
+++ b/site/src/testHelpers/websockets.ts
@@ -9,7 +9,7 @@ export type MockWebSocketServer = Readonly<{
 	publishOpen: (event: Event) => void;
 
 	readonly isConnectionOpen: boolean;
-	readonly socketSendArguments: readonly SocketSendData[];
+	readonly clientSentData: readonly SocketSendData[];
 }>;
 
 export function createMockWebSocket(
@@ -47,7 +47,7 @@ export function createMockWebSocket(
 		open: [],
 	};
 
-	const sendData: SocketSendData[] = [];
+	const sentData: SocketSendData[] = [];
 
 	const mockSocket: WebSocket = {
 		CONNECTING: 0,
@@ -71,7 +71,7 @@ export function createMockWebSocket(
 			if (!isOpen) {
 				return;
 			}
-			sendData.push(data);
+			sentData.push(data);
 		},
 
 		addEventListener: <E extends WebSocketEventType>(
@@ -113,8 +113,8 @@ export function createMockWebSocket(
 			return isOpen;
 		},
 
-		get socketSendArguments() {
-			return [...sendData];
+		get clientSentData() {
+			return [...sentData];
 		},
 
 		publishOpen: (event) => {

From cb42aa7857306cc776187eb750cc27b86c79c3fe Mon Sep 17 00:00:00 2001
From: Michael Smith <throwawayclover@gmail.com>
Date: Thu, 7 Aug 2025 20:06:13 +0000
Subject: [PATCH 08/12] fix: use sets

---
 site/src/testHelpers/websockets.ts | 53 +++++++++++-------------------
 1 file changed, 19 insertions(+), 34 deletions(-)

diff --git a/site/src/testHelpers/websockets.ts b/site/src/testHelpers/websockets.ts
index 5792fa945b52b..cb2bb9f662510 100644
--- a/site/src/testHelpers/websockets.ts
+++ b/site/src/testHelpers/websockets.ts
@@ -12,39 +12,31 @@ export type MockWebSocketServer = Readonly<{
 	readonly clientSentData: readonly SocketSendData[];
 }>;
 
+type EventMap = {
+	message: MessageEvent<string>;
+	error: Event;
+	close: CloseEvent;
+	open: Event;
+};
+
+type CallbackStore = {
+	[K in keyof EventMap]: Set<(event: EventMap[K]) => void>;
+};
+
 export function createMockWebSocket(
 	url: string,
-	protocols?: string | string[],
+	protocol?: string,
 ): readonly [WebSocket, MockWebSocketServer] {
-	type EventMap = {
-		message: MessageEvent<string>;
-		error: Event;
-		close: CloseEvent;
-		open: Event;
-	};
-	type CallbackStore = {
-		[K in keyof EventMap]: ((event: EventMap[K]) => void)[];
-	};
-
 	if (!url.startsWith("ws://") && !url.startsWith("wss://")) {
 		throw new Error("URL must start with ws:// or wss://");
 	}
 
-	let activeProtocol: string;
-	if (Array.isArray(protocols)) {
-		activeProtocol = protocols[0] ?? "";
-	} else if (typeof protocols === "string") {
-		activeProtocol = protocols;
-	} else {
-		activeProtocol = "";
-	}
-
 	let isOpen = true;
 	const store: CallbackStore = {
-		message: [],
-		error: [],
-		close: [],
-		open: [],
+		message: new Set(),
+		error: new Set(),
+		close: new Set(),
+		open: new Set(),
 	};
 
 	const sentData: SocketSendData[] = [];
@@ -56,7 +48,7 @@ export function createMockWebSocket(
 		CLOSED: 3,
 
 		url,
-		protocol: activeProtocol,
+		protocol: protocol ?? "",
 		readyState: 1,
 		binaryType: "blob",
 		bufferedAmount: 0,
@@ -81,11 +73,8 @@ export function createMockWebSocket(
 			if (!isOpen) {
 				return;
 			}
-
 			const subscribers = store[eventType];
-			if (!subscribers.includes(callback)) {
-				subscribers.push(callback);
-			}
+			subscribers.add(callback);
 		},
 
 		removeEventListener: <E extends WebSocketEventType>(
@@ -95,12 +84,8 @@ export function createMockWebSocket(
 			if (!isOpen) {
 				return;
 			}
-
 			const subscribers = store[eventType];
-			if (subscribers.includes(callback)) {
-				const updated = store[eventType].filter((c) => c !== callback);
-				store[eventType] = updated as CallbackStore[E];
-			}
+			subscribers.delete(callback);
 		},
 
 		close: () => {

From b63ff10ffa4b6677b89f875b3861d08d0d5a40ed Mon Sep 17 00:00:00 2001
From: Michael Smith <throwawayclover@gmail.com>
Date: Thu, 7 Aug 2025 20:17:04 +0000
Subject: [PATCH 09/12] refactor: clean up types

---
 site/src/testHelpers/websockets.test.ts |  8 ++-----
 site/src/testHelpers/websockets.ts      | 31 +++++++++++++++----------
 2 files changed, 21 insertions(+), 18 deletions(-)

diff --git a/site/src/testHelpers/websockets.test.ts b/site/src/testHelpers/websockets.test.ts
index eceefbc00b638..edd4191cffebe 100644
--- a/site/src/testHelpers/websockets.test.ts
+++ b/site/src/testHelpers/websockets.test.ts
@@ -75,16 +75,12 @@ describe(createMockWebSocket.name, () => {
 		for (const jd of jsonData) {
 			socket.addEventListener("message", onMessage);
 			server.publishMessage(
-				new MessageEvent<string>("message", {
-					data: JSON.stringify(jd),
-				}),
+				new MessageEvent("message", { data: JSON.stringify(jd) }),
 			);
 
 			expect(onMessage).toHaveBeenCalledTimes(1);
 			expect(onMessage).toHaveBeenCalledWith(
-				new MessageEvent<string>("message", {
-					data: JSON.stringify(jd),
-				}),
+				new MessageEvent("message", { data: JSON.stringify(jd) }),
 			);
 
 			socket.removeEventListener("message", onMessage);
diff --git a/site/src/testHelpers/websockets.ts b/site/src/testHelpers/websockets.ts
index cb2bb9f662510..3f473ead0711f 100644
--- a/site/src/testHelpers/websockets.ts
+++ b/site/src/testHelpers/websockets.ts
@@ -1,6 +1,6 @@
 import type { WebSocketEventType } from "utils/OneWayWebSocket";
 
-type SocketSendData = string | ArrayBufferLike | Blob | ArrayBufferView;
+type SocketSendData = Parameters<WebSocket["send"]>[0];
 
 export type MockWebSocketServer = Readonly<{
 	publishMessage: (event: MessageEvent<string>) => void;
@@ -12,21 +12,28 @@ export type MockWebSocketServer = Readonly<{
 	readonly clientSentData: readonly SocketSendData[];
 }>;
 
-type EventMap = {
-	message: MessageEvent<string>;
-	error: Event;
-	close: CloseEvent;
-	open: Event;
+type CallbackStore = {
+	[K in keyof WebSocketEventMap]: Set<(event: WebSocketEventMap[K]) => void>;
 };
 
-type CallbackStore = {
-	[K in keyof EventMap]: Set<(event: EventMap[K]) => void>;
+export type MockWebSocket = Omit<WebSocket, "send"> & {
+	/**
+	 * A version of the WebSocket `send` method that has been pre-wrapped inside
+	 * a Jest mock. Most of the time, you should not be using the Jest mock
+	 * features, and should instead be using the `clientSentData` property from
+	 * the `MockWebSocketServer` type.
+	 *
+	 * The Jest mock functionality only exists to help make sure that the
+	 * client-side socket method got called. It does nothing to make sure that
+	 * the mock server actually received anything.
+	 */
+	send: jest.Mock<void, [SocketSendData], unknown>;
 };
 
 export function createMockWebSocket(
 	url: string,
 	protocol?: string,
-): readonly [WebSocket, MockWebSocketServer] {
+): readonly [MockWebSocket, MockWebSocketServer] {
 	if (!url.startsWith("ws://") && !url.startsWith("wss://")) {
 		throw new Error("URL must start with ws:// or wss://");
 	}
@@ -41,7 +48,7 @@ export function createMockWebSocket(
 
 	const sentData: SocketSendData[] = [];
 
-	const mockSocket: WebSocket = {
+	const mockSocket: MockWebSocket = {
 		CONNECTING: 0,
 		OPEN: 1,
 		CLOSING: 2,
@@ -59,12 +66,12 @@ export function createMockWebSocket(
 		onopen: null,
 		dispatchEvent: jest.fn(),
 
-		send: (data) => {
+		send: jest.fn((data) => {
 			if (!isOpen) {
 				return;
 			}
 			sentData.push(data);
-		},
+		}),
 
 		addEventListener: <E extends WebSocketEventType>(
 			eventType: E,

From dab74d71cc6a536f2386056c4bf9d95f3a02d71c Mon Sep 17 00:00:00 2001
From: Michael Smith <throwawayclover@gmail.com>
Date: Thu, 7 Aug 2025 20:24:12 +0000
Subject: [PATCH 10/12] docs: beef up comment

---
 site/src/testHelpers/websockets.ts | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/site/src/testHelpers/websockets.ts b/site/src/testHelpers/websockets.ts
index 3f473ead0711f..6b975e3144318 100644
--- a/site/src/testHelpers/websockets.ts
+++ b/site/src/testHelpers/websockets.ts
@@ -19,13 +19,21 @@ type CallbackStore = {
 export type MockWebSocket = Omit<WebSocket, "send"> & {
 	/**
 	 * A version of the WebSocket `send` method that has been pre-wrapped inside
-	 * a Jest mock. Most of the time, you should not be using the Jest mock
-	 * features, and should instead be using the `clientSentData` property from
-	 * the `MockWebSocketServer` type.
+	 * a Jest mock.
 	 *
-	 * The Jest mock functionality only exists to help make sure that the
-	 * client-side socket method got called. It does nothing to make sure that
-	 * the mock server actually received anything.
+	 * 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>;
 };

From c17ee61a3433488754c71af06e3d15f8de2d56ba Mon Sep 17 00:00:00 2001
From: Michael Smith <throwawayclover@gmail.com>
Date: Thu, 7 Aug 2025 20:34:35 +0000
Subject: [PATCH 11/12] fix: update protocol back

---
 site/src/testHelpers/websockets.ts | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/site/src/testHelpers/websockets.ts b/site/src/testHelpers/websockets.ts
index 6b975e3144318..010f8bf8c7664 100644
--- a/site/src/testHelpers/websockets.ts
+++ b/site/src/testHelpers/websockets.ts
@@ -40,12 +40,16 @@ export type MockWebSocket = Omit<WebSocket, "send"> & {
 
 export function createMockWebSocket(
 	url: string,
-	protocol?: 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(),
@@ -63,7 +67,7 @@ export function createMockWebSocket(
 		CLOSED: 3,
 
 		url,
-		protocol: protocol ?? "",
+		protocol: activeProtocol,
 		readyState: 1,
 		binaryType: "blob",
 		bufferedAmount: 0,

From 5f9032d09a2951c5a159a36d44c7a0e8ecbe9db2 Mon Sep 17 00:00:00 2001
From: Michael Smith <throwawayclover@gmail.com>
Date: Thu, 7 Aug 2025 20:46:13 +0000
Subject: [PATCH 12/12] fix: knip

---
 site/src/testHelpers/websockets.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/site/src/testHelpers/websockets.ts b/site/src/testHelpers/websockets.ts
index 010f8bf8c7664..57584cd55e887 100644
--- a/site/src/testHelpers/websockets.ts
+++ b/site/src/testHelpers/websockets.ts
@@ -16,7 +16,7 @@ type CallbackStore = {
 	[K in keyof WebSocketEventMap]: Set<(event: WebSocketEventMap[K]) => void>;
 };
 
-export type MockWebSocket = Omit<WebSocket, "send"> & {
+type MockWebSocket = Omit<WebSocket, "send"> & {
 	/**
 	 * A version of the WebSocket `send` method that has been pre-wrapped inside
 	 * a Jest mock.