From edf1461c5a90da29c7e708aabddbdf4623a146e2 Mon Sep 17 00:00:00 2001
From: Artem Zakharchenko <kettanaito@gmail.com>
Date: Thu, 3 Oct 2024 13:55:27 +0200
Subject: [PATCH] fix(WebSocket): allow reusing the same listener

---
 .../WebSocket/WebSocketClientConnection.ts    | 21 ++---
 .../WebSocket/WebSocketServerConnection.ts    | 20 ++---
 .../websocket.reuse-listeners.test.ts         | 76 +++++++++++++++++++
 3 files changed, 99 insertions(+), 18 deletions(-)
 create mode 100644 test/modules/WebSocket/compliance/websocket.reuse-listeners.test.ts

diff --git a/src/interceptors/WebSocket/WebSocketClientConnection.ts b/src/interceptors/WebSocket/WebSocketClientConnection.ts
index 3f5e4b61..0e1749cf 100644
--- a/src/interceptors/WebSocket/WebSocketClientConnection.ts
+++ b/src/interceptors/WebSocket/WebSocketClientConnection.ts
@@ -86,18 +86,21 @@ export class WebSocketClientConnection
     listener: WebSocketEventListener<WebSocketClientEventMap[EventType]>,
     options?: AddEventListenerOptions | boolean
   ): void {
-    const boundListener = listener.bind(this.socket)
-
-    // Store the bound listener on the original listener
-    // so the exact bound function can be accessed in "removeEventListener()".
-    Object.defineProperty(listener, kBoundListener, {
-      value: boundListener,
-      enumerable: false,
-    })
+    if (!Reflect.has(listener, kBoundListener)) {
+      const boundListener = listener.bind(this.socket)
+
+      // Store the bound listener on the original listener
+      // so the exact bound function can be accessed in "removeEventListener()".
+      Object.defineProperty(listener, kBoundListener, {
+        value: boundListener,
+        enumerable: false,
+        configurable: false,
+      })
+    }
 
     this[kEmitter].addEventListener(
       type,
-      boundListener as EventListener,
+      Reflect.get(listener, kBoundListener) as EventListener,
       options
     )
   }
diff --git a/src/interceptors/WebSocket/WebSocketServerConnection.ts b/src/interceptors/WebSocket/WebSocketServerConnection.ts
index 0f90d9f6..e4c72c7a 100644
--- a/src/interceptors/WebSocket/WebSocketServerConnection.ts
+++ b/src/interceptors/WebSocket/WebSocketServerConnection.ts
@@ -169,18 +169,20 @@ export class WebSocketServerConnection {
     listener: WebSocketEventListener<WebSocketServerEventMap[EventType]>,
     options?: AddEventListenerOptions | boolean
   ): void {
-    const boundListener = listener.bind(this.client)
-
-    // Store the bound listener on the original listener
-    // so the exact bound function can be accessed in "removeEventListener()".
-    Object.defineProperty(listener, kBoundListener, {
-      value: boundListener,
-      enumerable: false,
-    })
+    if (!Reflect.has(listener, kBoundListener)) {
+      const boundListener = listener.bind(this.client)
+
+      // Store the bound listener on the original listener
+      // so the exact bound function can be accessed in "removeEventListener()".
+      Object.defineProperty(listener, kBoundListener, {
+        value: boundListener,
+        enumerable: false,
+      })
+    }
 
     this[kEmitter].addEventListener(
       event,
-      boundListener as EventListener,
+      Reflect.get(listener, kBoundListener) as EventListener,
       options
     )
   }
diff --git a/test/modules/WebSocket/compliance/websocket.reuse-listeners.test.ts b/test/modules/WebSocket/compliance/websocket.reuse-listeners.test.ts
new file mode 100644
index 00000000..d833d0c7
--- /dev/null
+++ b/test/modules/WebSocket/compliance/websocket.reuse-listeners.test.ts
@@ -0,0 +1,76 @@
+// @vitest-environment node-with-websocket
+import { vi, it, expect, beforeAll, afterEach, afterAll } from 'vitest'
+import { WebSocketServer } from 'ws'
+import { WebSocketInterceptor } from '../../../../src/interceptors/WebSocket'
+import { waitForWebSocketEvent } from '../utils/waitForWebSocketEvent'
+import { getWsUrl } from '../utils/getWsUrl'
+
+const wsServer = new WebSocketServer({
+  host: '127.0.0.1',
+  port: 0,
+})
+
+const interceptor = new WebSocketInterceptor()
+
+beforeAll(async () => {
+  interceptor.apply()
+})
+
+afterEach(() => {
+  interceptor.removeAllListeners()
+  wsServer.clients.forEach((client) => client.close())
+})
+
+afterAll(async () => {
+  interceptor.dispose()
+  wsServer.close()
+})
+
+it('allows reusing the same function for multiple client listeners', async () => {
+  const clientMessageListener = vi.fn()
+
+  interceptor.on('connection', ({ client }) => {
+    client.addEventListener('message', clientMessageListener)
+    client.addEventListener('message', clientMessageListener)
+    client.addEventListener('message', clientMessageListener)
+
+    queueMicrotask(() => client.close())
+  })
+
+  const socket = new WebSocket('wss://example.com')
+  socket.onopen = () => socket.send('hello world')
+
+  await waitForWebSocketEvent('close', socket)
+
+  /**
+   * @note The same event listner for the same event is deduped.
+   * It will only be called once. That is correct.
+   */
+  expect(clientMessageListener).toHaveBeenCalledTimes(1)
+})
+
+it('allows reusing the same function for multiple server listeners', async () => {
+  wsServer.once('connection', (ws) => {
+    ws.send('hello from server')
+    queueMicrotask(() => ws.close())
+  })
+
+  const serverMessageListener = vi.fn()
+
+  interceptor.on('connection', ({ server }) => {
+    server.connect()
+    server.addEventListener('message', serverMessageListener)
+    server.addEventListener('message', serverMessageListener)
+    server.addEventListener('message', serverMessageListener)
+  })
+
+  const socket = new WebSocket(getWsUrl(wsServer))
+
+  await waitForWebSocketEvent('close', socket)
+
+  /**
+   * @note The same event listner for the same event is deduped.
+   * It will only be called once. That is correct.
+   */
+  expect(serverMessageListener).toHaveBeenCalledTimes(1)
+})