From b3b1e8af54d5db45d61415085ebf5376755f374c Mon Sep 17 00:00:00 2001
From: naporitan <naporin0624@gmail.com>
Date: Tue, 6 Aug 2024 20:13:02 +0900
Subject: [PATCH] feat(client): Add WebSocket Provider Integration Tests and
 Enhance WebSocket Initialization (#3213)

* chore(client): add websocket provider testing

* feat(client): add websocket provider options

* chore(client): same function format

* chore(client): remove $ws method options

* chore(client): remove console.log
---
 src/client/client.test.ts | 49 +++++++++++++++++++++++++++++++++++++++
 src/client/client.ts      |  8 ++++++-
 src/client/types.ts       |  1 +
 3 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/src/client/client.test.ts b/src/client/client.test.ts
index e9c91877b..5e0b5218f 100644
--- a/src/client/client.test.ts
+++ b/src/client/client.test.ts
@@ -1240,3 +1240,52 @@ describe('Redirect response - only types', () => {
     }
   })
 })
+
+describe('WebSocket Provider Integration', () => {
+  const app = new Hono()
+  const route = app.get(
+    '/',
+    upgradeWebSocket((c) => ({
+      onMessage(event, ws) {
+        ws.send('Hello from server!')
+      },
+      onClose() {
+        console.log('Connection closed')
+      },
+    }))
+  )
+
+  type AppType = typeof route
+
+  const server = setupServer()
+  beforeAll(() => server.listen())
+  afterEach(() => {
+    vi.clearAllMocks()
+    server.resetHandlers()
+  })
+  afterAll(() => server.close())
+
+  it.each([
+    {
+      description: 'should initialize the WebSocket provider correctly',
+      url: 'http://localhost',
+      query: undefined,
+      expectedUrl: 'ws://localhost/index',
+    },
+    {
+      description: 'should correctly add query parameters to the WebSocket URL',
+      url: 'http://localhost',
+      query: { id: '123', type: 'test', tag: ['a', 'b'] },
+      expectedUrl: 'ws://localhost/index?id=123&type=test&tag=a&tag=b',
+    },
+  ])('$description', ({ url, expectedUrl, query }) => {
+    const webSocketMock = vi.fn()
+    const client = hc<AppType>(url, {
+      webSocket(url, options) {
+        return webSocketMock(url, options)
+      },
+    })
+    client.index.$ws({ query })
+    expect(webSocketMock).toHaveBeenCalledWith(expectedUrl, undefined)
+  })
+})
diff --git a/src/client/client.ts b/src/client/client.ts
index acd0ddadd..f9b487926 100644
--- a/src/client/client.ts
+++ b/src/client/client.ts
@@ -194,8 +194,14 @@ export const hc = <T extends Hono<any, any, any>>(
           }
         })
       }
+      const establishWebSocket = (...args: ConstructorParameters<typeof WebSocket>) => {
+        if (options?.webSocket !== undefined && typeof options.webSocket === 'function') {
+          return options.webSocket(...args)
+        }
+        return new WebSocket(...args)
+      }
 
-      return new WebSocket(targetUrl.toString())
+      return establishWebSocket(targetUrl.toString())
     }
 
     const req = new ClientRequestImpl(url, method)
diff --git a/src/client/types.ts b/src/client/types.ts
index 8f2c77c50..e64a98eb8 100644
--- a/src/client/types.ts
+++ b/src/client/types.ts
@@ -7,6 +7,7 @@ type HonoRequest = (typeof Hono.prototype)['request']
 
 export type ClientRequestOptions<T = unknown> = {
   fetch?: typeof fetch | HonoRequest
+  webSocket?: (...args: ConstructorParameters<typeof WebSocket>) => WebSocket
   /**
    * Standard `RequestInit`, caution that this take highest priority
    * and could be used to overwrite things that Hono sets for you, like `body | method | headers`.