Allow overriding WebSocket implementation and add flaky version for testing

frontend/test-client/src/agent/mock-agent.ts

@@ -6,6 +6,8 @@ import { LogLevel } from "sync-client";
 import { MockClient } from "./mock-client";
 import { sleep } from "../utils/sleep";
 import type { LogLine } from "sync-client/dist/types/tracing/logger";
+import { flakyFetchFactory } from "../utils/flaky-fetch";
+import { flakyWebSocketFactory } from "../utils/flaky-websocket";
 
 export class MockAgent extends MockClient {
 	private readonly writtenContents: string[] = [];
@@ -26,16 +28,8 @@ export class MockAgent extends MockClient {
 
 	public async init(): Promise<void> {
 		await super.init(
-			// flaky fetch implementation to use during testing
-			async (
-				input: string | URL | globalThis.Request,
-				init?: RequestInit
-			): Promise<Response> => {
-				await sleep(Math.random() * this.jitterScaleInSeconds * 1000);
-				const response = await fetch(input, init);
-				await sleep(Math.random() * this.jitterScaleInSeconds * 1000);
-				return response;
-			}
+			flakyFetchFactory(this.jitterScaleInSeconds),
+			flakyWebSocketFactory(this.jitterScaleInSeconds)
 		);
 
 		assert(

frontend/test-client/src/agent/mock-client.ts

@@ -30,7 +30,8 @@ export class MockClient implements FileSystemOperations {
 	}
 
 	public async init(
-		fetchImplementation: typeof globalThis.fetch
+		fetchImplementation: typeof globalThis.fetch,
+		webSocketImplementation: typeof globalThis.WebSocket
 	): Promise<void> {
 		this.client = await SyncClient.create({
 			fs: this,
@@ -38,7 +39,8 @@ export class MockClient implements FileSystemOperations {
 				load: async () => this.data,
 				save: async (data) => void (this.data = data)
 			},
-			fetch: fetchImplementation
+			fetch: fetchImplementation,
+			webSocket: webSocketImplementation
 		});
 
 		await this.client.start();

frontend/test-client/src/utils/flaky-fetch.ts

@@ -0,0 +1,20 @@
+import { sleep } from "./sleep";
+
+export const flakyFetchFactory =
+	(jitterScaleInSeconds: number) =>
+	async (
+		input: string | URL | globalThis.Request,
+		init?: RequestInit
+	): Promise<Response> => {
+		if (jitterScaleInSeconds > 0) {
+			await sleep(Math.random() * jitterScaleInSeconds * 1000);
+		}
+
+		const response = await fetch(input, init);
+
+		if (jitterScaleInSeconds > 0) {
+			await sleep(Math.random() * jitterScaleInSeconds * 1000);
+		}
+
+		return response;
+	};

frontend/test-client/src/utils/flaky-websocket.ts

@@ -0,0 +1,61 @@
+import { sleep } from "./sleep";
+
+export function flakyWebSocketFactory(
+	jitterScaleInSeconds: number
+): typeof WebSocket {
+	// eslint-disable-next-line
+	return class FlakyWebSocket extends require("ws") {
+		public set onopen(callback: (event: Event) => void) {
+			// eslint-disable-next-line
+			super.onopen = async (event: Event): Promise<void> => {
+				if (jitterScaleInSeconds > 0) {
+					await sleep(Math.random() * jitterScaleInSeconds * 1000);
+				}
+
+				callback(event);
+			};
+		}
+
+		public set onmessage(callback: (event: MessageEvent) => void) {
+			// eslint-disable-next-line
+			super.onmessage = async (event: MessageEvent): Promise<void> => {
+				if (jitterScaleInSeconds > 0) {
+					await sleep(Math.random() * jitterScaleInSeconds * 1000);
+				}
+
+				callback(event);
+			};
+		}
+
+		public set onclose(callback: (event: CloseEvent) => void) {
+			// eslint-disable-next-line
+			super.onclose = async (event: CloseEvent): Promise<void> => {
+				if (jitterScaleInSeconds > 0) {
+					await sleep(Math.random() * jitterScaleInSeconds * 1000);
+				}
+				callback(event);
+			};
+		}
+
+		public set onerror(callback: (event: Event) => void) {
+			// eslint-disable-next-line
+			super.onerror = async (event: Event): Promise<void> => {
+				if (jitterScaleInSeconds > 0) {
+					await sleep(Math.random() * jitterScaleInSeconds * 1000);
+				}
+				callback(event);
+			};
+		}
+
+		public async send(
+			data: string | ArrayBufferLike | Blob | ArrayBufferView
+		): Promise<void> {
+			if (jitterScaleInSeconds > 0) {
+				await sleep(Math.random() * jitterScaleInSeconds * 1000);
+			}
+
+			// eslint-disable-next-line
+			super.send(data);
+		}
+	} as unknown as typeof WebSocket;
+}