Intercept and transform websocket messages

transformWebsocketMessage.helper.ts

import { Page } from '@playwright/test';

export async function transformWebsocketMessage(
  page: Page,
  fn: <TMessage extends Record<string, unknown>>(
    url: string,
    message: TMessage
  ) => TMessage
) {
  fn = fn || ((url, message) => message);
  return page.evaluate((fn) => {
    const MESSAGE_DELIMITER = '\u001e';
    const transformFn = eval(`(${fn})`) as (
      url: string,
      message: Record<string, unknown>
    ) => Record<string, unknown>;
    const OriginalWebsocket = window.WebSocket;
    const ProxiedWebSocket = function () {
      const url = arguments[0];
      console.warn('Intercepting web socket creation', ...arguments);

      // @ts-ignore
      const ws = new OriginalWebsocket(...arguments);

      const tryJson = (data: string) => {
        try {
          const parsed = JSON.parse(
            data?.trim().replace(MESSAGE_DELIMITER, '')
          );
          return [true, parsed];
        } catch (err) {
          // @ts-ignore
          console.error('JSON.parse error', err.toString());
          return [false, data];
        }
      };

      const originalAddEventListener = ws.addEventListener;
      const proxiedAddEventListener = function () {
        console.warn('Intercepting addEventListener', ...arguments);
        if (arguments[0] === 'message') {
          const cb = arguments[1];
          arguments[1] = function (e: { data: string }) {
            // Here you can get the actual data from the incoming messages
            // Here you can even change the data before calling the real message listener
            const [isJson, parsed] = tryJson(e.data);
            const originalData = isJson ? parsed : e.data;
            const transformedData = isJson
              ? transformFn(url, originalData)
              : originalData;
            console.warn(
              'transformed',
              JSON.stringify({
                originalData,
                transformedData,
              })
            );
            if (isJson && Object.keys(transformedData).length) {
              Object.defineProperty(e, 'data', {
                value: JSON.stringify(transformedData) + MESSAGE_DELIMITER,
              });
            }
            return cb.apply(this, arguments);
          };
        }
        // @ts-ignore
        return originalAddEventListener.apply(this, arguments);
      };
      ws.addEventListener = proxiedAddEventListener;

      Object.defineProperty(ws, 'onmessage', {
        set(func) {
          // @ts-ignore
          return proxiedAddEventListener.apply(this, ['message', func, false]);
        },
      });
      return ws;
    };
    Object.defineProperty(ProxiedWebSocket, 'OPEN', { value: 1 });
    // @ts-ignore
    window.WebSocket = ProxiedWebSocket;
    console.warn('WebSocket intercepted');
  }, fn.toString());
}