From ad0cddaf117ca6a962d57236bc2c7474952de1ff Mon Sep 17 00:00:00 2001 From: Artem Zakharchenko Date: Mon, 30 Sep 2024 11:40:14 +0200 Subject: [PATCH] feat(ws): add "socket" property on `server` object (#647) --- .../WebSocket/WebSocketServerConnection.ts | 39 ++++++------ .../compliance/websocket.server.close.test.ts | 13 ++-- .../websocket.server.connect.test.ts | 6 +- .../websocket.server.socket.test.ts | 62 +++++++++++++++++++ .../exchange/websocket.server.connect.test.ts | 12 ++-- 5 files changed, 93 insertions(+), 39 deletions(-) create mode 100644 test/modules/WebSocket/compliance/websocket.server.socket.test.ts diff --git a/src/interceptors/WebSocket/WebSocketServerConnection.ts b/src/interceptors/WebSocket/WebSocketServerConnection.ts index 378264d8..6828ab57 100644 --- a/src/interceptors/WebSocket/WebSocketServerConnection.ts +++ b/src/interceptors/WebSocket/WebSocketServerConnection.ts @@ -38,7 +38,7 @@ export class WebSocketServerConnection { private [kEmitter]: EventTarget constructor( - private readonly socket: WebSocketOverride, + private readonly client: WebSocketOverride, private readonly transport: WebSocketClassTransport, private readonly createConnection: () => WebSocket ) { @@ -53,7 +53,7 @@ export class WebSocketServerConnection { this.transport.addEventListener('outgoing', (event) => { // Ignore client messages if the server connection // hasn't been established yet. Nowhere to forward. - if (this.readyState === -1) { + if (typeof this.realWebSocket === 'undefined') { return } @@ -75,17 +75,16 @@ export class WebSocketServerConnection { } /** - * Server ready state. - * Proxies the ready state of the original WebSocket instance, - * if set. If the original connection hasn't been established, - * defaults to `-1`. + * The `WebSocket` instance connected to the original server. + * Accessing this before calling `server.connect()` will throw. */ - public get readyState(): number { - if (this.realWebSocket) { - return this.realWebSocket.readyState - } + public get socket(): WebSocket { + invariant( + this.realWebSocket, + 'Cannot access "socket" on the original WebSocket server object: the connection is not open. Did you forget to call `server.connect()`?' + ) - return -1 + return this.realWebSocket } /** @@ -100,7 +99,7 @@ export class WebSocketServerConnection { const realWebSocket = this.createConnection() // Inherit the binary type from the mock WebSocket client. - realWebSocket.binaryType = this.socket.binaryType + realWebSocket.binaryType = this.client.binaryType // Allow the interceptor to listen to when the server connection // has been established. This isn't necessary to operate with the connection @@ -133,7 +132,7 @@ export class WebSocketServerConnection { // Close the original connection when the mock client closes. // E.g. "client.close()" was called. This is never forwarded anywhere. - this.socket.addEventListener('close', this.handleMockClose.bind(this), { + this.client.addEventListener('close', this.handleMockClose.bind(this), { signal: this.mockCloseController.signal, }) @@ -150,7 +149,7 @@ export class WebSocketServerConnection { // Forward original server errors to the WebSocket client. // This ensures the client is closed if the original server errors. - this.socket.dispatchEvent(bindEvent(this.socket, new Event('error'))) + this.client.dispatchEvent(bindEvent(this.client, new Event('error'))) }) this.realWebSocket = realWebSocket @@ -164,7 +163,7 @@ export class WebSocketServerConnection { listener: WebSocketEventListener, options?: AddEventListenerOptions | boolean ): void { - const boundListener = listener.bind(this.socket) + const boundListener = listener.bind(this.client) // Store the bound listener on the original listener // so the exact bound function can be accessed in "removeEventListener()". @@ -208,7 +207,7 @@ export class WebSocketServerConnection { invariant( realWebSocket, 'Failed to call "server.send()" for "%s": the connection is not open. Did you forget to call "server.connect()"?', - this.socket.url + this.client.url ) // Silently ignore writes on the closed original WebSocket. @@ -246,7 +245,7 @@ export class WebSocketServerConnection { invariant( realWebSocket, 'Failed to close server connection for "%s": the connection is not open. Did you forget to call "server.connect()"?', - this.socket.url + this.client.url ) // Remove the "close" event listener from the server @@ -310,14 +309,14 @@ export class WebSocketServerConnection { * Preventing the default on the message event stops this. */ if (!messageEvent.defaultPrevented) { - this.socket.dispatchEvent( + this.client.dispatchEvent( bindEvent( /** * @note Bind the forwarded original server events * to the mock WebSocket instance so it would * dispatch them straight away. */ - this.socket, + this.client, // Clone the message event again to prevent // the "already being dispatched" exception. new MessageEvent('message', { @@ -361,7 +360,7 @@ export class WebSocketServerConnection { // allow non-configurable status codes from the server. // If the socket has been closed by now, no harm calling // this again—it will have no effect. - this.socket[kClose](event.code, event.reason) + this.client[kClose](event.code, event.reason) } } } diff --git a/test/modules/WebSocket/compliance/websocket.server.close.test.ts b/test/modules/WebSocket/compliance/websocket.server.close.test.ts index bc3dc841..7308733b 100644 --- a/test/modules/WebSocket/compliance/websocket.server.close.test.ts +++ b/test/modules/WebSocket/compliance/websocket.server.close.test.ts @@ -1,6 +1,4 @@ -/** - * @vitest-environment node-with-websocket - */ +// @vitest-environment node-with-websocket import { DeferredPromise } from '@open-draft/deferred-promise' import { vi, it, expect, beforeAll, afterEach, afterAll } from 'vitest' import { WebSocketServer, Data } from 'ws' @@ -59,13 +57,10 @@ it('closes the actual server connection when called "server.close()"', async () interceptor.once('connection', ({ client, server }) => { server.connect() - serverCallback(server.readyState) + serverCallback(server.socket.readyState) - /** - * @fixme Tapping into internals isn't nice. - */ - server['realWebSocket']?.addEventListener('close', () => { - serverCallback(server.readyState) + server.socket.addEventListener('close', () => { + serverCallback(server.socket.readyState) }) client.addEventListener('message', (event) => { diff --git a/test/modules/WebSocket/compliance/websocket.server.connect.test.ts b/test/modules/WebSocket/compliance/websocket.server.connect.test.ts index c21aeedc..02674c44 100644 --- a/test/modules/WebSocket/compliance/websocket.server.connect.test.ts +++ b/test/modules/WebSocket/compliance/websocket.server.connect.test.ts @@ -1,6 +1,4 @@ -/** - * @vitest-environment node-with-websocket - */ +// @vitest-environment node-with-websocket import { vi, it, expect, beforeAll, afterEach, afterAll } from 'vitest' import { Data, WebSocketServer } from 'ws' import { WebSocketInterceptor } from '../../../../src/interceptors/WebSocket' @@ -40,7 +38,7 @@ it('forwards client messages to the server by default', async () => { }) }) - interceptor.once('connection', ({ client, server }) => { + interceptor.once('connection', ({ server }) => { server.connect() }) diff --git a/test/modules/WebSocket/compliance/websocket.server.socket.test.ts b/test/modules/WebSocket/compliance/websocket.server.socket.test.ts new file mode 100644 index 00000000..6425adf6 --- /dev/null +++ b/test/modules/WebSocket/compliance/websocket.server.socket.test.ts @@ -0,0 +1,62 @@ +// @vitest-environment node-with-websocket +import { it, expect, beforeAll, afterEach, afterAll } from 'vitest' +import { WebSocketInterceptor } from '../../../../src/interceptors/WebSocket' +import { waitForWebSocketEvent } from '../utils/waitForWebSocketEvent' +import { DeferredPromise } from '@open-draft/deferred-promise' + +const interceptor = new WebSocketInterceptor() + +beforeAll(() => { + interceptor.apply() +}) + +afterEach(() => { + interceptor.removeAllListeners() +}) + +afterAll(() => { + interceptor.dispose() +}) + +it('throws an error when accessing "server.socket" before calling "server.connect()"', async () => { + const socketPromise = new DeferredPromise() + interceptor.on('connection', ({ server }) => { + try { + // Accessing "server.socket" before calling "server.connect()" is a no-op. + const socket = server.socket + socketPromise.resolve(socket) + } catch (error) { + socketPromise.reject(error) + } + }) + + const clientSocket = new WebSocket('wss://localhost') + await waitForWebSocketEvent('open', clientSocket) + + await expect(socketPromise).rejects.toThrow( + 'Cannot access "socket" on the original WebSocket server object: the connection is not open. Did you forget to call `server.connect()`?' + ) + + // Client connection must remain open. + expect(clientSocket.readyState).toBe(WebSocket.OPEN) +}) + +it('returns the WebSocket instance after calling "server.connect()"', async () => { + const socketPromise = new DeferredPromise() + interceptor.on('connection', ({ server }) => { + server.connect() + try { + const socket = server.socket + socketPromise.resolve(socket) + } catch (error) { + socketPromise.reject(error) + } + }) + + await waitForWebSocketEvent('open', new WebSocket('wss://localhost')) + + const serverSocket = await socketPromise + expect(serverSocket).toBeInstanceOf(WebSocket) + expect(serverSocket.url).toBe('wss://localhost/') + expect(serverSocket.readyState).toBe(WebSocket.CONNECTING) +}) diff --git a/test/modules/WebSocket/exchange/websocket.server.connect.test.ts b/test/modules/WebSocket/exchange/websocket.server.connect.test.ts index 09d24352..02678990 100644 --- a/test/modules/WebSocket/exchange/websocket.server.connect.test.ts +++ b/test/modules/WebSocket/exchange/websocket.server.connect.test.ts @@ -1,6 +1,4 @@ -/** - * @vitest-environment node-with-websocket - */ +// @vitest-environment node-with-websocket import { vi, it, expect, beforeAll, afterEach, afterAll } from 'vitest' import { WebSocketServer } from 'ws' import { DeferredPromise } from '@open-draft/deferred-promise' @@ -80,11 +78,11 @@ it('forwards outgoing client data to the original server', async () => { it('closes the actual server connection when the client closes', async () => { const clientClosePromise = new DeferredPromise() - let realWebSocket: WebSocket | undefined + const serverSocketPromise = new DeferredPromise() interceptor.once('connection', ({ client, server }) => { server.connect() - realWebSocket = server['realWebSocket'] + serverSocketPromise.resolve(server.socket) client.addEventListener('message', (event) => { if (event.data === 'close') { @@ -101,8 +99,10 @@ it('closes the actual server connection when the client closes', async () => { ws.addEventListener('close', (event) => clientClosePromise.resolve(event)) await clientClosePromise + const serverSocket = await serverSocketPromise + expect(ws.readyState).toBe(WebSocket.CLOSED) - expect(realWebSocket?.readyState).toBe(WebSocket.CLOSING) + expect(serverSocket.readyState).toBe(WebSocket.CLOSING) }) it('throw an error when connecting to a non-existing server', async () => {