fix(WebSocket): buffer sending the data until connection is open (#678)

src/interceptors/WebSocket/WebSocketClassTransport.ts

@@ -64,23 +64,44 @@ export class WebSocketClassTransport
 
   public send(data: WebSocketData): void {
     queueMicrotask(() => {
-      this.socket.dispatchEvent(
-        bindEvent(
-          /**
-           * @note Setting this event's "target" to the
-           * WebSocket override instance is important.
-           * This way it can tell apart original incoming events
-           * (must be forwarded to the transport) from the
-           * mocked message events like the one below
-           * (must be dispatched on the client instance).
-           */
-          this.socket,
-          new MessageEvent('message', {
-            data,
-            origin: this.socket.url,
-          })
+      if (
+        this.socket.readyState === this.socket.CLOSING ||
+        this.socket.readyState === this.socket.CLOSED
+      ) {
+        return
+      }
+
+      const dispatchEvent = () => {
+        this.socket.dispatchEvent(
+          bindEvent(
+            /**
+             * @note Setting this event's "target" to the
+             * WebSocket override instance is important.
+             * This way it can tell apart original incoming events
+             * (must be forwarded to the transport) from the
+             * mocked message events like the one below
+             * (must be dispatched on the client instance).
+             */
+            this.socket,
+            new MessageEvent('message', {
+              data,
+              origin: this.socket.url,
+            })
+          )
         )
-      )
+      }
+
+      if (this.socket.readyState === this.socket.CONNECTING) {
+        this.socket.addEventListener(
+          'open',
+          () => {
+            dispatchEvent()
+          },
+          { once: true }
+        )
+      } else {
+        dispatchEvent()
+      }
     })
   }
 

test/modules/WebSocket/compliance/websocket.client.send.test.ts

@@ -0,0 +1,104 @@
+// @vitest-environment node-with-websocket
+import { beforeAll, afterEach, afterAll, vi, it, expect } from 'vitest'
+import { WebSocketInterceptor } from '../../../../src/interceptors/WebSocket'
+
+const interceptor = new WebSocketInterceptor()
+
+beforeAll(() => {
+  interceptor.apply()
+})
+
+afterEach(() => {
+  interceptor.removeAllListeners()
+})
+
+afterAll(() => {
+  interceptor.dispose()
+})
+
+it('buffers client sends until the connection is open', async () => {
+  const events: Array<string> = []
+
+  interceptor.on('connection', ({ client }) => {
+    client.socket.addEventListener('open', () => {
+      events.push('open')
+    })
+    client.socket.addEventListener('message', () => {
+      events.push('send')
+    })
+
+    client.send('hello world')
+  })
+
+  const socket = new WebSocket('ws://localhost')
+  const messageListener = vi.fn()
+  socket.addEventListener('message', messageListener)
+
+  await vi.waitFor(() => {
+    expect(messageListener).toHaveBeenCalledWith(
+      expect.objectContaining({
+        data: 'hello world',
+      })
+    )
+  })
+
+  expect(events).toEqual(['open', 'send'])
+})
+
+it('does not send data if the client connection is closing', async () => {
+  const events: Array<string> = []
+
+  interceptor.on('connection', ({ client }) => {
+    client.socket.addEventListener('close', () => {
+      events.push('close')
+    })
+    client.socket.addEventListener('message', () => {
+      events.push('send')
+    })
+
+    client.close()
+    client.send('hello world')
+  })
+
+  const socket = new WebSocket('ws://localhost')
+  const messageListener = vi.fn()
+  const closeListener = vi.fn()
+  socket.addEventListener('message', messageListener)
+  socket.addEventListener('close', closeListener)
+
+  await vi.waitFor(() => {
+    expect(closeListener).toHaveBeenCalledOnce()
+  })
+
+  expect(messageListener).not.toHaveBeenCalled()
+  expect(events).toEqual(['close'])
+})
+
+it('does not send data if the client connection is closed', async () => {
+  const events: Array<string> = []
+
+  interceptor.on('connection', ({ client }) => {
+    client.socket.addEventListener('close', () => {
+      events.push('close')
+    })
+    client.socket.addEventListener('message', () => {
+      events.push('send')
+    })
+
+    client.close()
+    queueMicrotask(() => client.send('hello world'))
+  })
+
+  const socket = new WebSocket('ws://localhost')
+  const messageListener = vi.fn()
+  const closeListener = vi.fn()
+  socket.addEventListener('message', messageListener)
+  socket.addEventListener('close', closeListener)
+
+  await vi.waitFor(() => {
+    expect(closeListener).toHaveBeenCalledOnce()
+  })
+
+  expect(messageListener).not.toHaveBeenCalled()
+  expect(events).toEqual(['close'])
+})

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

@@ -15,7 +15,7 @@ const wsServer = new WebSocketServer({
   port: 0,
 })
 
-beforeAll(async () => {
+beforeAll(() => {
   interceptor.apply()
 })
 

test/modules/WebSocket/intercept/websocket.dispose.test.ts

@@ -1,6 +1,4 @@
-/**
- * @vitest-environment node-with-websocket
- */
+// @vitest-environment node-with-websocket
 import { vi, it, expect, beforeAll, afterAll } from 'vitest'
 import { WebSocketServer } from 'ws'
 import { WebSocketInterceptor } from '../../../../src/interceptors/WebSocket/index'

test/modules/WebSocket/intercept/websocket.send.test.ts

@@ -1,6 +1,4 @@
-/**
- * @vitest-environment node-with-websocket
- */
+// @vitest-environment node-with-websocket
 import { it, expect, beforeAll, afterAll, afterEach } from 'vitest'
 import { DeferredPromise } from '@open-draft/deferred-promise'
 import { WebSocketInterceptor } from '../../../../src/interceptors/WebSocket'
@@ -30,7 +28,11 @@ it('intercepts text sent over websocket', async () => {
 
   interceptor.once('connection', ({ client }) => {
     client.addEventListener('message', (event) => {
-      messageReceivedPromise.resolve(event.data)
+      if (typeof event.data === 'string') {
+        messageReceivedPromise.resolve(event.data)
+      } else {
+        messageReceivedPromise.reject(new Error('Expected string data'))
+      }
     })
   })
 
@@ -45,7 +47,11 @@ it('intercepts Blob sent over websocket', async () => {
 
   interceptor.once('connection', ({ client }) => {
     client.addEventListener('message', (event) => {
-      messageReceivedPromise.resolve(event.data)
+      if (event.data instanceof Blob) {
+        messageReceivedPromise.resolve(event.data)
+      } else {
+        messageReceivedPromise.reject(new Error('Expected Blob data'))
+      }
     })
   })
 
@@ -57,11 +63,18 @@ it('intercepts Blob sent over websocket', async () => {
 })
 
 it('intercepts ArrayBuffer sent over websocket', async () => {
-  const messageReceivedPromise = new DeferredPromise<ArrayBuffer>()
+  const messageReceivedPromise = new DeferredPromise<Uint8Array>()
 
   interceptor.once('connection', ({ client }) => {
     client.addEventListener('message', (event) => {
-      messageReceivedPromise.resolve(event.data)
+      /**
+       * @note ArrayBuffer data is represented as Buffer in Node.js.
+       */
+      if (event.data instanceof Uint8Array) {
+        messageReceivedPromise.resolve(event.data)
+      } else {
+        messageReceivedPromise.reject(new Error('Expected ArrayBuffer data'))
+      }
     })
   })
 

test/modules/WebSocket/intercept/websocket.server.events.test.ts

@@ -1,6 +1,4 @@
-/**
- * @vitest-environment node-with-websocket
- */
+// @vitest-environment node-with-websocket
 import { vi, it, expect, beforeAll, afterAll } from 'vitest'
 import { WebSocketServer } from 'ws'
 import { WebSocketInterceptor } from '../../../../src/interceptors/WebSocket'