feat: include hooks not within route handler when wrapping routes (#65)

index.js

@@ -110,6 +110,17 @@ async function openTelemetryPlugin (fastify, opts = {}) {
     contextMap.set(request, trace.setSpan(activeContext, span))
   }
 
+  function onRequestWrapRoutes (request, reply, done) {
+    if (
+      !shouldIgnoreRoute(request.url, request.method) &&
+      (wrapRoutes === true || (Array.isArray(wrapRoutes) && wrapRoutes.includes(request.url)))
+    ) {
+      context.with(getContext(request), done)
+    } else {
+      done()
+    }
+  }
+
   async function onResponse (request, reply) {
     if (shouldIgnoreRoute(request.url, request.method)) return
 
@@ -136,28 +147,9 @@ async function openTelemetryPlugin (fastify, opts = {}) {
   }
 
   fastify.addHook('onRequest', onRequest)
+  if (wrapRoutes) fastify.addHook('onRequest', onRequestWrapRoutes)
   fastify.addHook('onResponse', onResponse)
   fastify.addHook('onError', onError)
-
-  if (wrapRoutes) {
-    function wrapRoute (routeHandler) {
-      return function (request, ...args) {
-        const reqContext = getContext(request)
-        return context.with(reqContext, routeHandler.bind(this, request, ...args))
-      }
-    }
-
-    fastify.addHook('onRoute', function (routeOpts) {
-      const { path, handler, method } = routeOpts
-      if (!shouldIgnoreRoute(path, method)) {
-        if (wrapRoutes === true) {
-          routeOpts.handler = wrapRoute(handler)
-        } else if (Array.isArray(wrapRoutes) && wrapRoutes.includes(path)) {
-          routeOpts.handler = wrapRoute(handler)
-        }
-      }
-    })
-  }
 }
 
 module.exports = fp(openTelemetryPlugin, {

test/index.test.js

@@ -238,11 +238,17 @@ test('should wrap all routes when wrapRoutes is true', async ({ equal, same, tea
 
 test('should only wrap routes provided in wrapRoutes array', async ({ same, equal, teardown }) => {
   const dummyContext = trace.setSpan(ROOT_CONTEXT, STUB_SPAN)
+  let next
 
   const fastify = require('fastify')()
 
   await fastify.register(openTelemetryPlugin, { wrapRoutes: ['/testTwo'] })
 
+  fastify.addHook('preHandler', (request, reply, done) => {
+    next = done
+    done()
+  })
+
   const testHandlerOne = async () => 'one'
   const testHandlerTwo = async () => 'two'
 
@@ -257,11 +263,14 @@ test('should only wrap routes provided in wrapRoutes array', async ({ same, equa
   })
 
   await fastify.inject({ ...injectArgs, url: '/testOne' })
+
+  equal(STUB_CONTEXT_API.with.called, false)
+
   await fastify.inject({ ...injectArgs, url: '/testTwo' })
 
   equal(STUB_CONTEXT_API.with.calledOnce, true)
   same(STUB_CONTEXT_API.with.args[0][0], dummyContext)
-  same(await STUB_CONTEXT_API.with.args[0][1](), await testHandlerTwo())
+  same(STUB_CONTEXT_API.with.args[0][1], next)
 })
 
 test('should ignore routes found in ignoreRoutes array', async ({ equal, same, teardown }) => {