From 8c633adc2e047153398e72d1c7cdd9778157f50c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABlle=20Huisman?= <danielle@huisman.me>
Date: Thu, 6 Jun 2024 19:57:03 +0200
Subject: [PATCH] feat(WebSocket): add connection "info" to the "connection"
 event payload (#577)

Co-authored-by: Artem Zakharchenko <kettanaito@gmail.com>
---
 README.md                                     |  9 ++++----
 src/interceptors/WebSocket/index.ts           | 13 +++++++++++
 .../compliance/websocket.connection.test.ts   | 22 +++++++++++++++++++
 3 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 8bc5adb5..ea124ccb 100644
--- a/README.md
+++ b/README.md
@@ -322,10 +322,11 @@ intereceptor.on('connection', ({ client }) => {
 
 The `connection` event exposes the following arguments:
 
-| Name     | Type                                                      | Description                                                      |
-| -------- | --------------------------------------------------------- | ---------------------------------------------------------------- |
-| `client` | [`WebSocketClientConnection`](#websocketclientconnection) | An object representing a connected WebSocket client instance.    |
-| `server` | [`WebSocketServerConnection`](#websocketserverconnection) | An object representing the original WebSocket server connection. |
+| Name     | Type                                                      | Description                                                                         |
+| -------- | --------------------------------------------------------- | ----------------------------------------------------------------------------------- |
+| `client` | [`WebSocketClientConnection`](#websocketclientconnection) | An object representing a connected WebSocket client instance.                       |
+| `server` | [`WebSocketServerConnection`](#websocketserverconnection) | An object representing the original WebSocket server connection.                    |
+| `info`   | `object`                                                  | Additional WebSocket connection information (like the original client `protocols`). |
 
 ### `WebSocketClientConnection`
 
diff --git a/src/interceptors/WebSocket/index.ts b/src/interceptors/WebSocket/index.ts
index 5950c1a4..4407b9f8 100644
--- a/src/interceptors/WebSocket/index.ts
+++ b/src/interceptors/WebSocket/index.ts
@@ -28,6 +28,16 @@ export type WebSocketConnectionData = {
    * The original WebSocket server connection.
    */
   server: WebSocketServerConnection
+
+  /**
+   * The connection information.
+   */
+  info: {
+    /**
+     * The protocols supported by the WebSocket client.
+     */
+    protocols: string | Array<string> | undefined
+  }
 }
 
 /**
@@ -82,6 +92,9 @@ export class WebSocketInterceptor extends Interceptor<WebSocketEventMap> {
               transport,
               createConnection
             ),
+            info: {
+              protocols,
+            },
           })
         })
 
diff --git a/test/modules/WebSocket/compliance/websocket.connection.test.ts b/test/modules/WebSocket/compliance/websocket.connection.test.ts
index 3c0ccad3..4b704593 100644
--- a/test/modules/WebSocket/compliance/websocket.connection.test.ts
+++ b/test/modules/WebSocket/compliance/websocket.connection.test.ts
@@ -57,6 +57,9 @@ it('emits the correct "connection" event on the interceptor', async () => {
         addEventListener: expect.any(Function),
         removeEventListener: expect.any(Function),
       }),
+      info: {
+        protocols: undefined,
+      },
     })
   )
 })
@@ -74,3 +77,22 @@ it('does not connect to the actual WebSocket server by default', async () => {
   expect(connectionListener).toHaveBeenCalledTimes(1)
   expect(realConnectionListener).not.toHaveBeenCalled()
 })
+
+it('includes connection information in the "connection" event payload', async () => {
+  const connectionListener = vi.fn()
+  interceptor.once('connection', connectionListener)
+
+  new WebSocket('wss://example.com', ['protocol1', 'protocol2'])
+  await waitForNextTick()
+
+  expect(connectionListener).toHaveBeenCalledTimes(1)
+  expect(connectionListener).toHaveBeenNthCalledWith(
+    1,
+    expect.objectContaining({
+      info: {
+        // Preserves the client protocols as-is.
+        protocols: ['protocol1', 'protocol2'],
+      },
+    })
+  )
+})