msw 2.2.2 → 2.3.0-ws.rc-1

package/src/browser/setupWorker/glossary.ts

@@ -5,13 +5,11 @@ import {
   SharedOptions,
 } from '~/core/sharedOptions'
 import { ServiceWorkerMessage } from './start/utils/createMessageChannel'
-import {
-  RequestHandler,
-  RequestHandlerDefaultInfo,
-} from '~/core/handlers/RequestHandler'
+import { RequestHandler } from '~/core/handlers/RequestHandler'
 import type { HttpRequestEventMap, Interceptor } from '@mswjs/interceptors'
-import { Path } from '~/core/utils/matching/matchRequestUrl'
-import { RequiredDeep } from '~/core/typeUtils'
+import type { Path } from '~/core/utils/matching/matchRequestUrl'
+import type { RequiredDeep } from '~/core/typeUtils'
+import type { WebSocketHandler } from '~/core/handlers/WebSocketHandler'
 
 export type ResolvedPath = Path | URL
 
@@ -84,7 +82,7 @@ export interface SetupWorkerInternalContext {
   startOptions: RequiredDeep<StartOptions>
   worker: ServiceWorker | null
   registration: ServiceWorkerRegistration | null
-  getRequestHandlers(): Array<RequestHandler>
+  getRequestHandlers(): Array<RequestHandler | WebSocketHandler>
   requests: Map<string, Request>
   emitter: Emitter<LifeCycleEventsMap>
   keepAliveInterval?: number
@@ -208,7 +206,7 @@ export interface SetupWorker {
    *
    * @see {@link https://mswjs.io/docs/api/setup-worker/use `worker.use()` API reference}
    */
-  use: (...handlers: RequestHandler[]) => void
+  use: (...handlers: Array<RequestHandler | WebSocketHandler>) => void
 
   /**
    * Marks all request handlers that respond using `res.once()` as unused.
@@ -223,14 +221,16 @@ export interface SetupWorker {
    *
    * @see {@link https://mswjs.io/docs/api/setup-worker/reset-handlers `worker.resetHandlers()` API reference}
    */
-  resetHandlers: (...nextHandlers: RequestHandler[]) => void
+  resetHandlers: (
+    ...nextHandlers: Array<RequestHandler | WebSocketHandler>
+  ) => void
 
   /**
    * Returns a readonly list of currently active request handlers.
    *
    * @see {@link https://mswjs.io/docs/api/setup-worker/list-handlers `worker.listHandlers()` API reference}
    */
-  listHandlers(): ReadonlyArray<RequestHandler<RequestHandlerDefaultInfo, any>>
+  listHandlers(): ReadonlyArray<RequestHandler | WebSocketHandler>
 
   /**
    * Life-cycle events.

package/src/browser/setupWorker/setupWorker.ts

@@ -18,9 +18,12 @@ import { createFallbackStop } from './stop/createFallbackStop'
 import { devUtils } from '~/core/utils/internal/devUtils'
 import { SetupApi } from '~/core/SetupApi'
 import { mergeRight } from '~/core/utils/internal/mergeRight'
-import { LifeCycleEventsMap } from '~/core/sharedOptions'
+import type { LifeCycleEventsMap } from '~/core/sharedOptions'
+import type { WebSocketHandler } from '~/core/handlers/WebSocketHandler'
 import { SetupWorker } from './glossary'
 import { supportsReadableStreamTransfer } from '../utils/supportsReadableStreamTransfer'
+import { webSocketInterceptor } from '~/core/ws/webSocketInterceptor'
+import { handleWebSocketEvent } from '~/core/utils/handleWebSocketEvent'
 
 interface Listener {
   target: EventTarget
@@ -37,7 +40,7 @@ export class SetupWorkerApi
   private stopHandler: StopHandler = null as any
   private listeners: Array<Listener>
 
-  constructor(...handlers: Array<RequestHandler>) {
+  constructor(...handlers: Array<RequestHandler | WebSocketHandler>) {
     super(...handlers)
 
     invariant(
@@ -176,6 +179,15 @@ export class SetupWorkerApi
       options,
     ) as SetupWorkerInternalContext['startOptions']
 
+    handleWebSocketEvent(() => {
+      return this.handlersController.currentHandlers()
+    })
+    webSocketInterceptor.apply()
+
+    this.subscriptions.push(() => {
+      webSocketInterceptor.dispose()
+    })
+
     return await this.startHandler(this.context.startOptions, options)
   }
 
@@ -193,6 +205,8 @@ export class SetupWorkerApi
  *
  * @see {@link https://mswjs.io/docs/api/setup-worker `setupWorker()` API reference}
  */
-export function setupWorker(...handlers: Array<RequestHandler>): SetupWorker {
+export function setupWorker(
+  ...handlers: Array<RequestHandler | WebSocketHandler>
+): SetupWorker {
   return new SetupWorkerApi(...handlers)
 }

package/src/core/SetupApi.ts

@@ -1,38 +1,42 @@
 import { invariant } from 'outvariant'
 import { EventMap, Emitter } from 'strict-event-emitter'
-import {
-  RequestHandler,
-  RequestHandlerDefaultInfo,
-} from './handlers/RequestHandler'
+import { RequestHandler } from './handlers/RequestHandler'
 import { LifeCycleEventEmitter } from './sharedOptions'
 import { devUtils } from './utils/internal/devUtils'
 import { pipeEvents } from './utils/internal/pipeEvents'
 import { toReadonlyArray } from './utils/internal/toReadonlyArray'
 import { Disposable } from './utils/internal/Disposable'
+import type { WebSocketHandler } from './handlers/WebSocketHandler'
 
 export abstract class HandlersController {
-  abstract prepend(runtimeHandlers: Array<RequestHandler>): void
-  abstract reset(nextHandles: Array<RequestHandler>): void
-  abstract currentHandlers(): Array<RequestHandler>
+  abstract prepend(
+    runtimeHandlers: Array<RequestHandler | WebSocketHandler>,
+  ): void
+  abstract reset(nextHandles: Array<RequestHandler | WebSocketHandler>): void
+  abstract currentHandlers(): Array<RequestHandler | WebSocketHandler>
 }
 
 export class InMemoryHandlersController implements HandlersController {
-  private handlers: Array<RequestHandler>
+  private handlers: Array<RequestHandler | WebSocketHandler>
 
-  constructor(private initialHandlers: Array<RequestHandler>) {
+  constructor(
+    private initialHandlers: Array<RequestHandler | WebSocketHandler>,
+  ) {
     this.handlers = [...initialHandlers]
   }
 
-  public prepend(runtimeHandles: Array<RequestHandler>): void {
+  public prepend(
+    runtimeHandles: Array<RequestHandler | WebSocketHandler>,
+  ): void {
     this.handlers.unshift(...runtimeHandles)
   }
 
-  public reset(nextHandlers: Array<RequestHandler>): void {
+  public reset(nextHandlers: Array<RequestHandler | WebSocketHandler>): void {
     this.handlers =
       nextHandlers.length > 0 ? [...nextHandlers] : [...this.initialHandlers]
   }
 
-  public currentHandlers(): Array<RequestHandler> {
+  public currentHandlers(): Array<RequestHandler | WebSocketHandler> {
     return this.handlers
   }
 }
@@ -47,7 +51,7 @@ export abstract class SetupApi<EventsMap extends EventMap> extends Disposable {
 
   public readonly events: LifeCycleEventEmitter<EventsMap>
 
-  constructor(...initialHandlers: Array<RequestHandler>) {
+  constructor(...initialHandlers: Array<RequestHandler | WebSocketHandler>) {
     super()
 
     invariant(
@@ -71,12 +75,14 @@ export abstract class SetupApi<EventsMap extends EventMap> extends Disposable {
     })
   }
 
-  private validateHandlers(handlers: ReadonlyArray<RequestHandler>): boolean {
+  private validateHandlers(handlers: ReadonlyArray<unknown>): boolean {
     // Guard against incorrect call signature of the setup API.
     return handlers.every((handler) => !Array.isArray(handler))
   }
 
-  public use(...runtimeHandlers: Array<RequestHandler>): void {
+  public use(
+    ...runtimeHandlers: Array<RequestHandler | WebSocketHandler>
+  ): void {
     invariant(
       this.validateHandlers(runtimeHandlers),
       devUtils.formatMessage(
@@ -89,17 +95,19 @@ export abstract class SetupApi<EventsMap extends EventMap> extends Disposable {
 
   public restoreHandlers(): void {
     this.handlersController.currentHandlers().forEach((handler) => {
-      handler.isUsed = false
+      if ('isUsed' in handler) {
+        handler.isUsed = false
+      }
     })
   }
 
-  public resetHandlers(...nextHandlers: Array<RequestHandler>): void {
+  public resetHandlers(
+    ...nextHandlers: Array<RequestHandler | WebSocketHandler>
+  ): void {
     this.handlersController.reset(nextHandlers)
   }
 
-  public listHandlers(): ReadonlyArray<
-    RequestHandler<RequestHandlerDefaultInfo, any, any>
-  > {
+  public listHandlers(): ReadonlyArray<RequestHandler | WebSocketHandler> {
     return toReadonlyArray(this.handlersController.currentHandlers())
   }
 

package/src/core/handlers/WebSocketHandler.ts

@@ -0,0 +1,89 @@
+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'
+
+type WebSocketHandlerParsedResult = {
+  match: Match
+}
+
+export type WebSocketHandlerEventMap = {
+  connection: [
+    args: {
+      client: WebSocketClientConnection
+      server: WebSocketServerConnection
+      params: PathParams
+    },
+  ]
+}
+
+type WebSocketHandlerIncomingEvent = MessageEvent<{
+  client: WebSocketClientConnection
+  server: WebSocketServerConnection
+}>
+
+export const kEmitter = Symbol('kEmitter')
+export const kDispatchEvent = Symbol('kDispatchEvent')
+export const kDefaultPrevented = Symbol('kDefaultPrevented')
+
+export class WebSocketHandler {
+  protected [kEmitter]: Emitter<WebSocketHandlerEventMap>
+
+  constructor(private readonly url: Path) {
+    this[kEmitter] = new Emitter()
+  }
+
+  public parse(args: {
+    event: WebSocketHandlerIncomingEvent
+  }): WebSocketHandlerParsedResult {
+    const connection = args.event.data
+    const match = matchRequestUrl(connection.client.url, this.url)
+
+    return {
+      match,
+    }
+  }
+
+  public predicate(args: {
+    event: WebSocketHandlerIncomingEvent
+    parsedResult: WebSocketHandlerParsedResult
+  }): boolean {
+    return args.parsedResult.match.matches
+  }
+
+  async [kDispatchEvent](event: MessageEvent<any>): Promise<void> {
+    const parsedResult = this.parse({ event })
+    const shouldIntercept = this.predicate({ event, parsedResult })
+
+    if (!shouldIntercept) {
+      return
+    }
+
+    // Account for other matching event handlers that've already prevented this event.
+    if (!Reflect.get(event, kDefaultPrevented)) {
+      // At this point, the WebSocket connection URL has matched the handler.
+      // Prevent the default behavior of establishing the connection as-is.
+      // Use internal symbol because we aren't actually dispatching this
+      // event. Events can only marked as cancelable and can be prevented
+      // when dispatched on an EventTarget.
+      Reflect.set(event, kDefaultPrevented, true)
+    }
+
+    const connection = event.data
+
+    // Emit the connection event on the handler.
+    // This is what the developer adds listeners for.
+    this[kEmitter].emit('connection', {
+      client: connection.client,
+      server: connection.server,
+      params: parsedResult.match.params || {},
+    })
+  }
+}

package/src/core/index.ts

@@ -9,6 +9,9 @@ export { HttpHandler, HttpMethods } from './handlers/HttpHandler'
 export { graphql } from './graphql'
 export { GraphQLHandler } from './handlers/GraphQLHandler'
 
+/* WebSocket */
+export { ws } from './ws/ws'
+
 /* Utils */
 export { matchRequestUrl } from './utils/matching/matchRequestUrl'
 export * from './utils/handleRequest'

package/src/core/utils/executeHandlers.ts

@@ -1,6 +1,6 @@
 import {
   RequestHandler,
-  RequestHandlerExecutionResult,
+  type RequestHandlerExecutionResult,
 } from '../handlers/RequestHandler'
 
 export interface HandlersExecutionResult {
@@ -18,7 +18,7 @@ export interface ResponseResolutionContext {
  * Returns the execution result object containing any matching request
  * handler and any mocked response it returned.
  */
-export const executeHandlers = async <Handlers extends Array<RequestHandler>>({
+export const executeHandlers = async <Handlers extends Array<unknown>>({
   request,
   requestId,
   handlers,
@@ -33,6 +33,10 @@ export const executeHandlers = async <Handlers extends Array<RequestHandler>>({
   let result: RequestHandlerExecutionResult<any> | null = null
 
   for (const handler of handlers) {
+    if (!(handler instanceof RequestHandler)) {
+      continue
+    }
+
     result = await handler.run({ request, requestId, resolutionContext })
 
     // If the handler produces some result for this request,

package/src/core/utils/handleRequest.ts

@@ -1,6 +1,5 @@
 import { until } from '@open-draft/until'
 import { Emitter } from 'strict-event-emitter'
-import { RequestHandler } from '../handlers/RequestHandler'
 import { LifeCycleEventsMap, SharedOptions } from '../sharedOptions'
 import { RequiredDeep } from '../typeUtils'
 import { HandlersExecutionResult, executeHandlers } from './executeHandlers'
@@ -45,7 +44,7 @@ export interface HandleRequestOptions {
 export async function handleRequest(
   request: Request,
   requestId: string,
-  handlers: Array<RequestHandler>,
+  handlers: Array<unknown>,
   options: RequiredDeep<SharedOptions>,
   emitter: Emitter<LifeCycleEventsMap>,
   handleRequestOptions?: HandleRequestOptions,

package/src/core/utils/handleWebSocketEvent.ts

@@ -0,0 +1,49 @@
+import { RequestHandler } from '../handlers/RequestHandler'
+import {
+  WebSocketHandler,
+  kDefaultPrevented,
+  kDispatchEvent,
+} from '../handlers/WebSocketHandler'
+import { webSocketInterceptor } from '../ws/webSocketInterceptor'
+
+export function handleWebSocketEvent(
+  getCurrentHandlers: () => Array<RequestHandler | WebSocketHandler>,
+) {
+  webSocketInterceptor.on('connection', (connection) => {
+    const handlers = getCurrentHandlers()
+
+    const connectionEvent = new MessageEvent('connection', {
+      data: connection,
+      /**
+       * @note This message event should be marked as "cancelable"
+       * to have its default prevented using "event.preventDefault()".
+       * There's a bug in Node.js that breaks the "cancelable" flag.
+       * @see https://github.com/nodejs/node/issues/51767
+       */
+    })
+
+    Object.defineProperty(connectionEvent, kDefaultPrevented, {
+      enumerable: false,
+      writable: true,
+      value: false,
+    })
+
+    // 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) {
+        handler[kDispatchEvent](connectionEvent)
+      }
+    }
+
+    // If none of the "ws" handlers matched,
+    // establish the WebSocket connection as-is.
+    if (!Reflect.get(connectionEvent, kDefaultPrevented)) {
+      connection.server.connect()
+      connection.client.addEventListener('message', (event) => {
+        connection.server.send(event.data)
+      })
+    }
+  })
+}

package/src/core/utils/matching/matchRequestUrl.test.ts

@@ -61,6 +61,50 @@ describe('matchRequestUrl', () => {
     expect(match).toHaveProperty('matches', false)
     expect(match).toHaveProperty('params', {})
   })
+
+  test('returns true for matching WebSocket URL', () => {
+    expect(
+      matchRequestUrl(new URL('ws://test.mswjs.io'), 'ws://test.mswjs.io'),
+    ).toEqual({
+      matches: true,
+      params: {},
+    })
+    expect(
+      matchRequestUrl(new URL('wss://test.mswjs.io'), 'wss://test.mswjs.io'),
+    ).toEqual({
+      matches: true,
+      params: {},
+    })
+  })
+
+  test('returns false for non-matching WebSocket URL', () => {
+    expect(
+      matchRequestUrl(new URL('ws://test.mswjs.io'), 'ws://foo.mswjs.io'),
+    ).toEqual({
+      matches: false,
+      params: {},
+    })
+    expect(
+      matchRequestUrl(new URL('wss://test.mswjs.io'), 'wss://completely.diff'),
+    ).toEqual({
+      matches: false,
+      params: {},
+    })
+  })
+
+  test('returns path parameters when matched a WebSocket URL', () => {
+    expect(
+      matchRequestUrl(
+        new URL('wss://test.mswjs.io'),
+        'wss://:service.mswjs.io',
+      ),
+    ).toEqual({
+      matches: true,
+      params: {
+        service: 'test',
+      },
+    })
+  })
 })
 
 describe('coercePath', () => {

package/src/core/utils/matching/matchRequestUrl.ts

@@ -71,3 +71,7 @@ export function matchRequestUrl(url: URL, path: Path, baseUrl?: string): Match {
     params,
   }
 }
+
+export function isPath(value: unknown): value is Path {
+  return typeof value === 'string' || value instanceof RegExp
+}

package/src/core/ws/WebSocketClientManager.test.ts

@@ -0,0 +1,159 @@
+/**
+ * @vitest-environment node-websocket
+ */
+import { randomUUID } from 'node:crypto'
+import {
+  WebSocketClientConnection,
+  WebSocketTransport,
+} from '@mswjs/interceptors/WebSocket'
+import {
+  WebSocketClientManager,
+  WebSocketBroadcastChannelMessage,
+  WebSocketRemoteClientConnection,
+} from './WebSocketClientManager'
+
+const channel = new BroadcastChannel('test:channel')
+vi.spyOn(channel, 'postMessage')
+
+const socket = new WebSocket('ws://localhost')
+const transport = {
+  onOutgoing: vi.fn(),
+  onIncoming: vi.fn(),
+  onClose: vi.fn(),
+  send: vi.fn(),
+  close: vi.fn(),
+} satisfies WebSocketTransport
+
+afterEach(() => {
+  vi.resetAllMocks()
+})
+
+it('adds a client from this runtime to the list of clients', () => {
+  const manager = new WebSocketClientManager(channel)
+  const connection = new WebSocketClientConnection(socket, transport)
+
+  manager.addConnection(connection)
+
+  // Must add the client to the list of clients.
+  expect(Array.from(manager.clients.values())).toEqual([connection])
+
+  // Must emit the connection open event to notify other runtimes.
+  expect(channel.postMessage).toHaveBeenCalledWith({
+    type: 'connection:open',
+    payload: {
+      clientId: connection.id,
+      url: socket.url,
+    },
+  } satisfies WebSocketBroadcastChannelMessage)
+})
+
+it('adds a client from another runtime to the list of clients', async () => {
+  const clientId = randomUUID()
+  const url = new URL('ws://localhost')
+  const manager = new WebSocketClientManager(channel)
+
+  channel.dispatchEvent(
+    new MessageEvent<WebSocketBroadcastChannelMessage>('message', {
+      data: {
+        type: 'connection:open',
+        payload: {
+          clientId,
+          url: url.href,
+        },
+      },
+    }),
+  )
+
+  await vi.waitFor(() => {
+    expect(Array.from(manager.clients.values())).toEqual([
+      new WebSocketRemoteClientConnection(clientId, url, channel),
+    ])
+  })
+})
+
+it('replays a "send" event coming from another runtime', async () => {
+  const manager = new WebSocketClientManager(channel)
+  const connection = new WebSocketClientConnection(socket, transport)
+  manager.addConnection(connection)
+  vi.spyOn(connection, 'send')
+
+  // Emulate another runtime signaling this connection to receive data.
+  channel.dispatchEvent(
+    new MessageEvent<WebSocketBroadcastChannelMessage>('message', {
+      data: {
+        type: 'extraneous:send',
+        payload: {
+          clientId: connection.id,
+          data: 'hello',
+        },
+      },
+    }),
+  )
+
+  await vi.waitFor(() => {
+    // Must execute the requested operation on the connection.
+    expect(connection.send).toHaveBeenCalledWith('hello')
+    expect(connection.send).toHaveBeenCalledTimes(1)
+  })
+})
+
+it('replays a "close" event coming from another runtime', async () => {
+  const manager = new WebSocketClientManager(channel)
+  const connection = new WebSocketClientConnection(socket, transport)
+  manager.addConnection(connection)
+  vi.spyOn(connection, 'close')
+
+  // Emulate another runtime signaling this connection to close.
+  channel.dispatchEvent(
+    new MessageEvent<WebSocketBroadcastChannelMessage>('message', {
+      data: {
+        type: 'extraneous:close',
+        payload: {
+          clientId: connection.id,
+          code: 1000,
+          reason: 'Normal closure',
+        },
+      },
+    }),
+  )
+
+  await vi.waitFor(() => {
+    // Must execute the requested operation on the connection.
+    expect(connection.close).toHaveBeenCalledWith(1000, 'Normal closure')
+    expect(connection.close).toHaveBeenCalledTimes(1)
+  })
+})
+
+it('removes the extraneous message listener when the connection closes', async () => {
+  const manager = new WebSocketClientManager(channel)
+  const connection = new WebSocketClientConnection(socket, transport)
+  vi.spyOn(connection, 'close').mockImplementationOnce(() => {
+    /**
+     * @note This is a nasty hack so we don't have to uncouple
+     * the connection from transport. Creating a mock transport
+     * is difficult because it relies on the `WebSocketOverride` class.
+     * All we care here is that closing the connection triggers
+     * the transport closure, which it always does.
+     */
+    connection['transport'].onClose()
+  })
+  vi.spyOn(connection, 'send')
+
+  manager.addConnection(connection)
+  connection.close()
+
+  // Signals from other runtimes have no effect on the closed connection.
+  channel.dispatchEvent(
+    new MessageEvent<WebSocketBroadcastChannelMessage>('message', {
+      data: {
+        type: 'extraneous:send',
+        payload: {
+          clientId: connection.id,
+          data: 'hello',
+        },
+      },
+    }),
+  )
+
+  expect(connection.send).not.toHaveBeenCalled()
+})