msw 2.2.3 → 2.3.0-ws.rc-1

package/src/core/ws/WebSocketClientManager.ts

@@ -0,0 +1,170 @@
+import type {
+  WebSocketData,
+  WebSocketClientConnection,
+  WebSocketClientConnectionProtocol,
+} from '@mswjs/interceptors/WebSocket'
+
+export type WebSocketBroadcastChannelMessage =
+  | {
+      type: 'connection:open'
+      payload: {
+        clientId: string
+        url: string
+      }
+    }
+  | {
+      type: 'extraneous:send'
+      payload: {
+        clientId: string
+        data: WebSocketData
+      }
+    }
+  | {
+      type: 'extraneous:close'
+      payload: {
+        clientId: string
+        code?: number
+        reason?: string
+      }
+    }
+
+export const kAddByClientId = Symbol('kAddByClientId')
+
+/**
+ * A manager responsible for accumulating WebSocket client
+ * connections across different browser runtimes.
+ */
+export class WebSocketClientManager {
+  /**
+   * All active WebSocket client connections.
+   */
+  public clients: Set<WebSocketClientConnectionProtocol>
+
+  constructor(private channel: BroadcastChannel) {
+    this.clients = new Set()
+
+    this.channel.addEventListener('message', (message) => {
+      const { type, payload } = message.data as WebSocketBroadcastChannelMessage
+
+      switch (type) {
+        case 'connection:open': {
+          // When another runtime notifies about a new connection,
+          // create a connection wrapper class and add it to the set.
+          this.onRemoteConnection(payload.clientId, new URL(payload.url))
+          break
+        }
+      }
+    })
+  }
+
+  /**
+   * Adds the given `WebSocket` client connection to the set
+   * of all connections. The given connection is always the complete
+   * connection object because `addConnection()` is called only
+   * for the opened connections in the same runtime.
+   */
+  public addConnection(client: WebSocketClientConnection): void {
+    this.clients.add(client)
+
+    // Signal to other runtimes about this connection.
+    this.channel.postMessage({
+      type: 'connection:open',
+      payload: {
+        clientId: client.id,
+        url: client.url.toString(),
+      },
+    } as WebSocketBroadcastChannelMessage)
+
+    // Instruct the current client how to handle events
+    // coming from other runtimes (e.g. when calling `.broadcast()`).
+    const handleExtraneousMessage = (
+      message: MessageEvent<WebSocketBroadcastChannelMessage>,
+    ) => {
+      const { type, payload } = message.data
+
+      // Ignore broadcasted messages for other clients.
+      if (
+        typeof payload === 'object' &&
+        'clientId' in payload &&
+        payload.clientId !== client.id
+      ) {
+        return
+      }
+
+      switch (type) {
+        case 'extraneous:send': {
+          client.send(payload.data)
+          break
+        }
+
+        case 'extraneous:close': {
+          client.close(payload.code, payload.reason)
+          break
+        }
+      }
+    }
+
+    const abortController = new AbortController()
+
+    this.channel.addEventListener('message', handleExtraneousMessage, {
+      signal: abortController.signal,
+    })
+
+    // Once closed, this connection cannot be operated on.
+    // This must include the extraneous runtimes as well.
+    client.addEventListener('close', () => abortController.abort(), {
+      once: true,
+    })
+  }
+
+  /**
+   * Adds a client connection wrapper to operate with
+   * WebSocket client connections in other runtimes.
+   */
+  private onRemoteConnection(id: string, url: URL): void {
+    this.clients.add(
+      // Create a connection-compatible instance that can
+      // operate with this client from a different runtime
+      // using the BroadcastChannel messages.
+      new WebSocketRemoteClientConnection(id, url, this.channel),
+    )
+  }
+}
+
+/**
+ * A wrapper class to operate with WebSocket client connections
+ * from other runtimes. This class maintains 1-1 public API
+ * compatibility to the `WebSocketClientConnection` but relies
+ * on the given `BroadcastChannel` to communicate instructions
+ * with the client connections from other runtimes.
+ */
+export class WebSocketRemoteClientConnection
+  implements WebSocketClientConnectionProtocol
+{
+  constructor(
+    public readonly id: string,
+    public readonly url: URL,
+    private channel: BroadcastChannel,
+  ) {}
+
+  send(data: WebSocketData): void {
+    this.channel.postMessage({
+      type: 'extraneous:send',
+      payload: {
+        clientId: this.id,
+        data,
+      },
+    } as WebSocketBroadcastChannelMessage)
+  }
+
+  close(code?: number | undefined, reason?: string | undefined): void {
+    this.channel.postMessage({
+      type: 'extraneous:close',
+      payload: {
+        clientId: this.id,
+        code,
+        reason,
+      },
+    } as WebSocketBroadcastChannelMessage)
+  }
+}

package/src/core/ws/webSocketInterceptor.ts

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

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

@@ -0,0 +1,23 @@
+/**
+ * @vitest-environment node-websocket
+ */
+import { ws } from './ws'
+
+it('exports the "link()" method', () => {
+  expect(ws).toHaveProperty('link')
+  expect(ws.link).toBeInstanceOf(Function)
+})
+
+it('throws an error when calling "ws.link()" without a URL argument', () => {
+  expect(() =>
+    // @ts-expect-error Intentionally invalid call.
+    ws.link(),
+  ).toThrow('Expected a WebSocket server URL but got undefined')
+})
+
+it('throws an error when given a non-path argument to "ws.link()"', () => {
+  expect(() =>
+    // @ts-expect-error Intentionally invalid argument.
+    ws.link(2),
+  ).toThrow('Expected a WebSocket server URL but got number')
+})

package/src/core/ws/ws.ts

@@ -0,0 +1,108 @@
+import { invariant } from 'outvariant'
+import type {
+  WebSocketClientConnectionProtocol,
+  WebSocketData,
+} from '@mswjs/interceptors/WebSocket'
+import {
+  WebSocketHandler,
+  kEmitter,
+  type WebSocketHandlerEventMap,
+} from '../handlers/WebSocketHandler'
+import { Path, isPath } from '../utils/matching/matchRequestUrl'
+import { WebSocketClientManager } from './WebSocketClientManager'
+
+const wsBroadcastChannel = new BroadcastChannel('msw:ws-client-manager')
+
+/**
+ * Intercepts outgoing WebSocket connections to the given URL.
+ *
+ * @example
+ * const chat = ws.link('wss://chat.example.com')
+ * chat.on('connection', ({ client }) => {
+ *   client.send('hello from server!')
+ * })
+ */
+function createWebSocketLinkHandler(url: Path) {
+  invariant(url, 'Expected a WebSocket server URL but got undefined')
+
+  invariant(
+    isPath(url),
+    'Expected a WebSocket server URL but got %s',
+    typeof url,
+  )
+
+  const clientManager = new WebSocketClientManager(wsBroadcastChannel)
+
+  return {
+    clients: clientManager.clients,
+    on<K extends keyof WebSocketHandlerEventMap>(
+      event: K,
+      listener: (...args: WebSocketHandlerEventMap[K]) => void,
+    ): WebSocketHandler {
+      const handler = new WebSocketHandler(url)
+
+      // Add the connection event listener for when the
+      // handler matches and emits a connection event.
+      // When that happens, store that connection in the
+      // set of all connections for reference.
+      handler[kEmitter].on('connection', ({ client }) => {
+        clientManager.addConnection(client)
+      })
+
+      // The "handleWebSocketEvent" function will invoke
+      // the "run()" method on the WebSocketHandler.
+      // If the handler matches, it will emit the "connection"
+      // event. Attach the user-defined listener to that event.
+      handler[kEmitter].on(event, listener)
+
+      return handler
+    },
+
+    /**
+     * Broadcasts the given data to all WebSocket clients.
+     *
+     * @example
+     * const service = ws.link('wss://example.com')
+     * service.on('connection', () => {
+     *   service.broadcast('hello, everyone!')
+     * })
+     */
+    broadcast(data: WebSocketData): void {
+      // This will invoke "send()" on the immediate clients
+      // in this runtime and post a message to the broadcast channel
+      // to trigger send for the clients in other runtimes.
+      this.broadcastExcept([], data)
+    },
+
+    /**
+     * Broadcasts the given data to all WebSocket clients
+     * except the ones provided in the `clients` argument.
+     *
+     * @example
+     * const service = ws.link('wss://example.com')
+     * service.on('connection', ({ client }) => {
+     *   service.broadcastExcept(client, 'hi, the rest of you!')
+     * })
+     */
+    broadcastExcept(
+      clients:
+        | WebSocketClientConnectionProtocol
+        | Array<WebSocketClientConnectionProtocol>,
+      data: WebSocketData,
+    ): void {
+      const ignoreClients = Array.prototype
+        .concat(clients)
+        .map((client) => client.id)
+
+      clientManager.clients.forEach((otherClient) => {
+        if (!ignoreClients.includes(otherClient.id)) {
+          otherClient.send(data)
+        }
+      })
+    },
+  }
+}
+
+export const ws = {
+  link: createWebSocketLinkHandler,
+}

package/src/node/SetupServerApi.ts

@@ -4,14 +4,15 @@ import { XMLHttpRequestInterceptor } from '@mswjs/interceptors/XMLHttpRequest'
 import { FetchInterceptor } from '@mswjs/interceptors/fetch'
 import { HandlersController } from '~/core/SetupApi'
 import type { RequestHandler } from '~/core/handlers/RequestHandler'
+import type { WebSocketHandler } from '~/core/handlers/WebSocketHandler'
 import type { SetupServer } from './glossary'
 import { SetupServerCommonApi } from './SetupServerCommonApi'
 
 const store = new AsyncLocalStorage<RequestHandlersContext>()
 
 type RequestHandlersContext = {
-  initialHandlers: Array<RequestHandler>
-  handlers: Array<RequestHandler>
+  initialHandlers: Array<RequestHandler | WebSocketHandler>
+  handlers: Array<RequestHandler | WebSocketHandler>
 }
 
 /**
@@ -22,7 +23,7 @@ type RequestHandlersContext = {
 class AsyncHandlersController implements HandlersController {
   private rootContext: RequestHandlersContext
 
-  constructor(initialHandlers: Array<RequestHandler>) {
+  constructor(initialHandlers: Array<RequestHandler | WebSocketHandler>) {
     this.rootContext = { initialHandlers, handlers: [] }
   }
 
@@ -30,18 +31,18 @@ class AsyncHandlersController implements HandlersController {
     return store.getStore() || this.rootContext
   }
 
-  public prepend(runtimeHandlers: Array<RequestHandler>) {
+  public prepend(runtimeHandlers: Array<RequestHandler | WebSocketHandler>) {
     this.context.handlers.unshift(...runtimeHandlers)
   }
 
-  public reset(nextHandlers: Array<RequestHandler>) {
+  public reset(nextHandlers: Array<RequestHandler | WebSocketHandler>) {
     const context = this.context
     context.handlers = []
     context.initialHandlers =
       nextHandlers.length > 0 ? nextHandlers : context.initialHandlers
   }
 
-  public currentHandlers(): Array<RequestHandler> {
+  public currentHandlers(): Array<RequestHandler | WebSocketHandler> {
     const { initialHandlers, handlers } = this.context
     return handlers.concat(initialHandlers)
   }
@@ -51,7 +52,7 @@ export class SetupServerApi
   extends SetupServerCommonApi
   implements SetupServer
 {
-  constructor(handlers: Array<RequestHandler>) {
+  constructor(handlers: Array<RequestHandler | WebSocketHandler>) {
     super(
       [ClientRequestInterceptor, XMLHttpRequestInterceptor, FetchInterceptor],
       handlers,

package/src/node/SetupServerCommonApi.ts

@@ -14,9 +14,12 @@ import type { LifeCycleEventsMap, SharedOptions } from '~/core/sharedOptions'
 import { SetupApi } from '~/core/SetupApi'
 import { handleRequest } from '~/core/utils/handleRequest'
 import type { RequestHandler } from '~/core/handlers/RequestHandler'
+import type { WebSocketHandler } from '~/core/handlers/WebSocketHandler'
 import { mergeRight } from '~/core/utils/internal/mergeRight'
 import { devUtils } from '~/core/utils/internal/devUtils'
 import type { SetupServerCommon } from './glossary'
+import { handleWebSocketEvent } from '~/core/utils/handleWebSocketEvent'
+import { webSocketInterceptor } from '~/core/ws/webSocketInterceptor'
 
 export const DEFAULT_LISTEN_OPTIONS: RequiredDeep<SharedOptions> = {
   onUnhandledRequest: 'warn',
@@ -34,7 +37,7 @@ export class SetupServerCommonApi
 
   constructor(
     interceptors: Array<{ new (): Interceptor<HttpRequestEventMap> }>,
-    handlers: Array<RequestHandler>,
+    handlers: Array<RequestHandler | WebSocketHandler>,
   ) {
     super(...handlers)
 
@@ -81,6 +84,10 @@ export class SetupServerCommonApi
         )
       },
     )
+
+    handleWebSocketEvent(() => {
+      return this.handlersController.currentHandlers()
+    })
   }
 
   public listen(options: Partial<SharedOptions> = {}): void {
@@ -91,9 +98,11 @@ export class SetupServerCommonApi
 
     // Apply the interceptor when starting the server.
     this.interceptor.apply()
+    webSocketInterceptor.apply()
 
     this.subscriptions.push(() => {
       this.interceptor.dispose()
+      webSocketInterceptor.dispose()
     })
 
     // Assert that the interceptor has been applied successfully.

package/src/node/glossary.ts

@@ -1,8 +1,6 @@
 import type { PartialDeep } from 'type-fest'
-import type {
-  RequestHandler,
-  RequestHandlerDefaultInfo,
-} from '~/core/handlers/RequestHandler'
+import type { RequestHandler } from '~/core/handlers/RequestHandler'
+import type { WebSocketHandler } from '~/core/handlers/WebSocketHandler'
 import type {
   LifeCycleEventEmitter,
   LifeCycleEventsMap,
@@ -29,7 +27,7 @@ export interface SetupServerCommon {
    *
    * @see {@link https://mswjs.io/docs/api/setup-server/use `server.use()` API reference}
    */
-  use(...handlers: Array<RequestHandler>): void
+  use(...handlers: Array<RequestHandler | WebSocketHandler>): void
 
   /**
    * Marks all request handlers that respond using `res.once()` as unused.
@@ -43,14 +41,14 @@ export interface SetupServerCommon {
    *
    * @see {@link https://mswjs.io/docs/api/setup-server/reset-handlers `server.reset-handlers()` API reference}
    */
-  resetHandlers(...nextHandlers: Array<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-server/list-handlers `server.listHandlers()` API reference}
    */
-  listHandlers(): ReadonlyArray<RequestHandler<RequestHandlerDefaultInfo, any>>
+  listHandlers(): ReadonlyArray<RequestHandler | WebSocketHandler>
 
   /**
    * Life-cycle events.

package/src/node/setupServer.ts

@@ -1,4 +1,5 @@
 import type { RequestHandler } from '~/core/handlers/RequestHandler'
+import type { WebSocketHandler } from '~/core/handlers/WebSocketHandler'
 import { SetupServerApi } from './SetupServerApi'
 
 /**
@@ -8,7 +9,7 @@ import { SetupServerApi } from './SetupServerApi'
  * @see {@link https://mswjs.io/docs/api/setup-server `setupServer()` API reference}
  */
 export const setupServer = (
-  ...handlers: Array<RequestHandler>
+  ...handlers: Array<RequestHandler | WebSocketHandler>
 ): SetupServerApi => {
   return new SetupServerApi(handlers)
 }