feat: support ws.onUpgrade for handling connection upgrades

Author: kettanaito

src/core/experimental/frames/http-frame.test.ts

@@ -4,7 +4,8 @@ import { ws } from '../../ws'
 import { bypass } from '../../bypass'
 import type { HttpNetworkFrameEventMap } from './http-frame'
 import { HttpNetworkFrame } from './http-frame'
-import { InMemoryHandlersController } from '#core/experimental/handlers-controller'
+import { InMemoryHandlersController } from '../../experimental/handlers-controller'
+import { getSiblingHandlers } from '../../utils/internal/attachSiblingHandlers'
 
 beforeAll(() => {
   vi.spyOn(console, 'error').mockImplementation(() => {})
@@ -46,6 +47,9 @@ it('filters only request type handlers', async () => {
   const webSocketHandlers = [
     ws.link('ws://localhost').addEventListener('connection', () => {}),
   ]
+  const webSocketSiblingHandlers = webSocketHandlers.flatMap((handler) =>
+    getSiblingHandlers(handler),
+  )
 
   const controller = new InMemoryHandlersController([
     ...httpHandlers,
@@ -55,6 +59,7 @@ it('filters only request type handlers', async () => {
 
   expect(frame.getHandlers(controller)).toEqual([
     ...httpHandlers,
+    ...webSocketSiblingHandlers,
     ...graphqlHandlers,
   ])
   expect(frame.getHandlers(new InMemoryHandlersController([]))).toEqual([])

src/core/experimental/handlers-controller.test.ts

@@ -1,8 +1,56 @@
 import { http } from '../http'
 import { graphql } from '../graphql'
 import { ws } from '../ws'
+import { getSiblingHandlers } from '../utils/internal/attachSiblingHandlers'
 import { InMemoryHandlersController } from './handlers-controller'
 
+describe('constructor', () => {
+  it('places the sibling in its own kind bucket', () => {
+    const wsHandler = ws.link('*').addEventListener('connection', () => {})
+    const [upgradeHandler] = getSiblingHandlers(wsHandler)
+
+    const controller = new InMemoryHandlersController([wsHandler])
+
+    expect(controller.getHandlersByKind('websocket')).toEqual([wsHandler])
+    expect(controller.getHandlersByKind('request')).toEqual([upgradeHandler])
+  })
+
+  it('interleaves the sibling at the owner position when grouping by kind', () => {
+    const httpOne = http.get('/', () => {})
+    const wsHandler = ws.link('*').addEventListener('connection', () => {})
+    const [upgradeHandler] = getSiblingHandlers(wsHandler)
+    const httpTwo = http.get('/', () => {})
+
+    const controller = new InMemoryHandlersController([
+      httpOne,
+      wsHandler,
+      httpTwo,
+    ])
+
+    expect(controller.getHandlersByKind('request')).toEqual([
+      httpOne,
+      upgradeHandler,
+      httpTwo,
+    ])
+    expect(controller.getHandlersByKind('websocket')).toEqual([wsHandler])
+  })
+
+  it('extracts siblings from every owner in the input list', () => {
+    const wsOne = ws.link('*').addEventListener('connection', () => {})
+    const wsTwo = ws.link('*').addEventListener('connection', () => {})
+    const [upgradeOne] = getSiblingHandlers(wsOne)
+    const [upgradeTwo] = getSiblingHandlers(wsTwo)
+
+    const controller = new InMemoryHandlersController([wsOne, wsTwo])
+
+    expect(controller.getHandlersByKind('websocket')).toEqual([wsOne, wsTwo])
+    expect(controller.getHandlersByKind('request')).toEqual([
+      upgradeOne,
+      upgradeTwo,
+    ])
+  })
+})
+
 describe(InMemoryHandlersController.prototype.use, () => {
   it('prepends a handler to an empty controller', () => {
     const controller = new InMemoryHandlersController([])
@@ -51,6 +99,31 @@ describe(InMemoryHandlersController.prototype.use, () => {
 
     expect(controller.currentHandlers()).toEqual([graphqlOne, httpTwo, httpOne])
   })
+
+  it('propagates siblings to their kind buckets at runtime', () => {
+    const controller = new InMemoryHandlersController([])
+    const wsHandler = ws.link('*').addEventListener('connection', () => {})
+    const [upgradeHandler] = getSiblingHandlers(wsHandler)
+
+    controller.use([wsHandler])
+
+    expect(controller.getHandlersByKind('websocket')).toEqual([wsHandler])
+    expect(controller.getHandlersByKind('request')).toEqual([upgradeHandler])
+  })
+
+  it('prepends incoming siblings before existing handlers of the same kind', () => {
+    const existingHttp = http.get('/existing', () => {})
+    const controller = new InMemoryHandlersController([existingHttp])
+    const wsHandler = ws.link('*').addEventListener('connection', () => {})
+    const [upgradeHandler] = getSiblingHandlers(wsHandler)
+
+    controller.use([wsHandler])
+
+    expect(controller.getHandlersByKind('request')).toEqual([
+      upgradeHandler,
+      existingHttp,
+    ])
+  })
 })
 
 describe(InMemoryHandlersController.prototype.reset, () => {
@@ -96,6 +169,29 @@ describe(InMemoryHandlersController.prototype.reset, () => {
      */
     expect(controller.currentHandlers()).toEqual([httpTwo])
   })
+
+  it('places siblings into their kind buckets when resetting to next handlers', () => {
+    const controller = new InMemoryHandlersController([])
+    const wsHandler = ws.link('*').addEventListener('connection', () => {})
+    const [upgradeHandler] = getSiblingHandlers(wsHandler)
+
+    controller.reset([wsHandler])
+
+    expect(controller.getHandlersByKind('websocket')).toEqual([wsHandler])
+    expect(controller.getHandlersByKind('request')).toEqual([upgradeHandler])
+  })
+
+  it('restores siblings when resetting to the initial handlers', () => {
+    const wsHandler = ws.link('*').addEventListener('connection', () => {})
+    const [upgradeHandler] = getSiblingHandlers(wsHandler)
+    const controller = new InMemoryHandlersController([wsHandler])
+
+    controller.use([http.get('/runtime', () => {})])
+    controller.reset([])
+
+    expect(controller.getHandlersByKind('websocket')).toEqual([wsHandler])
+    expect(controller.getHandlersByKind('request')).toEqual([upgradeHandler])
+  })
 })
 
 describe(InMemoryHandlersController.prototype.getHandlersByKind, () => {
@@ -112,11 +208,10 @@ describe(InMemoryHandlersController.prototype.getHandlersByKind, () => {
       ]).getHandlersByKind('websocket'),
     ).toEqual([])
 
+    const wsHandler = ws.link('*').addEventListener('connection', () => {})
     expect(
-      new InMemoryHandlersController([
-        ws.link('*').addEventListener('connection', () => {}),
-      ]).getHandlersByKind('request'),
-    ).toEqual([])
+      new InMemoryHandlersController([wsHandler]).getHandlersByKind('request'),
+    ).toEqual(getSiblingHandlers(wsHandler))
   })
 
   it('returns all handlers if they all match', () => {
@@ -142,14 +237,15 @@ describe(InMemoryHandlersController.prototype.getHandlersByKind, () => {
     const httpHandler = http.get('/', () => {})
     const graphqlHandler = graphql.query('', () => {})
     const wsHandler = ws.link('*').addEventListener('connection', () => {})
+    const wsHandlerSiblings = getSiblingHandlers(wsHandler)
 
     expect(
       new InMemoryHandlersController([
         httpHandler,
         graphqlHandler,
         wsHandler,
       ]).getHandlersByKind('request'),
-    ).toEqual([httpHandler, graphqlHandler])
+    ).toEqual([httpHandler, graphqlHandler, ...wsHandlerSiblings])
 
     expect(
       new InMemoryHandlersController([

src/core/experimental/handlers-controller.ts

@@ -2,6 +2,7 @@ import { invariant } from 'outvariant'
 import { type RequestHandler } from '../handlers/RequestHandler'
 import { type WebSocketHandler } from '../handlers/WebSocketHandler'
 import { devUtils } from '../utils/internal/devUtils'
+import { getSiblingHandlers } from '../utils/internal/attachSiblingHandlers'
 
 export type AnyHandler = RequestHandler | WebSocketHandler
 export type HandlersMap = Partial<Record<AnyHandler['kind'], Array<AnyHandler>>>
@@ -14,6 +15,10 @@ export function groupHandlersByKind(handlers: Array<AnyHandler>): HandlersMap {
    */
   for (const handler of handlers) {
     ;(groups[handler.kind] ||= []).push(handler)
+
+    for (const sibling of getSiblingHandlers(handler)) {
+      ;(groups[sibling.kind] ||= []).push(sibling)
+    }
   }
 
   return groups
@@ -69,14 +74,16 @@ export abstract class HandlersController {
     }
 
     const { handlers } = this.getState()
-
-    // Iterate over next handlers and prepend them to their respective lists.
-    // Iterate in a reverse order to the keep the order of the runtime handlers as provided.
-    for (let i = nextHandlers.length - 1; i >= 0; i--) {
-      const handler = nextHandlers[i]
-      handlers[handler.kind] = handlers[handler.kind]
-        ? [handler, ...handlers[handler.kind]!]
-        : [handler]
+    const overrides = groupHandlersByKind(nextHandlers)
+
+    // Prepend overrides to their respective kind buckets so they take
+    // priority over existing handlers while preserving input order.
+    for (const kind in overrides) {
+      const overridesForKind = overrides[kind as AnyHandler['kind']]!
+      const existingForKind = handlers[kind as AnyHandler['kind']]
+      handlers[kind as AnyHandler['kind']] = existingForKind
+        ? [...overridesForKind, ...existingForKind]
+        : overridesForKind
     }
 
     this.setState({ handlers })

src/core/utils/internal/attachSiblingHandlers.ts

@@ -0,0 +1,28 @@
+import { invariant } from 'outvariant'
+import type { AnyHandler } from '../../experimental/handlers-controller'
+
+const kSiblingHandlers = Symbol('kSiblingHandlers')
+
+export function attachSiblingHandlers<T extends AnyHandler>(
+  owner: T,
+  siblings: Array<AnyHandler>,
+): T {
+  invariant(
+    getSiblingHandlers(owner).length === 0,
+    'Failed to merge handlers: the owner "%s" handler is already merged',
+    owner.kind,
+  )
+
+  Object.defineProperty(owner, kSiblingHandlers, {
+    value: siblings,
+    enumerable: false,
+    writable: false,
+    configurable: false,
+  })
+
+  return owner
+}
+
+export function getSiblingHandlers(owner: AnyHandler): Array<AnyHandler> {
+  return Reflect.get(owner, kSiblingHandlers) || []
+}

src/core/ws.ts

@@ -1,16 +1,24 @@
 import { invariant } from 'outvariant'
+import { FetchResponse } from '@mswjs/interceptors'
 import type {
   WebSocketData,
   WebSocketClientConnectionProtocol,
 } from '@mswjs/interceptors/WebSocket'
+import type { ResponseResolver } from './handlers/RequestHandler'
 import {
   WebSocketHandler,
   kEmitter,
   type WebSocketHandlerEventMap,
 } from './handlers/WebSocketHandler'
 import { hasRefCounted } from './utils/internal/hasRefCounted'
-import { type Path, isPath } from './utils/matching/matchRequestUrl'
+import {
+  type Path,
+  isPath,
+  matchRequestUrl,
+} from './utils/matching/matchRequestUrl'
 import { WebSocketClientManager } from './ws/WebSocketClientManager'
+import { http } from './http'
+import { attachSiblingHandlers } from './utils/internal/attachSiblingHandlers'
 
 const webSocketChannel = new BroadcastChannel('msw:websocket-client-manager')
 
@@ -105,23 +113,30 @@ function createWebSocketLinkHandler(url: Path): WebSocketLink {
       return clientManager.clients
     },
     addEventListener(event, listener) {
-      const handler = new WebSocketHandler(url)
+      const webSocketHandler = 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', async ({ client }) => {
+      webSocketHandler[kEmitter].on('connection', async ({ client }) => {
         await 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)
+      webSocketHandler[kEmitter].on(event, listener)
+
+      const upgradeHandler = http.get(({ request }) => {
+        return (
+          request.headers.get('upgrade') === 'websocket' &&
+          matchRequestUrl(new URL(request.url), url).matches
+        )
+      }, ws.onUpgrade)
 
-      return handler
+      return attachSiblingHandlers(webSocketHandler, [upgradeHandler])
     },
 
     broadcast(data) {
@@ -145,6 +160,13 @@ function createWebSocketLinkHandler(url: Path): WebSocketLink {
   }
 }
 
+const WEBSOCKET_GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
+
+interface WebSocketNamespace {
+  link: typeof createWebSocketLinkHandler
+  onUpgrade: ResponseResolver
+}
+
 /**
  * A namespace to intercept and mock WebSocket connections.
  *
@@ -154,8 +176,28 @@ function createWebSocketLinkHandler(url: Path): WebSocketLink {
  * @see {@link https://mswjs.io/docs/api/ws `ws` API reference}
  * @see {@link https://mswjs.io/docs/basics/handling-websocket-events Handling WebSocket events}
  */
-export const ws = {
+export const ws: WebSocketNamespace = {
   link: createWebSocketLinkHandler,
+  async onUpgrade({ request }) {
+    const key = request.headers.get('sec-websocket-key')
+
+    if (!key) {
+      return
+    }
+
+    const keyBytes = new TextEncoder().encode(key + WEBSOCKET_GUID)
+    const digest = await crypto.subtle.digest('SHA-1', keyBytes)
+    const acceptValue = btoa(String.fromCharCode(...new Uint8Array(digest)))
+
+    return new FetchResponse(null, {
+      status: 101,
+      headers: {
+        upgrade: 'websocket',
+        connection: 'upgrade',
+        'sec-websocket-accept': acceptValue,
+      },
+    })
+  },
 }
 
 export { type WebSocketData }