feat: provide lifecycle hooks for injectWS (#341)

* feat: provide lifecycle hooks for injectWS

* test: add typescript test

Author: climba03003

index.js

@@ -51,7 +51,8 @@ function fastifyWebsocket (fastify, opts, next) {
   const wss = new WebSocket.Server(wssOptions)
   fastify.decorate('websocketServer', wss)
 
-  async function injectWS (path = '/', upgradeContext = {}) {
+  // TODO: place upgrade context as options
+  async function injectWS (path = '/', upgradeContext = {}, options = {}) {
     const server2Client = new PassThrough()
     const client2Server = new PassThrough()
 
@@ -64,7 +65,10 @@ function fastifyWebsocket (fastify, opts, next) {
     let resolve, reject
     const promise = new Promise((_resolve, _reject) => { resolve = _resolve; reject = _reject })
 
+    typeof options.onInit === 'function' && options.onInit(ws)
+
     ws.on('open', () => {
+      typeof options.onOpen === 'function' && options.onOpen(ws)
       clientStream.removeListener('data', onData)
       resolve(ws)
     })

test/inject.test.js

@@ -132,3 +132,43 @@ test('rejects if the websocket is not upgraded', async (t) => {
   await fastify.ready()
   await t.assert.rejects(fastify.injectWS('/'), new Error('Unexpected server response: 401'))
 })
+
+test('inject hooks', async (t) => {
+  const fastify = buildFastify(t)
+  const message = 'hi from client'
+
+  let _resolve
+  const promise = new Promise((resolve) => { _resolve = resolve })
+
+  fastify.register(
+    async function (instance) {
+      instance.get('/ws', { websocket: true }, function (socket) {
+        socket.once('message', chunk => {
+          _resolve(chunk.toString())
+        })
+      })
+    })
+
+  await fastify.ready()
+
+  let order = 0
+  let initWS, openWS
+  const ws = await fastify.injectWS('/ws', {}, {
+    onInit (ws) {
+      t.assert.strictEqual(order, 0)
+      order++
+      initWS = ws
+    },
+    onOpen (ws) {
+      t.assert.strictEqual(order, 1)
+      order++
+      openWS = ws
+    }
+  })
+  ws.send(message)
+  t.assert.strictEqual(order, 2)
+  t.assert.deepStrictEqual(ws, initWS)
+  t.assert.deepStrictEqual(ws, openWS)
+  t.assert.deepStrictEqual(await promise, message)
+  ws.terminate()
+})

types/index.d.ts

@@ -1,11 +1,11 @@
 /// <reference types="node" />
-import { IncomingMessage, ServerResponse, Server } from 'node:http'
-import { FastifyRequest, FastifyPluginCallback, RawServerBase, RawServerDefault, RawRequestDefaultExpression, RawReplyDefaultExpression, RequestGenericInterface, ContextConfigDefault, FastifyInstance, FastifySchema, FastifyTypeProvider, FastifyTypeProviderDefault, FastifyBaseLogger } from 'fastify'
 import * as fastify from 'fastify'
-import * as WebSocket from 'ws'
+import { ContextConfigDefault, FastifyBaseLogger, FastifyInstance, FastifyPluginCallback, FastifyRequest, FastifySchema, FastifyTypeProvider, FastifyTypeProviderDefault, RawReplyDefaultExpression, RawRequestDefaultExpression, RawServerBase, RawServerDefault, RequestGenericInterface } from 'fastify'
+import { preCloseAsyncHookHandler, preCloseHookHandler } from 'fastify/types/hooks'
 import { FastifyReply } from 'fastify/types/reply'
-import { preCloseHookHandler, preCloseAsyncHookHandler } from 'fastify/types/hooks'
 import { RouteGenericInterface } from 'fastify/types/route'
+import { IncomingMessage, Server, ServerResponse } from 'node:http'
+import * as WebSocket from 'ws'
 
 interface WebsocketRouteOptions<
   RawServer extends RawServerBase = RawServerDefault,
@@ -27,8 +27,13 @@ declare module 'fastify' {
     websocket?: boolean;
   }
 
+  interface InjectWSOption {
+    onInit?: (ws: WebSocket.WebSocket) => void
+    onOpen?: (ws: WebSocket.WebSocket) => void
+  }
+
   type InjectWSFn<RawRequest> =
-    ((path?: string, upgradeContext?: Partial<RawRequest>) => Promise<WebSocket>)
+    ((path?: string, upgradeContext?: Partial<RawRequest>, options?: InjectWSOption) => Promise<WebSocket>)
 
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
   interface FastifyInstance<RawServer, RawRequest, RawReply, Logger, TypeProvider> {

types/index.test-d.ts

@@ -1,12 +1,12 @@
-// eslint-disable-next-line import-x/no-named-default -- Testing default export
-import fastifyWebsocket, { WebsocketHandler, fastifyWebsocket as namedFastifyWebsocket, default as defaultFastifyWebsocket, WebSocket } from '..'
+import { TypeBoxTypeProvider } from '@fastify/type-provider-typebox'
+import { Type } from '@sinclair/typebox'
+import fastify, { FastifyBaseLogger, FastifyInstance, FastifyReply, FastifyRequest, FastifySchema, RawRequestDefaultExpression, RawServerDefault, RequestGenericInterface, RouteOptions } from 'fastify'
+import { RouteGenericInterface } from 'fastify/types/route'
 import type { IncomingMessage } from 'node:http'
-import fastify, { RouteOptions, FastifyRequest, FastifyInstance, FastifyReply, RequestGenericInterface, FastifyBaseLogger, RawServerDefault, FastifySchema, RawRequestDefaultExpression } from 'fastify'
 import { expectType } from 'tsd'
 import { Server } from 'ws'
-import { RouteGenericInterface } from 'fastify/types/route'
-import { TypeBoxTypeProvider } from '@fastify/type-provider-typebox'
-import { Type } from '@sinclair/typebox'
+// eslint-disable-next-line import-x/no-named-default -- Test default export
+import fastifyWebsocket, { default as defaultFastifyWebsocket, fastifyWebsocket as namedFastifyWebsocket, WebSocket, WebsocketHandler } from '..'
 
 const app: FastifyInstance = fastify()
 app.register(fastifyWebsocket)
@@ -22,8 +22,8 @@ app.register(fastifyWebsocket, {
   }
 })
 app.register(fastifyWebsocket, { options: { perMessageDeflate: true } })
-app.register(fastifyWebsocket, { preClose: function syncPreclose () {} })
-app.register(fastifyWebsocket, { preClose: async function asyncPreclose () {} })
+app.register(fastifyWebsocket, { preClose: function syncPreclose () { } })
+app.register(fastifyWebsocket, { preClose: async function asyncPreclose () { } })
 
 app.get('/websockets-via-inferrence', { websocket: true }, async function (socket, request) {
   expectType<FastifyInstance>(this)
@@ -89,7 +89,7 @@ app.get<{ Params: { foo: string }, Body: { bar: string }, Querystring: { search:
   expectType<{ foo: string }>(request.params)
   expectType<{ bar: string }>(request.body)
   expectType<{ search: string }>(request.query)
-  expectType< IncomingMessage['headers'] & { auth: string }>(request.headers)
+  expectType<IncomingMessage['headers'] & { auth: string }>(request.headers)
 })
 
 app.route<{ Params: { foo: string }, Body: { bar: string }, Querystring: { search: string }, Headers: { auth: string } }>({
@@ -162,3 +162,15 @@ server.get('/websockets-no-type-inference',
 
 expectType<typeof fastifyWebsocket>(namedFastifyWebsocket)
 expectType<typeof fastifyWebsocket>(defaultFastifyWebsocket)
+
+app.injectWS('/', {}, {})
+app.injectWS('/', {}, {
+  onInit (ws) {
+    expectType<WebSocket>(ws)
+  },
+})
+app.injectWS('/', {}, {
+  onOpen (ws) {
+    expectType<WebSocket>(ws)
+  },
+})