Skip to content

Commit

Permalink
fix(WebSocket): allow reusing the same event listener (#653)
Browse files Browse the repository at this point in the history
  • Loading branch information
kettanaito authored Oct 3, 2024
1 parent 4fab617 commit 3c1cc28
Show file tree
Hide file tree
Showing 3 changed files with 99 additions and 18 deletions.
21 changes: 12 additions & 9 deletions src/interceptors/WebSocket/WebSocketClientConnection.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
}
Expand Down
20 changes: 11 additions & 9 deletions src/interceptors/WebSocket/WebSocketServerConnection.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
}
Expand Down
Original file line number Diff line number Diff line change
@@ -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)
})

0 comments on commit 3c1cc28

Please sign in to comment.