Add #injectWS

README.md

@@ -261,6 +261,8 @@ fastify.register(require('@fastify/websocket'), {
 Testing the ws handler can be quite tricky, luckily `fastify-websocket` decorates fastify instance with `injectWS`.
 It allows to test easily a websocket endpoint.
 
+The signature of injectWS is the following: `([path], [upgradeContext])`.
+
 #### App.js
 
 ```js
@@ -272,7 +274,14 @@ const FastifyWebSocket = require('@fastify/websocket')
 const App = Fastify()
 
 App.register(FastifyWebSocket);
+
 App.register(async function(fastify) {
+  fastify.addHook('preValidation', async (request, reply) => {
+    if (request.headers['api-key'] !== 'some-random-key') {
+      return reply.code(401).send()
+    }
+  })
+
   fastify.get('/', { websocket: true }, (connection) => {
     connection.socket.on('message', message => {
       connection.socket.send('hi from server')
@@ -299,7 +308,7 @@ test('connect to /', async (t) => {
   fastify.register(App)
   t.teardown(fastify.close.bind(fastify))
 
-  const ws = await fastify.injectWS('/')
+  const ws = await fastify.injectWS('/', {headers: { "api-key" : "some-random-key" }})
   let resolve;
   const promise = new Promise(r => { resolve = r })
 
@@ -318,7 +327,7 @@ test('connect to /', async (t) => {
 - Websocket need to be closed manually at the end of each test.
 - `fastify.ready()` needs to be awaited to ensure that fastify has been decorated.
 - You need to register the event listener before sending the message if you need to process server response.
-
+- It's possible to use 
 ## Options
 
 `@fastify/websocket` accept these options for [`ws`](https://github.com/websockets/ws/blob/master/doc/ws.md#new-websocketserveroptions-callback) :

index.js

@@ -50,7 +50,7 @@ function fastifyWebsocket (fastify, opts, next) {
   const wss = new WebSocket.Server(wssOptions)
   fastify.decorate('websocketServer', wss)
 
-  async function injectWS (path = '/') {
+  async function injectWS (path = '/', upgradeContext = {}) {
     const server2Client = new PassThrough()
     const client2Server = new PassThrough()
 
@@ -80,8 +80,10 @@ function fastifyWebsocket (fastify, opts, next) {
     clientStream.on('data', onData)
 
     const req = {
+      ...upgradeContext,
       method: 'GET',
       headers: {
+        ...upgradeContext.headers,
         connection: 'upgrade',
         upgrade: 'websocket',
         'sec-websocket-version': 13,

test/inject.test.js

@@ -86,3 +86,32 @@ test('routes correctly the message between two routes', async (t) => {
   t.same(await promise, message)
   ws.terminate()
 })
+
+test('use the upgrade context to upgrade if there is some hook', 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.addHook('preValidation', async (request, reply) => {
+        if (request.headers['api-key'] !== 'some-random-key') {
+          return reply.code(401).send()
+        }
+      })
+
+      instance.get('/', { websocket: true }, function (conn) {
+        conn.once('data', chunk => {
+          _resolve(chunk.toString())
+        })
+      })
+    })
+
+  await fastify.ready()
+  const ws = await fastify.injectWS('/', { headers: { 'api-key': 'some-random-key' } })
+  ws.send(message)
+  t.same(await promise, message)
+  ws.terminate()
+})

types/index.d.ts

@@ -27,10 +27,13 @@ declare module 'fastify' {
     websocket?: boolean;
   }
 
+  type InjectWSFn<RawRequest> =
+    ((path?: string, upgradeContext?: Partial<RawRequest>) => Promise<WebSocket>)
+
   interface FastifyInstance<RawServer, RawRequest, RawReply, Logger, TypeProvider> {
     get: RouteShorthandMethod<RawServer, RawRequest, RawReply, TypeProvider>,
     websocketServer: WebSocket.Server,
-    injectWS: (path?: string) => Promise<WebSocket>
+    injectWS: InjectWSFn<RawRequest>
   }
 
   interface FastifyRequest {

types/index.test-d.ts

@@ -1,7 +1,7 @@
 import fastifyWebsocket, { WebsocketHandler, SocketStream, fastifyWebsocket as namedFastifyWebsocket, default as defaultFastifyWebsocket } from '..';
 import type { IncomingMessage } from "http";
 import fastify, { RouteOptions, FastifyRequest, FastifyInstance, FastifyReply, RequestGenericInterface, FastifyBaseLogger, RawServerDefault, FastifySchema, RawRequestDefaultExpression, RawServerBase, ContextConfigDefault, RawReplyDefaultExpression } from 'fastify';
-import { expectAssignable, expectType } from 'tsd';
+import { expectAssignable, expectError, expectType } from 'tsd';
 import { Server } from 'ws';
 import { RouteGenericInterface } from 'fastify/types/route';
 import { TypeBoxTypeProvider } from '@fastify/type-provider-typebox';
@@ -22,10 +22,15 @@ 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 (connection, request) {
+app.injectWS()
+app.injectWS('/test')
+app.injectWS('/test', { headers: { 'test': 'test' } })
+expectError(app.injectWS({ headers: { 'test': "test" } }))
+
+app.get('/websockets-via-inferrence', { websocket: true }, async function(connection, request) {
   expectType<FastifyInstance>(this);
   expectType<SocketStream>(connection);
   expectType<Server>(app.websocketServer);
@@ -89,7 +94,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);
 });
 
 
@@ -100,7 +105,7 @@ app.route<{ Params: { foo: string }, Body: { bar: string }, Querystring: { searc
     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);
   },
   wsHandler: (connection, request) => {
     expectType<SocketStream>(connection);
@@ -187,7 +192,7 @@ server.route({
 
 server.get('/websockets-no-type-inference',
   { websocket: true },
-  async function (connection, request) {
+  async function(connection, request) {
     expectType<FastifyInstance>(this);
     expectType<SocketStream>(connection);
     expectType<Server>(app.websocketServer);
@@ -199,6 +204,6 @@ server.get('/websockets-no-type-inference',
     expectType<IncomingMessage['headers']>(request.headers);
   });
 
-  expectType<typeof fastifyWebsocket>(namedFastifyWebsocket);
-  expectType<typeof fastifyWebsocket>(defaultFastifyWebsocket);
-  
+expectType<typeof fastifyWebsocket>(namedFastifyWebsocket);
+expectType<typeof fastifyWebsocket>(defaultFastifyWebsocket);
+