Add support for websockets

package.json

@@ -19,6 +19,7 @@
     "@types/esprima": "^4.0.6",
     "@types/mime-types": "^2.1.4",
     "@types/node": "^18.14.6",
+    "@types/ws": "^8.5.12",
     "@typescript-eslint/eslint-plugin": "^6.14.0",
     "@typescript-eslint/parser": "^6.14.0",
     "eslint": "^8.55.0",
@@ -40,7 +41,9 @@
     "esprima": "^4.0.1",
     "mime-types": "^2.1.35",
     "path-to-regexp": "^6.2.1",
-    "vue": "^3.3.8"
+    "portfinder": "^1.0.32",
+    "vue": "^3.3.8",
+    "ws": "^8.18.0"
   },
   "lint-staged": {
     "*.ts": [

src/index.ts

@@ -1,18 +1,21 @@
+import { reactive, ref } from 'vue';
 import type { InterceptOptions } from './types/options/intercept';
 import type { InterceptSurface } from './types/intercept-surface';
 import type { InterceptSetup } from './types/intercept-setup';
 import type { GlobalOptions } from './types/options/global';
+import type { WebsocketOptions } from './types/options/websocket';
+import type { WaitOptions } from './types/options/wait';
+import type { UpdateOptions } from './types/options/update';
+import type { InterceptWebsocketSurface } from './types/intercept-websocket-surface';
 import type { Page, Route, Request } from '@playwright/test';
 import { routeHandler } from './lib/route-handler';
 import { createMatchFunction } from './lib/create-match-function';
 import { getMimeType } from './lib/get-mime-type';
 import { normalizeBody } from './lib/normalize-body';
-import { reactive, ref } from 'vue';
-import type { WaitOptions } from './types/options/wait';
-import type { UpdateOptions } from './types/options/update';
 import { wait } from './lib/wait';
 import { update } from './lib/update';
 import { getRouteRegex } from './lib/get-route-regex';
+import { websocketSetup } from './lib/websocket-setup';
 
 export { createMatchFunction, getMimeType, normalizeBody };
 
@@ -110,6 +113,10 @@ export class Intercept {
   delete(options: InterceptOptions): InterceptSurface {
     return this.setup(options, 'DELETE');
   }
+
+  async wss(options: WebsocketOptions): Promise<InterceptWebsocketSurface> {
+    return await websocketSetup(this.page, options);
+  }
 }
 
 export type { InterceptOptions, InterceptSurface };

src/lib/route-handler.ts

@@ -33,7 +33,7 @@ export async function routeHandler({
 }: RouteHandlerParams) {
   const intercept = intercepts.find(
     (intercept) =>
-      intercept.match(route.request().url()) &&
+      intercept.match?.(route.request().url()) &&
       intercept.method === route.request().method()
   );
 
@@ -46,10 +46,10 @@ export async function routeHandler({
   // Ensure the response has been downloaded by client and then
   // as an anti-flake measure, push the request upon next tick.
   page.once('requestfinished', () => {
-    nextTick(() => intercept.requests.push(route.request()));
+    nextTick(() => intercept.requests?.push(route.request()));
   });
 
-  const { params } = intercept.match(route.request().url()) as MatchResult;
+  const { params } = intercept.match?.(route.request().url()) as MatchResult;
 
   if ('body' in intercept || 'fixture' in intercept) {
     let body = '';

src/lib/wait.ts

@@ -15,20 +15,22 @@ export function wait(
   return expect
     .poll(
       () => {
+        const requests = intercept.wssPayloads ?? intercept.requests;
+
         const hasFirstRequest = Boolean(
           intercept.previousRequestCount.value === 0 &&
-            intercept.requests.length > 0
+            requests!.length > 0
         );
 
         const hasNewRequests = Boolean(
           intercept.previousRequestCount.value > 0 &&
-            intercept.previousRequestCount.value < intercept.requests.length
+            intercept.previousRequestCount.value < requests!.length
         );
 
         const shouldPass = Boolean(hasFirstRequest || hasNewRequests);
 
         if (shouldPass) {
-          intercept.previousRequestCount.value = intercept.requests.length;
+          intercept.previousRequestCount.value = requests!.length;
           return true;
         }
 

src/lib/websocket-setup.ts

@@ -0,0 +1,74 @@
+import type { Page } from '@playwright/test';
+import type { WebsocketOptions } from '../types/options/websocket';
+import type { InterceptWebsocketSurface } from '../types/intercept-websocket-surface';
+import WebSocket, { WebSocketServer } from 'ws';
+import portfinder from 'portfinder';
+import type { InterceptSetup } from '../types/intercept-setup';
+import { reactive, ref } from 'vue';
+import { wait } from './wait';
+import type { WaitOptions } from '../types/options/wait';
+
+export async function websocketSetup(page: Page, options: WebsocketOptions): Promise<InterceptWebsocketSurface> {
+  const interceptSetup = <InterceptSetup>{
+    wssPayloads: reactive([]),
+    previousRequestCount: ref(0),
+    method: 'WSS',
+    ...options,
+  };
+
+  const port = await portfinder.getPortPromise();
+
+  const localWssUrl = `ws://localhost:${port}`;
+
+  // @todo: ...ugh
+  await page.route("**", async (route) => {
+    try {
+      const response = await route.fetch();
+      const pageHtml = await response.text();
+
+      if (!pageHtml.includes(options.url)) {
+        route.continue();
+        return;
+      }
+
+      const body = pageHtml.replaceAll(options.url, localWssUrl);
+      const headers = response.headers();
+      const csp = headers['content-security-policy'];
+      headers['content-security-policy'] = csp?.replace(
+        "connect-src 'self'",
+        "connect-src 'self' localhost:* ws://localhost:*"
+      );
+
+      await route.fulfill({ response, body, headers });
+    } catch (e) {
+      try {
+        route?.continue();
+      } catch {
+        //
+      }
+    }
+  });
+
+  const wss = new WebSocketServer({ port });
+
+  let send: (message: string) => void;
+
+  wss.on('connection', (ws: WebSocket) => {
+    send = ws.send.bind(ws);
+
+    const handshakeMessage = 'intercept:handshake';
+    interceptSetup.wssPayloads?.push(handshakeMessage);
+    options.handler({ message: handshakeMessage, send });
+
+    ws.on('message', (data) => {
+      interceptSetup.wssPayloads?.push(data.toString());
+      options.handler({ message: data.toString(), send });
+    });
+  });
+
+  return {
+    wait: (waitOptions?: WaitOptions) => wait(interceptSetup, waitOptions),
+    send: (message: string) => send(message),
+    wssPayloads: interceptSetup.wssPayloads!,
+  }
+}

src/types/intercept-setup.ts

@@ -8,8 +8,9 @@ import type { Ref } from 'vue';
  */
 
 export type InterceptSetup = InterceptOptions & {
-  method: 'GET' | 'POST' | 'PUT' | 'PATCH' | 'DELETE';
-  match: MatchFunction;
-  requests: Request[];
+  method: 'GET' | 'POST' | 'PUT' | 'PATCH' | 'DELETE' | 'WSS';
+  match?: MatchFunction;
+  requests?: Request[];
+  wssPayloads?: string[];
   previousRequestCount: Ref<number>;
 };

src/types/intercept-websocket-surface.ts

@@ -0,0 +1,7 @@
+import type { WaitOptions } from './options/wait';
+
+export interface InterceptWebsocketSurface {
+  wait: (options?: WaitOptions | undefined) => Promise<void>;
+  send: (message: string) => void;
+  wssPayloads: string[];
+}

src/types/options/intercept.ts

@@ -33,13 +33,22 @@ type FixtureOption = {
     | ((args: { route: Route; params: Record<string, string> }) => string);
 };
 
+export type WebsocketHandlerArgs = {
+  message?: string;
+  send?: (message: string) => void;
+};
+
+type WebsocketHandlerOption = {
+  handler: (args: WebsocketHandlerArgs) => void;
+};
+
 type HandlerOption = {
   handler: (args: {
     route: Route;
     params: Record<string, string>;
     request: Request;
   }) => void;
-};
+} | WebsocketHandlerOption;
 
 export type InterceptOptions = BaseOptions &
   (

src/types/options/websocket.ts

@@ -0,0 +1,6 @@
+import type { WebsocketHandlerArgs } from "./intercept";
+
+export interface WebsocketOptions {
+  url: string;
+  handler: (args: WebsocketHandlerArgs) => void;
+}