feat(ws): outgoing events forwarding

package.json

@@ -124,6 +124,7 @@
     "@types/node": "^16.11.26",
     "@types/node-fetch": "2.5.12",
     "@types/supertest": "^2.0.11",
+    "@types/ws": "^8.5.10",
     "axios": "^1.6.0",
     "body-parser": "^1.19.0",
     "commitizen": "^4.2.4",
@@ -148,7 +149,8 @@
     "vitest-environment-miniflare": "^2.14.1",
     "wait-for-expect": "^3.0.2",
     "web-encoding": "^1.1.5",
-    "webpack-http-server": "^0.5.0"
+    "webpack-http-server": "^0.5.0",
+    "ws": "^8.16.0"
   },
   "dependencies": {
     "@open-draft/deferred-promise": "^2.2.0",

src/interceptors/WebSocket/WebSocketServer.ts

@@ -8,9 +8,9 @@ import type { WebSocketMessageListener } from './implementations/WebSocketClass/
  */
 export class WebSocketServer {
   /**
-   * Connect to the actual WebSocket server.
+   * Connect to the original WebSocket server.
    */
-  public connect(): Promise<void> {
+  public connect(): void {
     throw new Error('WebSocketServer#connect is not implemented')
   }
 

src/interceptors/WebSocket/implementations/WebSocketClass/WebSocketClassInterceptor.ts

@@ -3,10 +3,7 @@ import type { WebSocketEventsMap } from '../../index'
 import { Interceptor } from '../../../../Interceptor'
 import { WebSocketClassClient } from './WebSocketClassClient'
 import { WebSocketClassServer } from './WebSocketClassServer'
-import type {
-  WebSocketSendData,
-  WebSocketTransportOnIncomingCallback,
-} from '../../WebSocketTransport'
+import type { WebSocketSendData } from '../../WebSocketTransport'
 import { bindEvent } from '../../utils/bindEvent'
 import { WebSocketClassTransport } from './WebSocketClassTransport'
 
@@ -49,7 +46,7 @@ export class WebSocketClassInterceptor extends Interceptor<WebSocketEventsMap> {
         // as soon as the WebSocket instance is constructed.
         this.emitter.emit('connection', {
           client: new WebSocketClassClient(mockWs, transport),
-          server: new WebSocketClassServer(mockWs, createConnection),
+          server: new WebSocketClassServer(mockWs, createConnection, transport),
         })
 
         return mockWs
@@ -68,7 +65,7 @@ const WEBSOCKET_CLOSE_CODE_RANGE_ERROR =
   'InvalidAccessError: close code out of user configurable range'
 
 export const kOnSend = Symbol('kOnSend')
-export const kOnReceive = Symbol('kOnReceive')
+// export const kOnReceive = Symbol('kOnReceive')
 
 type WebSocketEventListener = (this: WebSocket, event: Event) => void
 export type WebSocketMessageListener = (
@@ -100,7 +97,6 @@ export class WebSocketClassOverride extends EventTarget implements WebSocket {
   private _onclose: WebSocketCloseListener | null = null
 
   private [kOnSend]?: (data: WebSocketSendData) => void
-  private [kOnReceive]?: WebSocketTransportOnIncomingCallback
 
   constructor(url: string | URL, protocols?: string | Array<string>) {
     super()
@@ -198,29 +194,6 @@ export class WebSocketClassOverride extends EventTarget implements WebSocket {
     })
   }
 
-  public dispatchEvent(event: Event): boolean {
-    console.log('WebSocketClassOverride#dispatchEvent', event.type, event)
-
-    /**
-     * @note This override class never forwards the incoming
-     * events to the actual client instance. Instead, it
-     * forwards the incoming events to the connection
-     * and lets the "server" API handle the forwarding.
-     */
-    if (
-      event.type === 'message' &&
-      // Ignore mocked events sent from the connection.
-      // This condition is for the original server-sent events only.
-      !(event.target instanceof WebSocketClassOverride)
-    ) {
-      this[kOnReceive]?.(event as MessageEvent)
-      return true
-    }
-
-    // Dispatch the other events (open, close, etc).
-    return super.dispatchEvent(event)
-  }
-
   public close(code: number = 1000, reason?: string): void {
     invariant(code, WEBSOCKET_CLOSE_CODE_RANGE_ERROR)
     invariant(

src/interceptors/WebSocket/implementations/WebSocketClass/WebSocketClassServer.ts

@@ -6,70 +6,114 @@ import type {
   WebSocketMessageListener,
 } from './WebSocketClassInterceptor'
 import type { WebSocketSendData } from '../../WebSocketTransport'
+import type { WebSocketClassTransport } from './WebSocketClassTransport'
 import { bindEvent } from '../../utils/bindEvent'
 
+const kEmitter = Symbol('kEmitter')
+
 export class WebSocketClassServer extends WebSocketServer {
+  /**
+   * A WebSocket instance connected to the original server.
+   */
   private prodWs?: WebSocket
+  private [kEmitter]: EventTarget
 
   constructor(
     private readonly mockWs: WebSocketClassOverride,
-    private readonly createConnection: () => WebSocket
+    private readonly createConnection: () => WebSocket,
+    private readonly transport: WebSocketClassTransport
   ) {
     super()
+    this[kEmitter] = new EventTarget()
+
+    // Handle incoming events from the actual server.
+    // The (mock) WebSocket instance will call this
+    // whenever a "message" event from the actual server
+    // is dispatched on it (the dispatch will be skipped).
+    this.transport.onIncoming = (event) => {
+      /**
+       * @note Emit "message" event on the WebSocketClassServer
+       * instance to let the interceptor know about these
+       * incoming events from the original server. In that listener,
+       * the interceptor can modify or skip the event forwarding
+       * to the mock WebSocket instance.
+       */
+      this[kEmitter].dispatchEvent(event)
+    }
   }
 
-  public connect(): Promise<void> {
-    const connectionPromise = new DeferredPromise<void>()
+  public connect(): void {
+    invariant(
+      !this.prodWs,
+      'Failed to call "connect()" on the original WebSocket instance: the connection already open'
+    )
+
     const ws = this.createConnection()
 
-    ws.addEventListener('open', () => connectionPromise.resolve(), {
-      once: true,
-    })
-    ws.addEventListener('error', () => connectionPromise.reject(), {
-      once: true,
-    })
+    // Once the connection is open, forward any incoming
+    // events directly to the (override) WebSocket instance.
+    ws.addEventListener('message', (event) => {
+      // Clone the event to dispatch it on this class
+      // once again and prevent the "already being dispatched"
+      // exception. Clone it here so we can observe this event
+      // being prevented in the "server.on()" listeners.
+      const messageEvent = bindEvent(
+        this.prodWs!,
+        new MessageEvent('message', {
+          data: event.data,
+        })
+      )
+      this.transport.onIncoming(messageEvent)
 
-    return connectionPromise
-      .then(() => {
-        this.prodWs = ws
-      })
-      .catch((error) => {
-        console.error(
-          'Failed to connect to the original WebSocket server at "%s"',
-          this.mockWs.url
+      // Unless the default is prevented, forward the
+      // messages from the original server to the mock client.
+      // This is the only way the user can receive them.
+      if (!messageEvent.defaultPrevented) {
+        this.mockWs.dispatchEvent(
+          bindEvent(
+            /**
+             * @note Bind the forwarded original server events
+             * to the mock WebSocket instance so it would
+             * dispatch them straight away.
+             */
+            this.mockWs,
+            // Clone the message event again to prevent
+            // the "already being dispatched" exception.
+            new MessageEvent('message', { data: event.data })
+          )
         )
-        console.error(error)
-      })
+      }
+    })
+
+    this.prodWs = ws
   }
 
   public send(data: WebSocketSendData): void {
+    const { prodWs } = this
     invariant(
-      this.prodWs,
+      prodWs,
       'Failed to call "server.send()" for "%s": the connection is not open. Did you forget to call "await server.connect()"?',
       this.mockWs.url
     )
 
-    // Send the data using the original WebSocket connection.
-    this.prodWs.send(data)
+    // Delegate the send to when the original connection is open.
+    // Unlike the mock, connecting to the original server may take time
+    // so we cannot call this on the next tick.
+    if (prodWs.readyState === prodWs.CONNECTING) {
+      prodWs.addEventListener('open', () => prodWs.send(data), { once: true })
+      return
+    }
+
+    // Send the data to the original WebSocket server.
+    prodWs.send(data)
   }
 
   public on(event: 'message', callback: WebSocketMessageListener): void {
-    invariant(
-      this.prodWs,
-      'Failed to call "server.on(%s)" for "%s": the connection is not open. Did you forget to call "await server.connect()"?',
-      this.mockWs.url
-    )
+    console.log('WebSocketClassServer#on', event)
 
-    const { prodWs } = this
-
-    prodWs.addEventListener(event, (messageEvent) => {
-      callback.call(prodWs, messageEvent)
-
-      // Unless the default is prevented, forward the
-      // messages from the original server to the mock client.
-      // This is the only way the user can receive them.
-      if (!messageEvent.defaultPrevented) {
-        this.mockWs.dispatchEvent(bindEvent(this.mockWs, messageEvent))
+    this[kEmitter].addEventListener('message', (event) => {
+      if (event instanceof MessageEvent) {
+        callback.call(this.prodWs!, event)
       }
     })
   }

src/interceptors/WebSocket/implementations/WebSocketClass/WebSocketClassTransport.ts

@@ -5,21 +5,15 @@ import {
   WebSocketTransportOnIncomingCallback,
   WebSocketTransportOnOutgoingCallback,
 } from '../../WebSocketTransport'
-import {
-  kOnReceive,
-  kOnSend,
-  WebSocketClassOverride,
-} from './WebSocketClassInterceptor'
+import { kOnSend, WebSocketClassOverride } from './WebSocketClassInterceptor'
 
 export class WebSocketClassTransport extends WebSocketTransport {
   public onOutgoing: WebSocketTransportOnOutgoingCallback = () => {}
   public onIncoming: WebSocketTransportOnIncomingCallback = () => {}
 
   constructor(protected readonly ws: WebSocketClassOverride) {
     super()
-
     this.ws[kOnSend] = (...args) => this.onOutgoing(...args)
-    this.ws[kOnReceive] = (...args) => this.onIncoming(...args)
   }
 
   public send(data: WebSocketSendData): void {
@@ -46,8 +40,6 @@ export class WebSocketClassTransport extends WebSocketTransport {
         })
       )
 
-      console.log('message', message, message.target)
-
       this.ws.dispatchEvent(message)
     })
   }