fix(WebSocket): await connection event listeners (#748)

Author: kettanaito

src/interceptors/WebSocket/index.ts

@@ -17,6 +17,7 @@ import {
 } from './WebSocketOverride'
 import { bindEvent } from './utils/bindEvent'
 import { hasConfigurableGlobal } from '../../utils/hasConfigurableGlobal'
+import { emitAsync } from '../../utils/emitAsync'
 
 export { type WebSocketData, WebSocketTransport } from './WebSocketTransport'
 export {
@@ -102,18 +103,21 @@ export class WebSocketInterceptor extends Interceptor<WebSocketEventMap> {
         // Emit the "connection" event to the interceptor on the next tick
         // so the client can modify WebSocket options, like "binaryType"
         // while the connection is already pending.
-        queueMicrotask(() => {
+        queueMicrotask(async () => {
           try {
             const server = new WebSocketServerConnection(
               socket,
               transport,
               createConnection
             )
 
+            const hasConnectionListeners =
+              this.emitter.listenerCount('connection') > 0
+
             // The "globalThis.WebSocket" class stands for
             // the client-side connection. Assume it's established
             // as soon as the WebSocket instance is constructed.
-            const hasConnectionListeners = this.emitter.emit('connection', {
+            await emitAsync(this.emitter, 'connection', {
               client: new WebSocketClientConnection(socket, transport),
               server,
               info: {

test/modules/WebSocket/compliance/websocket.events.test.ts

@@ -8,6 +8,7 @@ import { DeferredPromise } from '@open-draft/deferred-promise'
 import { WebSocketServer } from 'ws'
 import { WebSocketInterceptor } from '../../../../src/interceptors/WebSocket'
 import { getWsUrl } from '../utils/getWsUrl'
+import { sleep } from '../../../helpers'
 
 const wsServer = new WebSocketServer({
   host: '127.0.0.1',
@@ -21,6 +22,7 @@ beforeAll(() => {
 })
 
 afterEach(() => {
+  vi.restoreAllMocks()
   interceptor.removeAllListeners()
   wsServer.removeAllListeners()
   wsServer.clients.forEach((client) => client.close())
@@ -245,6 +247,79 @@ it('emits "error" event on passthrough client connection failure', async () => {
   expect(closeListener).toHaveBeenCalledOnce()
 })
 
+it('allows erroring the connection in a synchronous listener', async () => {
+  vi.spyOn(console, 'error').mockImplementation(() => {})
+
+  interceptor.once('connection', () => {
+    throw new Error('mock error')
+  })
+
+  const ws = new WebSocket('wss://localhost/non-existing-url')
+
+  const openListener = vi.fn()
+  const errorListener = vi.fn()
+  const closeListener = vi.fn()
+  ws.onopen = openListener
+  ws.onerror = errorListener
+  ws.onclose = closeListener
+
+  await expect.poll(() => errorListener).toHaveBeenCalledTimes(1)
+  expect(errorListener).toHaveBeenCalledWith(
+    expect.objectContaining({
+      type: 'error',
+    })
+  )
+
+  await expect.poll(() => ws.readyState).toBe(ws.CLOSED)
+  expect(openListener).not.toHaveBeenCalled()
+  expect(closeListener).toHaveBeenCalledOnce()
+  expect(closeListener).toHaveBeenCalledWith(
+    expect.objectContaining({
+      type: 'close',
+      code: 1011,
+      reason: 'mock error',
+    })
+  )
+})
+
+it('allows erroring the connection from an asynchronous listener', async ({
+  onTestFinished,
+}) => {
+  vi.spyOn(console, 'error').mockImplementation(() => {})
+
+  interceptor.once('connection', async () => {
+    await sleep(200)
+    throw new Error('mock error')
+  })
+
+  const ws = new WebSocket('wss://localhost/non-existing-url')
+
+  const openListener = vi.fn()
+  const errorListener = vi.fn()
+  const closeListener = vi.fn()
+  ws.onopen = openListener
+  ws.onerror = errorListener
+  ws.onclose = closeListener
+
+  await expect.poll(() => errorListener).toHaveBeenCalledTimes(1)
+  expect(errorListener).toHaveBeenCalledWith(
+    expect.objectContaining({
+      type: 'error',
+    })
+  )
+
+  await expect.poll(() => ws.readyState).toBe(ws.CLOSED)
+  expect(openListener).not.toHaveBeenCalled()
+  expect(closeListener).toHaveBeenCalledOnce()
+  expect(closeListener).toHaveBeenCalledWith(
+    expect.objectContaining({
+      type: 'close',
+      code: 1011,
+      reason: 'mock error',
+    })
+  )
+})
+
 it('does not emit "error" event on mocked error code closures', async () => {
   interceptor.once('connection', ({ client }) => {
     /**

test/modules/WebSocket/utils/getWsUrl.ts

@@ -1,7 +1,11 @@
+import { invariant } from 'outvariant'
 import type { WebSocketServer } from 'ws'
 
 export function getWsUrl(ws: WebSocketServer): string {
   const address = ws.address()
+
+  invariant(address != null, 'Failed to get WebSocket address: address is null')
+
   if (typeof address === 'string') {
     return address
   }