chore: add agnostic "Handler" class

src/core/handlers/WebSocketHandler.ts

@@ -0,0 +1,138 @@
+import { Emitter } from 'strict-event-emitter'
+import type {
+  WebSocketClientConnection,
+  WebSocketServerConnection,
+} from '@mswjs/interceptors/WebSocket'
+import {
+  type Match,
+  type Path,
+  type PathParams,
+  matchRequestUrl,
+} from '../utils/matching/matchRequestUrl'
+
+export type HandlerOptions = {
+  once?: boolean
+}
+
+export abstract class Handler<Input = unknown> {
+  public isUsed: boolean
+
+  constructor(protected readonly options: HandlerOptions = {}) {
+    this.isUsed = false
+  }
+
+  abstract parse(args: { input: Input }): unknown
+  abstract predicate(args: { input: Input; parsedResult: unknown }): boolean
+  protected abstract handle(args: {
+    input: Input
+    parsedResult: unknown
+  }): Promise<unknown | null>
+
+  public async run(input: Input): Promise<unknown | null> {
+    if (this.options?.once && this.isUsed) {
+      return null
+    }
+
+    const parsedResult = this.parse({ input })
+    const shouldHandle = this.predicate({
+      input,
+      parsedResult,
+    })
+
+    if (!shouldHandle) {
+      return null
+    }
+
+    const result = await this.handle({
+      input,
+      parsedResult,
+    })
+
+    this.isUsed = true
+
+    return result
+  }
+}
+
+type WebSocketHandlerParsedResult = {
+  match: Match
+}
+
+type WebSocketHandlerEventMap = {
+  connection: [
+    args: {
+      client: WebSocketClientConnection
+      server: WebSocketServerConnection
+      params: PathParams
+    },
+  ]
+}
+
+export class WebSocketHandler extends Handler<MessageEvent<any>> {
+  public on: <K extends keyof WebSocketHandlerEventMap>(
+    event: K,
+    listener: (...args: WebSocketHandlerEventMap[K]) => void,
+  ) => void
+
+  public off: <K extends keyof WebSocketHandlerEventMap>(
+    event: K,
+    listener: (...args: WebSocketHandlerEventMap[K]) => void,
+  ) => void
+
+  public removeAllListeners: <K extends keyof WebSocketHandlerEventMap>(
+    event?: K,
+  ) => void
+
+  protected emitter: Emitter<WebSocketHandlerEventMap>
+
+  constructor(private readonly url: Path) {
+    super()
+    this.emitter = new Emitter()
+
+    // Forward some of the emitter API to the public API
+    // of the event handler.
+    this.on = this.emitter.on.bind(this.emitter)
+    this.off = this.emitter.off.bind(this.emitter)
+    this.removeAllListeners = this.emitter.removeAllListeners.bind(this.emitter)
+  }
+
+  public parse(args: {
+    input: MessageEvent<any>
+  }): WebSocketHandlerParsedResult {
+    const connection = args.input.data
+    const match = matchRequestUrl(connection.client.url, this.url)
+
+    return {
+      match,
+    }
+  }
+
+  public predicate(args: {
+    input: MessageEvent<any>
+    parsedResult: WebSocketHandlerParsedResult
+  }): boolean {
+    const { match } = args.parsedResult
+    return match.matches
+  }
+
+  protected async handle(args: {
+    input: MessageEvent<any>
+    parsedResult: WebSocketHandlerParsedResult
+  }): Promise<void> {
+    const connectionEvent = args.input
+
+    // At this point, the WebSocket connection URL has matched the handler.
+    // Prevent the default behavior of establishing the connection as-is.
+    connectionEvent.preventDefault()
+
+    const connection = connectionEvent.data
+
+    // Emit the connection event on the handler.
+    // This is what the developer adds listeners for.
+    this.emitter.emit('connection', {
+      client: connection.client,
+      server: connection.server,
+      params: args.parsedResult.match.params || {},
+    })
+  }
+}

src/core/utils/handleWebSocketEvent.ts

@@ -0,0 +1,31 @@
+import { type Handler, WebSocketHandler } from '../handlers/WebSocketHandler'
+import { webSocketInterceptor } from '../ws/webSocketInterceptor'
+
+export function handleWebSocketEvent(handlers: Array<Handler>) {
+  webSocketInterceptor.on('connection', (connection) => {
+    const connectionEvent = new MessageEvent('connection', {
+      data: connection,
+      cancelable: true,
+    })
+
+    // 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 handlers) {
+      if (handler instanceof WebSocketHandler) {
+        // Never await the run function because event handlers
+        // are side-effectful and don't block the event loop.
+        handler.run(connectionEvent)
+      }
+    }
+
+    // If none of the "ws" handlers matched,
+    // establish the WebSocket connection as-is.
+    if (!connectionEvent.defaultPrevented) {
+      connection.server.connect()
+      connection.client.on('message', (event) => {
+        connection.server.send(event.data)
+      })
+    }
+  })
+}

src/core/ws/webSocketInterceptor.ts

@@ -0,0 +1,3 @@
+import { WebSocketInterceptor } from '@mswjs/interceptors/WebSocket'
+
+export const webSocketInterceptor = new WebSocketInterceptor()

src/core/ws/ws.ts

@@ -0,0 +1,19 @@
+import { WebSocketHandler } from '../handlers/WebSocketHandler'
+import type { Path } from '../utils/matching/matchRequestUrl'
+import { webSocketInterceptor } from './webSocketInterceptor'
+
+/**
+ * Intercepts outgoing WebSocket connections to the given URL.
+ *
+ * @example
+ * const chat = ws.link('wss://chat.example.com')
+ * chat.on('connection', (connection) => {})
+ */
+function createWebSocketLinkHandler(url: Path) {
+  webSocketInterceptor.apply()
+  return new WebSocketHandler(url)
+}
+
+export const ws = {
+  link: createWebSocketLinkHandler,
+}