feat(WebSocketHandler): add run method (#2527)

kettanaito · web-flow · commit 94fc78ea50bd · 2025-06-06T18:08:08.000+02:00
diff --git a/src/core/handlers/WebSocketHandler.test.ts b/src/core/handlers/WebSocketHandler.test.ts
@@ -1,17 +1,10 @@
-import type { WebSocketConnectionData } from '@mswjs/interceptors/WebSocket'
 import { WebSocketHandler } from './WebSocketHandler'
 
 describe('parse', () => {
   it('matches an exact url', () => {
     expect(
       new WebSocketHandler('ws://localhost:3000').parse({
-        event: new MessageEvent('connection', {
-          data: {
-            client: {
-              url: new URL('ws://localhost:3000'),
-            },
-          } as WebSocketConnectionData,
-        }),
+        url: new URL('ws://localhost:3000'),
       }),
     ).toEqual({
       match: {
@@ -24,13 +17,7 @@ describe('parse', () => {
   it('ignores trailing slash', () => {
     expect(
       new WebSocketHandler('ws://localhost:3000').parse({
-        event: new MessageEvent('connection', {
-          data: {
-            client: {
-              url: new URL('ws://localhost:3000/'),
-            },
-          } as WebSocketConnectionData,
-        }),
+        url: new URL('ws://localhost:3000/'),
       }),
     ).toEqual({
       match: {
@@ -41,13 +28,7 @@ describe('parse', () => {
 
     expect(
       new WebSocketHandler('ws://localhost:3000/').parse({
-        event: new MessageEvent('connection', {
-          data: {
-            client: {
-              url: new URL('ws://localhost:3000'),
-            },
-          } as WebSocketConnectionData,
-        }),
+        url: new URL('ws://localhost:3000/'),
       }),
     ).toEqual({
       match: {
@@ -60,13 +41,7 @@ describe('parse', () => {
   it('supports path parameters', () => {
     expect(
       new WebSocketHandler('ws://localhost:3000/:serviceName').parse({
-        event: new MessageEvent('connection', {
-          data: {
-            client: {
-              url: new URL('ws://localhost:3000/auth'),
-            },
-          } as WebSocketConnectionData,
-        }),
+        url: new URL('ws://localhost:3000/auth'),
       }),
     ).toEqual({
       match: {
@@ -81,15 +56,9 @@ describe('parse', () => {
   it('ignores "/socket.io/" prefix in the client url', () => {
     expect(
       new WebSocketHandler('ws://localhost:3000').parse({
-        event: new MessageEvent('connection', {
-          data: {
-            client: {
-              url: new URL(
-                'ws://localhost:3000/socket.io/?EIO=4&transport=websocket',
-              ),
-            },
-          } as WebSocketConnectionData,
-        }),
+        url: new URL(
+          'ws://localhost:3000/socket.io/?EIO=4&transport=websocket',
+        ),
       }),
     ).toEqual({
       match: {
@@ -100,15 +69,9 @@ describe('parse', () => {
 
     expect(
       new WebSocketHandler('ws://localhost:3000/non-matching').parse({
-        event: new MessageEvent('connection', {
-          data: {
-            client: {
-              url: new URL(
-                'ws://localhost:3000/socket.io/?EIO=4&transport=websocket',
-              ),
-            },
-          } as WebSocketConnectionData,
-        }),
+        url: new URL(
+          'ws://localhost:3000/socket.io/?EIO=4&transport=websocket',
+        ),
       }),
     ).toEqual({
       match: {
@@ -125,13 +88,7 @@ describe('parse', () => {
      */
     expect(
       new WebSocketHandler('ws://localhost:3000/clients/socket.io/123').parse({
-        event: new MessageEvent('connection', {
-          data: {
-            client: {
-              url: new URL('ws://localhost:3000/clients/socket.io/123'),
-            },
-          } as WebSocketConnectionData,
-        }),
+        url: new URL('ws://localhost:3000/clients/socket.io/123'),
       }),
     ).toEqual({
       match: {
@@ -142,13 +99,7 @@ describe('parse', () => {
 
     expect(
       new WebSocketHandler('ws://localhost:3000').parse({
-        event: new MessageEvent('connection', {
-          data: {
-            client: {
-              url: new URL('ws://localhost:3000/clients/socket.io/123'),
-            },
-          } as WebSocketConnectionData,
-        }),
+        url: new URL('ws://localhost:3000/clients/socket.io/123'),
       }),
     ).toEqual({
       match: {
diff --git a/src/core/handlers/WebSocketHandler.ts b/src/core/handlers/WebSocketHandler.ts
@@ -23,7 +23,6 @@ export interface WebSocketHandlerConnection extends WebSocketConnectionData {
 }
 
 export const kEmitter = Symbol('kEmitter')
-export const kDispatchEvent = Symbol('kDispatchEvent')
 export const kSender = Symbol('kSender')
 const kStopPropagationPatched = Symbol('kStopPropagationPatched')
 const KOnStopPropagation = Symbol('KOnStopPropagation')
@@ -44,11 +43,8 @@ export class WebSocketHandler {
     this.__kind = 'EventHandler'
   }
 
-  public parse(args: {
-    event: MessageEvent<WebSocketConnectionData>
-  }): WebSocketHandlerParsedResult {
-    const { data: connection } = args.event
-    const { url: clientUrl } = connection.client
+  public parse(args: { url: URL }): WebSocketHandlerParsedResult {
+    const clientUrl = new URL(args.url)
 
     /**
      * @note Remove the Socket.IO path prefix from the WebSocket
@@ -65,23 +61,30 @@ export class WebSocketHandler {
   }
 
   public predicate(args: {
-    event: MessageEvent<WebSocketConnectionData>
+    url: URL
     parsedResult: WebSocketHandlerParsedResult
   }): boolean {
     return args.parsedResult.match.matches
   }
 
-  async [kDispatchEvent](
-    event: MessageEvent<WebSocketConnectionData>,
-  ): Promise<void> {
-    const parsedResult = this.parse({ event })
-    const connection = event.data
+  public async run(connection: WebSocketConnectionData): Promise<boolean> {
+    const parsedResult = this.parse({
+      url: connection.client.url,
+    })
+
+    if (!this.predicate({ url: connection.client.url, parsedResult })) {
+      return false
+    }
 
     const resolvedConnection: WebSocketHandlerConnection = {
       ...connection,
       params: parsedResult.match.params || {},
     }
 
+    return this.connect(resolvedConnection)
+  }
+
+  protected connect(connection: WebSocketHandlerConnection): boolean {
     // Support `event.stopPropagation()` for various client/server events.
     connection.client.addEventListener(
       'message',
@@ -111,7 +114,7 @@ export class WebSocketHandler {
 
     // Emit the connection event on the handler.
     // This is what the developer adds listeners for.
-    this[kEmitter].emit('connection', resolvedConnection)
+    return this[kEmitter].emit('connection', connection)
   }
 }
 
diff --git a/src/core/ws/handleWebSocketEvent.ts b/src/core/ws/handleWebSocketEvent.ts
@@ -1,6 +1,6 @@
 import type { WebSocketConnectionData } from '@mswjs/interceptors/lib/browser/interceptors/WebSocket'
 import { RequestHandler } from '../handlers/RequestHandler'
-import { WebSocketHandler, kDispatchEvent } from '../handlers/WebSocketHandler'
+import { WebSocketHandler } from '../handlers/WebSocketHandler'
 import { webSocketInterceptor } from './webSocketInterceptor'
 import {
   onUnhandledRequest,
@@ -17,67 +17,48 @@ interface HandleWebSocketEventOptions {
 
 export function handleWebSocketEvent(options: HandleWebSocketEventOptions) {
   webSocketInterceptor.on('connection', async (connection) => {
-    const handlers = options.getHandlers()
+    const handlers = options.getHandlers().filter(isHandlerKind('EventHandler'))
 
-    const connectionEvent = new MessageEvent('connection', {
-      data: connection,
-    })
+    // Ignore this connection if the user hasn't defined any handlers.
+    if (handlers.length > 0) {
+      options?.onMockedConnection(connection)
 
-    // First, filter only those WebSocket handlers that
-    // match the "ws.link()" endpoint predicate. Don't dispatch
-    // anything yet so the logger can be attached to the connection
-    // before it potentially sends events.
-    const matchingHandlers: Array<WebSocketHandler> = []
+      await Promise.all(
+        handlers.map((handler) => {
+          // Iterate over the handlers and forward the connection
+          // event to WebSocket event handlers. This is equivalent
+          // to dispatching that event onto multiple listeners.
+          return handler.run(connection)
+        }),
+      )
 
-    for (const handler of handlers) {
-      if (
-        isHandlerKind('EventHandler')(handler) &&
-        handler.predicate({
-          event: connectionEvent,
-          parsedResult: handler.parse({
-            event: connectionEvent,
-          }),
-        })
-      ) {
-        matchingHandlers.push(handler)
-      }
+      return
     }
 
-    if (matchingHandlers.length > 0) {
-      options?.onMockedConnection(connection)
-
-      // Iterate over the handlers and forward the connection
-      // event to WebSocket event handlers. This is equivalent
-      // to dispatching that event onto multiple listeners.
-      for (const handler of matchingHandlers) {
-        handler[kDispatchEvent](connectionEvent)
-      }
-    } else {
-      // Construct a request representing this WebSocket connection.
-      const request = new Request(connection.client.url, {
-        headers: {
-          upgrade: 'websocket',
-          connection: 'upgrade',
-        },
-      })
-      await onUnhandledRequest(
-        request,
-        options.getUnhandledRequestStrategy(),
-      ).catch((error) => {
-        const errorEvent = new Event('error')
-        Object.defineProperty(errorEvent, 'cause', {
-          enumerable: true,
-          configurable: false,
-          value: error,
-        })
-        connection.client.socket.dispatchEvent(errorEvent)
+    // Construct a request representing this WebSocket connection.
+    const request = new Request(connection.client.url, {
+      headers: {
+        upgrade: 'websocket',
+        connection: 'upgrade',
+      },
+    })
+    await onUnhandledRequest(
+      request,
+      options.getUnhandledRequestStrategy(),
+    ).catch((error) => {
+      const errorEvent = new Event('error')
+      Object.defineProperty(errorEvent, 'cause', {
+        enumerable: true,
+        configurable: false,
+        value: error,
       })
+      connection.client.socket.dispatchEvent(errorEvent)
+    })
 
-      options?.onPassthroughConnection(connection)
+    options?.onPassthroughConnection(connection)
 
-      // If none of the "ws" handlers matched,
-      // establish the WebSocket connection as-is.
-      connection.server.connect()
-    }
+    // If none of the "ws" handlers matched,
+    // establish the WebSocket connection as-is.
+    connection.server.connect()
   })
 }
diff --git a/src/core/ws/utils/getMessageLength.ts b/src/core/ws/utils/getMessageLength.ts
@@ -15,5 +15,5 @@ export function getMessageLength(data: WebSocketData): number {
     return data.byteLength
   }
 
-  return new Blob([data]).size
+  return new Blob([data as any]).size
 }
diff --git a/src/core/ws/utils/getPublicData.ts b/src/core/ws/utils/getPublicData.ts
@@ -9,7 +9,7 @@ export async function getPublicData(data: WebSocketData): Promise<string> {
 
   // Handle all ArrayBuffer-like objects.
   if (typeof data === 'object' && 'byteLength' in data) {
-    const text = new TextDecoder().decode(data)
+    const text = new TextDecoder().decode(data as ArrayBuffer)
     return `ArrayBuffer(${truncateMessage(text)})`
   }
 
diff --git a/test/node/ws-api/ws.server.connect.test.ts b/test/node/ws-api/ws.server.connect.test.ts
@@ -86,7 +86,7 @@ it('throws an error when connecting to a non-existing server', async () => {
   )
 
   const errorListener = vi.fn()
-  const ws = new WebSocket('wss://localhost:9876')
+  const ws = new WebSocket('ws://localhost:9876')
   ws.onerror = errorListener
 
   await vi.waitFor(() => {

Original file line number	Diff line number	Diff line change
`@@ -15,5 +15,5 @@ export function getMessageLength(data: WebSocketData): number {`
`15`	`15`	`return data.byteLength`
`16`	`16`	`}`
`17`	`17`
`18`		`- return new Blob([data]).size`
	`18`	`+ return new Blob([data as any]).size`
`19`	`19`	`}`
Original file line number	Diff line number	Diff line change
`@@ -9,7 +9,7 @@ export async function getPublicData(data: WebSocketData): Promise<string> {`
`9`	`9`
`10`	`10`	`// Handle all ArrayBuffer-like objects.`
`11`	`11`	`if (typeof data === 'object' && 'byteLength' in data) {`
`12`		`- const text = new TextDecoder().decode(data)`
	`12`	`+ const text = new TextDecoder().decode(data as ArrayBuffer)`
`13`	`13`	return `ArrayBuffer(${truncateMessage(text)})`
`14`	`14`	`}`
`15`	`15`
Original file line number	Diff line number	Diff line change
`@@ -86,7 +86,7 @@ it('throws an error when connecting to a non-existing server', async () => {`
`86`	`86`	`)`
`87`	`87`
`88`	`88`	`const errorListener = vi.fn()`
`89`		`- const ws = new WebSocket('wss://localhost:9876')`
	`89`	`+ const ws = new WebSocket('ws://localhost:9876')`
`90`	`90`	`ws.onerror = errorListener`
`91`	`91`
`92`	`92`	`await vi.waitFor(() => {`