Skip to content

Commit

Permalink
test(ws): add interception tests
Browse files Browse the repository at this point in the history
  • Loading branch information
kettanaito committed Feb 9, 2024
1 parent 38f5d27 commit 35fccd7
Show file tree
Hide file tree
Showing 8 changed files with 225 additions and 31 deletions.
5 changes: 4 additions & 1 deletion package.json
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,7 @@
"devDependencies": {
"@commitlint/cli": "^18.4.4",
"@commitlint/config-conventional": "^18.4.4",
"@fastify/websocket": "^8.3.1",
"@open-draft/test-server": "^0.4.2",
"@ossjs/release": "^0.8.0",
"@playwright/test": "^1.40.1",
Expand All @@ -143,6 +144,7 @@
"@types/glob": "^8.1.0",
"@types/json-bigint": "^1.0.4",
"@types/node": "18.x",
"@types/ws": "^8.5.10",
"@typescript-eslint/eslint-plugin": "^5.11.0",
"@typescript-eslint/parser": "^5.11.0",
"@web/dev-server": "^0.1.38",
Expand All @@ -158,6 +160,7 @@
"eslint-config-prettier": "^9.1.0",
"eslint-plugin-prettier": "^5.1.3",
"express": "^4.18.2",
"fastify": "^4.26.0",
"fs-extra": "^11.2.0",
"fs-teardown": "^0.3.0",
"glob": "^10.3.10",
Expand All @@ -174,7 +177,7 @@
"typescript": "^5.0.2",
"undici": "^5.20.0",
"url-loader": "^4.1.1",
"vitest": "^0.34.6",
"vitest": "^1.2.2",
"vitest-environment-miniflare": "^2.14.1",
"webpack": "^5.89.0",
"webpack-http-server": "^0.5.0"
Expand Down
31 changes: 6 additions & 25 deletions src/core/handlers/WebSocketHandler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ type WebSocketHandlerParsedResult = {
match: Match
}

type WebSocketHandlerEventMap = {
export type WebSocketHandlerEventMap = {
connection: [
args: {
client: WebSocketClientConnection
Expand All @@ -29,33 +29,14 @@ type WebSocketHandlerIncomingEvent = MessageEvent<{
server: WebSocketServerConnection
}>

export const kRun = Symbol('run')
export const kEmitter = Symbol('kEmitter')
export const kRun = Symbol('kRun')

export class WebSocketHandler {
public on: <K extends keyof WebSocketHandlerEventMap>(
event: K,
listener: (...args: WebSocketHandlerEventMap[K]) => void,
) => void

public off: <K extends keyof WebSocketHandlerEventMap>(
event: K,
listener: (...args: WebSocketHandlerEventMap[K]) => void,
) => void

public removeAllListeners: <K extends keyof WebSocketHandlerEventMap>(
event?: K,
) => void

protected emitter: Emitter<WebSocketHandlerEventMap>
protected [kEmitter]: Emitter<WebSocketHandlerEventMap>

constructor(private readonly url: Path) {
this.emitter = new Emitter()

// Forward some of the emitter API to the public API
// of the event handler.
this.on = this.emitter.on.bind(this.emitter)
this.off = this.emitter.off.bind(this.emitter)
this.removeAllListeners = this.emitter.removeAllListeners.bind(this.emitter)
this[kEmitter] = new Emitter()
}

public parse(args: {
Expand Down Expand Up @@ -95,7 +76,7 @@ export class WebSocketHandler {

// Emit the connection event on the handler.
// This is what the developer adds listeners for.
this.emitter.emit('connection', {
this[kEmitter].emit('connection', {
client: connection.client,
server: connection.server,
params: parsedResult.match.params || {},
Expand Down
24 changes: 22 additions & 2 deletions src/core/ws/ws.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,8 @@
import { WebSocketHandler } from '../handlers/WebSocketHandler'
import {
WebSocketHandler,
kEmitter,
type WebSocketHandlerEventMap,
} from '../handlers/WebSocketHandler'
import type { Path } from '../utils/matching/matchRequestUrl'
import { webSocketInterceptor } from './webSocketInterceptor'

Expand All @@ -11,7 +15,23 @@ import { webSocketInterceptor } from './webSocketInterceptor'
*/
function createWebSocketLinkHandler(url: Path) {
webSocketInterceptor.apply()
return new WebSocketHandler(url)

return {
on<K extends keyof WebSocketHandlerEventMap>(
event: K,
listener: (...args: WebSocketHandlerEventMap[K]) => void,
): WebSocketHandler {
const handler = new WebSocketHandler(url)

// The "handleWebSocketEvent" function will invoke
// the "run()" method on the WebSocketHandler.
// If the handler matches, it will emit the "connection"
// event. Attach the user-defined listener to that event.
handler[kEmitter].on(event, listener)

return handler
},
}
}

export const ws = {
Expand Down
10 changes: 7 additions & 3 deletions src/node/SetupServerApi.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,14 @@ import {
import { invariant } from 'outvariant'
import { SetupApi } from '~/core/SetupApi'
import { RequestHandler } from '~/core/handlers/RequestHandler'
import { LifeCycleEventsMap, SharedOptions } from '~/core/sharedOptions'
import { RequiredDeep } from '~/core/typeUtils'
import type { LifeCycleEventsMap, SharedOptions } from '~/core/sharedOptions'
import type { RequiredDeep } from '~/core/typeUtils'
import { handleRequest } from '~/core/utils/handleRequest'
import { devUtils } from '~/core/utils/internal/devUtils'
import { mergeRight } from '~/core/utils/internal/mergeRight'
import { SetupServer } from './glossary'
import type { WebSocketHandler } from '~/core/handlers/WebSocketHandler'
import { handleWebSocketEvent } from '~/core/utils/handleWebSocketEvent'
import type { SetupServer } from './glossary'

const DEFAULT_LISTEN_OPTIONS: RequiredDeep<SharedOptions> = {
onUnhandledRequest: 'warn',
Expand Down Expand Up @@ -79,6 +80,9 @@ export class SetupServerApi
)
},
)

// Handle outgoing WebSocket connections.
handleWebSocketEvent(this.currentHandlers)
}

public listen(options: Partial<SharedOptions> = {}): void {
Expand Down
2 changes: 2 additions & 0 deletions test/node/vitest.config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@ export default defineConfig({
dir: './test/node',
globals: true,
alias: {
'vitest-environment-node-websocket':
'./test/support/environments/vitest-environment-node-websocket',
'msw/node': path.resolve(LIB_DIR, 'node/index.mjs'),
'msw/native': path.resolve(LIB_DIR, 'native/index.mjs'),
'msw/browser': path.resolve(LIB_DIR, 'browser/index.mjs'),
Expand Down
109 changes: 109 additions & 0 deletions test/node/ws-api/ws.intercept.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,109 @@
/**
* @vitest-environment node-websocket
*/
import { ws } from 'msw'
import { setupServer } from 'msw/node'
import { WebSocketServer } from '../../support/WebSocketServer'
import { waitFor } from '../../support/waitFor'

const server = setupServer()
const wsServer = new WebSocketServer()

const service = ws.link('ws://*')

beforeAll(async () => {
server.listen()
await wsServer.listen()
})

afterEach(() => {
wsServer.closeAllClients()
wsServer.removeAllListeners()
})

afterAll(async () => {
server.close()
await wsServer.close()
})

it('intercepts outgoing client text message', async () => {
const mockMessageListener = vi.fn()
const realConnectionListener = vi.fn()

server.use(
service.on('connection', ({ client }) => {
client.addEventListener('message', mockMessageListener)
}),
)
wsServer.on('connection', realConnectionListener)

const ws = new WebSocket(wsServer.url)
ws.onopen = () => ws.send('hello')

await waitFor(() => {
// Must intercept the outgoing client message event.
expect(mockMessageListener).toHaveBeenCalledTimes(1)

const messageEvent = mockMessageListener.mock.calls[0][0] as MessageEvent
expect(messageEvent.type).toBe('message')
expect(messageEvent.data).toBe('hello')
expect(messageEvent.target).toBe(ws)

// Must not connect to the actual server by default.
expect(realConnectionListener).not.toHaveBeenCalled()
})
})

it('intercepts outgoing client Blob message', async () => {
const mockMessageListener = vi.fn()
const realConnectionListener = vi.fn()

server.use(
service.on('connection', ({ client }) => {
client.addEventListener('message', mockMessageListener)
}),
)
wsServer.on('connection', realConnectionListener)

const ws = new WebSocket(wsServer.url)
ws.onopen = () => ws.send(new Blob(['hello']))

await waitFor(() => {
expect(mockMessageListener).toHaveBeenCalledTimes(1)

const messageEvent = mockMessageListener.mock.calls[0][0] as MessageEvent
expect(messageEvent.type).toBe('message')
expect(messageEvent.data.size).toBe(5)
expect(messageEvent.target).toEqual(ws)

// Must not connect to the actual server by default.
expect(realConnectionListener).not.toHaveBeenCalled()
})
})

it('intercepts outgoing client ArrayBuffer data', async () => {
const mockMessageListener = vi.fn()
const realConnectionListener = vi.fn()

server.use(
service.on('connection', ({ client }) => {
client.addEventListener('message', mockMessageListener)
}),
)
wsServer.on('connection', realConnectionListener)

const ws = new WebSocket(wsServer.url)
ws.onopen = () => ws.send(new TextEncoder().encode('hello'))

await waitFor(() => {
expect(mockMessageListener).toHaveBeenCalledTimes(1)

const messageEvent = mockMessageListener.mock.calls[0][0] as MessageEvent
expect(messageEvent.type).toBe('message')
expect(messageEvent.data).toEqual(new TextEncoder().encode('hello'))
expect(messageEvent.target).toEqual(ws)

// Must not connect to the actual server by default.
expect(realConnectionListener).not.toHaveBeenCalled()
})
})
55 changes: 55 additions & 0 deletions test/support/WebSocketServer.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
import { invariant } from 'outvariant'
import { Emitter } from 'strict-event-emitter'
import fastify, { FastifyInstance } from 'fastify'
import fastifyWebSocket, { SocketStream } from '@fastify/websocket'

type FastifySocket = SocketStream['socket']

type WebSocketEventMap = {
connection: [client: FastifySocket]
}

export class WebSocketServer extends Emitter<WebSocketEventMap> {
private _url?: string
private app: FastifyInstance
private clients: Set<FastifySocket>

constructor() {
super()
this.clients = new Set()

this.app = fastify()
this.app.register(fastifyWebSocket)
this.app.register(async (fastify) => {
fastify.get('/', { websocket: true }, (connection) => {
this.clients.add(connection.socket)
this.emit('connection', connection.socket)
})
})
}

get url(): string {
invariant(
this._url,
'Failed to get "url" on WebSocketServer: server is not running. Did you forget to "await server.listen()"?',
)
return this._url
}

public async listen(): Promise<void> {
const address = await this.app.listen({ port: 0 })
const url = new URL(address)
url.protocol = url.protocol.replace(/^http/, 'ws')
this._url = url.href
}

public closeAllClients(): void {
this.clients.forEach((client) => {
client.close()
})
}

public async close(): Promise<void> {
return this.app.close()
}
}
20 changes: 20 additions & 0 deletions test/support/environments/vitest-environment-node-websocket.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
/**
* Node.js environment superset that has a global WebSocket API.
*/
import type { Environment } from 'vitest'
import { builtinEnvironments } from 'vitest/environments'
import { WebSocket } from 'undici'

export default <Environment>{
name: 'node-with-websocket',
transformMode: 'ssr',
async setup(global, options) {
const { teardown } = await builtinEnvironments.jsdom.setup(global, options)

Reflect.set(globalThis, 'WebSocket', WebSocket)

return {
teardown,
}
},
}

0 comments on commit 35fccd7

Please sign in to comment.