Repository: mercurius-js/mercurius (js)

Issue:
Uniquely identify batched queries through all their lifecycle
While trying to give support to https://github.com/nearform/mercurius-apollo-tracing/issues/125 we noticed that we were unable to uniquely identify each operation definition contained in batched queries through their whole lifecycle. This is needed in order to know which tracer to start or stop at any given time.

We tried two approaches and faced two different issues with them:

1. Using the `document` object contained in the hooks to extract the relevant operation information. This simply can't be done since the `onResolution` hook doesn't receive the `document` object, only the `execution` and `context` ones. See https://github.com/mercurius-js/mercurius/blob/master/docs/hooks.md#onresolution

Also, in the specific case of mercurius-apollo-tracing, it hooks into schema resolvers that don't provide all the `document` details (and they don't need to), just some `info` object. So it's best to go with the next approach...

2. Using the `context.__currentQuery` attribute. Turns out that this attribute is only updated during the `preParsing` stage and left untouched afterwards, so it's basically useless for batched queries during the rest of the lifecycle.

For reproducing this we can setup a basic Fastify + Mercurius server with `allowBatchedQueries: true` and a typical `add` query, and include these debugs on all hooks:
```
app.graphql.addHook('preParsing', async (_, __, context) => {
  app.log.info(`preParsing: ${context.__currentQuery}`)
})

app.graphql.addHook('preValidation', async (_, __, context) => {
  app.log.info(`preValidation: ${context.__currentQuery}`)
})

app.graphql.addHook('preExecution', async (_, __, context) => {
  app.log.info(`preExecution: ${context.__currentQuery}`)
})

app.graphql.addHook('onResolution', async (__, context) => {
  app.log.info(`onResolution: ${context.__currentQuery}`)
})
```

Then fire the following queries from a client that supports batched queries:
```
query { add(x: 1, y: 1) }
query { add(x: 2, y: 2) }
```

One would expect to have each hook print out the two queries, but as we can see only the `preParsing` hook does that while the rest of them print out the second query twice:
```
{"level":30,"time":1660745778776,"pid":74281,"hostname":"BBaidal-SPMAC21","reqId":"req-2","req":{"method":"POST","url":"/graphql","hostname":"localhost:3000","remoteAddress":"::1","remotePort":50784},"msg":"incoming request"}
{"level":30,"time":1660745778780,"pid":74281,"hostname":"BBaidal-SPMAC21","msg":"preParsing: {\n  add(x: 1, y: 1)\n}"}
{"level":30,"time":1660745778781,"pid":74281,"hostname":"BBaidal-SPMAC21","msg":"preParsing: {\n  add(x: 2, y: 2)\n}"}
{"level":30,"time":1660745778781,"pid":74281,"hostname":"BBaidal-SPMAC21","msg":"preValidation: {\n  add(x: 2, y: 2)\n}"}
{"level":30,"time":1660745778782,"pid":74281,"hostname":"BBaidal-SPMAC21","msg":"preValidation: {\n  add(x: 2, y: 2)\n}"}
{"level":30,"time":1660745778786,"pid":74281,"hostname":"BBaidal-SPMAC21","msg":"preExecution: {\n  add(x: 2, y: 2)\n}"}
{"level":30,"time":1660745778787,"pid":74281,"hostname":"BBaidal-SPMAC21","msg":"preExecution: {\n  add(x: 2, y: 2)\n}"}
{"level":30,"time":1660745778788,"pid":74281,"hostname":"BBaidal-SPMAC21","msg":"onResolution: {\n  add(x: 2, y: 2)\n}"}
{"level":30,"time":1660745778788,"pid":74281,"hostname":"BBaidal-SPMAC21","msg":"onResolution: {\n  add(x: 2, y: 2)\n}"}
{"level":30,"time":1660745778789,"pid":74281,"hostname":"BBaidal-SPMAC21","reqId":"req-2","res":{"statusCode":200},"responseTime":12.858789920806885,"msg":"request completed"}
```

I'm sure that not only mercurius-apollo-tracing would benefit from having a way to uniquely identify batched queries through all their lifecycle, but also other potential plugins that could handle batched queries.

Pull request description:
feat: uniquely identify batched queries
resolves #842

Fixes `context.__currentQuery` being overwritten by subsequent batched query operations.
Stores `context.operationsCount` and `context.operationId` (same naming as with subscriptions) to uniquely identify batched query operations at a later stage.

Patch:
diff --git a/index.d.ts b/index.d.ts
index c2ac03a..434a970 100644
--- a/index.d.ts
+++ b/index.d.ts
@@ -25,6 +25,9 @@ export interface PubSub {
 export interface MercuriusContext {
   app: FastifyInstance;
   reply: FastifyReply;
+  operationsCount?: number;
+  operationId?: number;
+  __currentQuery: string;
   /**
    * __Caution__: Only available if `subscriptions` are enabled
    */
@@ -255,12 +258,12 @@ export interface MercuriusPlugin {
    * Define additional resolvers
    * @param resolvers object with resolver functions
    */
-  defineResolvers<TContext = MercuriusContext>(resolvers: IResolvers<any, TContext>): void;
+  defineResolvers<TContext extends Record<string, any> = MercuriusContext>(resolvers: IResolvers<any, TContext>): void;
   /**
    * Define data loaders
    * @param loaders object with data loader functions
    */
-  defineLoaders<TContext = MercuriusContext>(loaders: MercuriusLoaders<TContext>): void;
+  defineLoaders<TContext extends Record<string, any> = MercuriusContext>(loaders: MercuriusLoaders<TContext>): void;
   /**
    * Transform the existing schema
    */
diff --git a/lib/routes.js b/lib/routes.js
index 56a5483..0c40ab0 100644
--- a/lib/routes.js
+++ b/lib/routes.js
@@ -352,13 +352,23 @@ module.exports = async function (app, opts) {
 
     if (allowBatchedQueries && Array.isArray(request.body)) {
       // Batched query
-      return Promise.all(request.body.map(r =>
-        execute(r, request, reply)
-          .catch(e => {
-            const { response } = errorFormatter(e, request[kRequestContext])
-            return response
-          })
-      ))
+      const operationsCount = request.body.length
+
+      Object.assign(request[kRequestContext], { operationsCount })
+
+      return Promise.all(request.body.map(async (r, operationId) => {
+        // Create individual reqs for multiple operations, otherwise reference the original req
+        const operationReq = operationsCount > 1 ? Object.create(request) : request
+
+        Object.assign(operationReq[kRequestContext], { operationId })
+
+        try {
+          return await execute(r, operationReq, reply)
+        } catch (e) {
+          const { response } = errorFormatter(e, request[kRequestContext])
+          return response
+        }
+      }))
     } else {
       // Regular query
       return execute(request.body, request, reply)
diff --git a/lib/subscription-client.js b/lib/subscription-client.js
index 294046c..bc14bbb 100644
--- a/lib/subscription-client.js
+++ b/lib/subscription-client.js
@@ -220,7 +220,11 @@ class SubscriptionClient {
         if (operation) {
           operation.handler(null)
           this.operations.delete(operationId)
-          this.sendMessage(operationId, this.protocolMessageTypes.GQL_ERROR, data.payload)
+          this.sendMessage(
+            operationId,
+            this.protocolMessageTypes.GQL_ERROR,
+            data.payload
+          )
         }
         break
       case this.protocolMessageTypes.GQL_COMPLETE:
@@ -238,11 +242,27 @@ class SubscriptionClient {
         }
         break
       case this.protocolMessageTypes.GQL_CONNECTION_KEEP_ALIVE:
+        if (this.socket) {
+          this.sendMessage(
+            operationId,
+            this.protocolMessageTypes.GQL_CONNECTION_KEEP_ALIVE_ACK
+          )
+        }
         break
-      /* istanbul ignore next */
       default:
+        // GQL_CONNECTION_KEEP_ALIVE_ACK is only defined in the graphql-ws protocol
+        /* istanbul ignore next */
+        if (
+          this.protocolMessageTypes.GQL_CONNECTION_KEEP_ALIVE_ACK &&
+          data.type === this.protocolMessageTypes.GQL_CONNECTION_KEEP_ALIVE_ACK
+        ) {
+          break
+        }
+
         /* istanbul ignore next */
-        throw new MER_ERR_GQL_SUBSCRIPTION_MESSAGE_INVALID(`Invalid message type "${data.type}"`)
+        throw new MER_ERR_GQL_SUBSCRIPTION_MESSAGE_INVALID(
+          `Invalid message type "${data.type}"`
+        )
     }
   }
 
diff --git a/lib/subscription-connection.js b/lib/subscription-connection.js
index f08b290..51ed9dc 100644
--- a/lib/subscription-connection.js
+++ b/lib/subscription-connection.js
@@ -74,11 +74,19 @@ module.exports = class SubscriptionConnection {
         break
       case this.protocolMessageTypes.GQL_START: {
         if (this.isReady) {
-          this.handleGQLStart(data).catch(e => {
-            this.sendMessage(this.protocolMessageTypes.GQL_ERROR, id, e.message)
+          this.handleGQLStart(data).catch((e) => {
+            this.sendMessage(
+              this.protocolMessageTypes.GQL_ERROR,
+              id,
+              e.message
+            )
           })
         } else {
-          this.sendMessage(this.protocolMessageTypes.GQL_CONNECTION_ERROR, undefined, { message: 'Connection has not been established yet.' })
+          this.sendMessage(
+            this.protocolMessageTypes.GQL_CONNECTION_ERROR,
+            undefined,
+            { message: 'Connection has not been established yet.' }
+          )
           return this.handleConnectionClose()
         }
         break
@@ -86,8 +94,30 @@ module.exports = class SubscriptionConnection {
       case this.protocolMessageTypes.GQL_STOP:
         await this.handleGQLStop(data)
         break
+      case this.protocolMessageTypes.GQL_CONNECTION_KEEP_ALIVE:
+        // GQL_CONNECTION_KEEP_ALIVE_ACK is only defined in the graphql-ws protocol
+        /* istanbul ignore next */
+        if (this.protocolMessageTypes.GQL_CONNECTION_KEEP_ALIVE_ACK) {
+          this.sendMessage(
+            this.protocolMessageTypes.GQL_CONNECTION_KEEP_ALIVE_ACK,
+            id
+          )
+        }
+        break
       default:
-        this.sendMessage(this.protocolMessageTypes.GQL_ERROR, id, 'Invalid payload type')
+        // GQL_CONNECTION_KEEP_ALIVE_ACK is only defined in the graphql-ws protocol
+        if (
+          this.protocolMessageTypes.GQL_CONNECTION_KEEP_ALIVE_ACK &&
+          type === this.protocolMessageTypes.GQL_CONNECTION_KEEP_ALIVE_ACK
+        ) {
+          break
+        }
+
+        this.sendMessage(
+          this.protocolMessageTypes.GQL_ERROR,
+          id,
+          'Invalid payload type'
+        )
     }
   }
 
diff --git a/lib/subscription-protocol.js b/lib/subscription-protocol.js
index 79654b9..94a5878 100644
--- a/lib/subscription-protocol.js
+++ b/lib/subscription-protocol.js
@@ -13,7 +13,8 @@ module.exports.getProtocolByName = function (name) {
         GQL_CONNECTION_INIT: 'connection_init', // Client -> Server
         GQL_CONNECTION_ACK: 'connection_ack', // Server -> Client
         GQL_CONNECTION_ERROR: 'connection_error', // Server -> Client
-        GQL_CONNECTION_KEEP_ALIVE: 'ka', // Server -> Client
+        GQL_CONNECTION_KEEP_ALIVE: 'ping', // Bidirectional
+        GQL_CONNECTION_KEEP_ALIVE_ACK: 'pong', // Bidirectional
         GQL_CONNECTION_TERMINATE: 'connection_terminate', // Client -> Server
         GQL_START: 'subscribe', // Client -> Server
         GQL_DATA: 'next', // Server -> Client
diff --git a/package.json b/package.json
index dbb6997..5773bdc 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "mercurius",
-  "version": "10.2.0",
+  "version": "10.4.0",
   "description": "Fastify GraphQL adapter with gateway and subscription support",
   "main": "index.js",
   "types": "index.d.ts",
@@ -32,7 +32,7 @@
     "@graphql-tools/merge": "^8.0.0",
     "@graphql-tools/schema": "^8.0.0",
     "@graphql-tools/utils": "^8.0.0",
-    "@sinonjs/fake-timers": "^9.0.0",
+    "@sinonjs/fake-timers": "^9.1.2",
     "@types/node": "^18.0.0",
     "@types/ws": "^8.2.0",
     "@types/isomorphic-form-data": "^2.0.0",
@@ -48,8 +48,8 @@
     "snazzy": "^9.0.0",
     "split2": "^4.0.0",
     "standard": "^17.0.0",
-    "tap": "^16.0.0",
-    "tsd": "^0.22.0",
+    "tap": "^16.3.0",
+    "tsd": "^0.23.0",
     "typescript": "^4.3.5",
     "wait-on": "^6.0.0"
   },
diff --git a/static/main.js b/static/main.js
index dee49eb..485d39e 100644
--- a/static/main.js
+++ b/static/main.js
@@ -42,7 +42,7 @@ function render () {
 
 function importDependencies () {
   const link = document.createElement('link')
-  link.href = 'https://unpkg.com/graphiql@1.4.2/graphiql.css'
+  link.href = 'https://unpkg.com/graphiql@2.0.2/graphiql.min.css'
   link.type = 'text/css'
   link.rel = 'stylesheet'
   link.media = 'screen,print'
@@ -50,10 +50,9 @@ function importDependencies () {
   document.getElementsByTagName('head')[0].appendChild(link)
 
   return importer.urls([
-    'https://unpkg.com/react@16.8.0/umd/react.production.min.js',
-    'https://unpkg.com/react-dom@16.8.0/umd/react-dom.production.min.js',
-    'https://unpkg.com/graphiql@1.4.2/graphiql.min.js',
-    'https://unpkg.com/subscriptions-transport-ws@0.9.19/browser/client.js'
+    'https://unpkg.com/react@18.2.0/umd/react.production.min.js',
+    'https://unpkg.com/react-dom@18.2.0/umd/react-dom.production.min.js',
+    'https://unpkg.com/graphiql@2.0.2/graphiql.min.js'
   ])
 }
 
diff --git a/static/sw.js b/static/sw.js
index db99ee9..2c9f3ca 100644
--- a/static/sw.js
+++ b/static/sw.js
@@ -2,14 +2,13 @@
 
 self.addEventListener('install', function (e) {
   e.waitUntil(
-    caches.open('graphiql-v1.4.0').then(function (cache) {
+    caches.open('graphiql-v2.0.2').then(function (cache) {
       return cache.addAll([
         './main.js',
-        'https://unpkg.com/graphiql@1.4.2/graphiql.css',
-        'https://unpkg.com/react@16.8.0/umd/react.production.min.js',
-        'https://unpkg.com/react-dom@16.8.0/umd/react-dom.production.min.js',
-        'https://unpkg.com/graphiql@1.4.2/graphiql.min.js',
-        'https://unpkg.com/subscriptions-transport-ws@0.9.19/browser/client.js'
+        'https://unpkg.com/graphiql@2.0.2/graphiql.css',
+        'https://unpkg.com/react@18.2.0/umd/react.production.min.js',
+        'https://unpkg.com/react-dom@18.2.0/umd/react-dom.production.min.js',
+        'https://unpkg.com/graphiql@2.0.2/graphiql.min.js'
       ])
     })
   )


Test patch:
diff --git a/test/batched.js b/test/batched.js
index 949ff48..cdf98c1 100644
--- a/test/batched.js
+++ b/test/batched.js
@@ -2,6 +2,7 @@
 
 const { test } = require('tap')
 const Fastify = require('fastify')
+const sinon = require('sinon')
 const GQL = require('..')
 
 test('POST regular query', async (t) => {
@@ -346,3 +347,44 @@ test('POST batched query with a resolver which succeeds and a resolver which thr
 
   t.same(JSON.parse(res.body), [{ data: { add: 3 } }, { data: null, errors: [{ message: 'Internal Server Error' }] }])
 })
+
+test('POST batched query has an individual context for each operation', async (t) => {
+  const app = Fastify()
+
+  const contextSpy = sinon.spy()
+
+  const schema = `
+      type Query {
+        test: String
+      }
+    `
+
+  const resolvers = {
+    test: (_, ctx) => contextSpy(ctx.operationId, ctx.operationsCount, ctx.__currentQuery)
+  }
+
+  app.register(GQL, {
+    schema,
+    resolvers,
+    allowBatchedQueries: true
+  })
+
+  await app.inject({
+    method: 'POST',
+    url: '/graphql',
+    body: [
+      {
+        operationName: 'TestQuery',
+        query: 'query TestQuery { test }'
+      },
+      {
+        operationName: 'DoubleQuery',
+        query: 'query DoubleQuery { test }'
+      }
+    ]
+  })
+
+  sinon.assert.calledTwice(contextSpy)
+  sinon.assert.calledWith(contextSpy, 0, 2, sinon.match(/TestQuery/))
+  sinon.assert.calledWith(contextSpy, 1, 2, sinon.match(/DoubleQuery/))
+})
diff --git a/test/gateway/application-hooks.js b/test/gateway/application-hooks.js
index 9e84e5b..4d3c83a 100644
--- a/test/gateway/application-hooks.js
+++ b/test/gateway/application-hooks.js
@@ -1,6 +1,6 @@
 'use strict'
 
-const { test } = require('tap')
+const { test, t } = require('tap')
 const FakeTimers = require('@sinonjs/fake-timers')
 const { GraphQLSchema } = require('graphql')
 const { promisify } = require('util')
@@ -10,18 +10,24 @@ const buildFederationSchema = require('../../lib/federation')
 
 const immediate = promisify(setImmediate)
 
+t.beforeEach(({ context }) => {
+  context.clock = FakeTimers.install({
+    shouldClearNativeTimers: true,
+    shouldAdvanceTime: true,
+    advanceTimeDelta: 100
+  })
+})
+
+t.afterEach(({ context }) => {
+  context.clock.uninstall()
+})
+
 // ----------------------
 // onGatewayReplaceSchema
 // ----------------------
 test('onGatewayReplaceSchema - polling interval with a new schema should trigger onGatewayReplaceSchema hook', async (t) => {
   t.plan(2)
 
-  const clock = FakeTimers.install({
-    shouldAdvanceTime: true,
-    advanceTimeDelta: 40
-  })
-  t.teardown(() => clock.uninstall())
-
   const resolvers = {
     Query: {
       me: (root, args, context, info) => user
@@ -95,7 +101,9 @@ test('onGatewayReplaceSchema - polling interval with a new schema should trigger
   )
   userService.graphql.defineResolvers(resolvers)
 
-  await clock.tickAsync(2000)
+  for (let i = 0; i < 10; i++) {
+    await t.context.clock.tickAsync(200)
+  }
 
   // We need the event loop to actually spin twice to
   // be able to propagate the change
@@ -106,12 +114,6 @@ test('onGatewayReplaceSchema - polling interval with a new schema should trigger
 test('onGatewayReplaceSchema - should log an error should any errors occur in the hook', async (t) => {
   t.plan(2)
 
-  const clock = FakeTimers.install({
-    shouldAdvanceTime: true,
-    advanceTimeDelta: 40
-  })
-  t.teardown(() => clock.uninstall())
-
   const resolvers = {
     Query: {
       me: (root, args, context, info) => user
@@ -194,7 +196,9 @@ test('onGatewayReplaceSchema - should log an error should any errors occur in th
   )
   userService.graphql.defineResolvers(resolvers)
 
-  await clock.tickAsync(2000)
+  for (let i = 0; i < 10; i++) {
+    await t.context.clock.tickAsync(200)
+  }
 
   // We need the event loop to actually spin twice to
   // be able to propagate the change
diff --git a/test/gateway/pollingInterval.js b/test/gateway/pollingInterval.js
index d8f33d8..f4cf4e1 100644
--- a/test/gateway/pollingInterval.js
+++ b/test/gateway/pollingInterval.js
@@ -1,6 +1,6 @@
 'use strict'
 
-const { test } = require('tap')
+const { test, t } = require('tap')
 
 const FakeTimers = require('@sinonjs/fake-timers')
 
@@ -13,13 +13,19 @@ const WebSocket = require('ws')
 const buildFederationSchema = require('../../lib/federation')
 const GQL = require('../..')
 
-test('Polling schemas with disable cache', async (t) => {
-  const clock = FakeTimers.install({
+t.beforeEach(({ context }) => {
+  context.clock = FakeTimers.install({
+    shouldClearNativeTimers: true,
     shouldAdvanceTime: true,
     advanceTimeDelta: 40
   })
-  t.teardown(() => clock.uninstall())
+})
+
+t.afterEach(({ context }) => {
+  context.clock.uninstall()
+})
 
+test('Polling schemas with disable cache', async (t) => {
   const resolvers = {
     Query: {
       me: (root, args, context, info) => user
@@ -103,12 +109,6 @@ test('Polling schemas with disable cache', async (t) => {
 })
 
 test('Polling schemas', async (t) => {
-  const clock = FakeTimers.install({
-    shouldAdvanceTime: true,
-    advanceTimeDelta: 40
-  })
-  t.teardown(() => clock.uninstall())
-
   const resolvers = {
     Query: {
       me: (root, args, context, info) => user
@@ -234,7 +234,9 @@ test('Polling schemas', async (t) => {
   )
   userService.graphql.defineResolvers(resolvers)
 
-  await clock.tickAsync(2000)
+  for (let i = 0; i < 10; i++) {
+    await t.context.clock.tickAsync(200)
+  }
 
   // We need the event loop to actually spin twice to
   // be able to propagate the change
@@ -336,12 +338,6 @@ test('Polling schemas (gateway.polling interval is not a number)', async (t) =>
 })
 
 test("Polling schemas (if service is down, schema shouldn't be changed)", async (t) => {
-  const clock = FakeTimers.install({
-    shouldAdvanceTime: true,
-    advanceTimeDelta: 40
-  })
-  t.teardown(() => clock.uninstall())
-
   const resolvers = {
     Query: {
       me: (root, args, context, info) => user
@@ -381,7 +377,7 @@ test("Polling schemas (if service is down, schema shouldn't be changed)", async
   })
 
   await userService.listen({ port: 0 })
-  await clock.tickAsync()
+  await t.context.clock.tickAsync()
 
   const userServicePort = userService.server.address().port
 
@@ -397,7 +393,7 @@ test("Polling schemas (if service is down, schema shouldn't be changed)", async
     }
   })
 
-  await clock.tickAsync()
+  await t.context.clock.tickAsync()
 
   {
     const { body } = await gateway.inject({
@@ -418,7 +414,7 @@ test("Polling schemas (if service is down, schema shouldn't be changed)", async
       })
     })
 
-    await clock.tickAsync()
+    await t.context.clock.tickAsync()
 
     t.same(JSON.parse(body), {
       data: {
@@ -463,7 +459,7 @@ test("Polling schemas (if service is down, schema shouldn't be changed)", async
   }
 
   await userService.close()
-  await clock.tickAsync(500)
+  await t.context.clock.tickAsync(500)
 
   {
     const { body } = await gateway.inject({
@@ -622,12 +618,6 @@ test('Polling schemas (if service is mandatory, exception should be thrown)', as
 })
 
 test('Polling schemas (cache should be cleared)', async (t) => {
-  const clock = FakeTimers.install({
-    shouldAdvanceTime: true,
-    advanceTimeDelta: 40
-  })
-  t.teardown(() => clock.uninstall())
-
   const user = {
     id: 'u1',
     name: 'John',
@@ -729,7 +719,9 @@ test('Polling schemas (cache should be cleared)', async (t) => {
     }
   })
 
-  await clock.tickAsync(10000)
+  for (let i = 0; i < 100; i++) {
+    await t.context.clock.tickAsync(100)
+  }
 
   // We need the event loop to actually spin twice to
   // be able to propagate the change
@@ -793,11 +785,6 @@ test('Polling schemas (cache should be cleared)', async (t) => {
 })
 
 test('Polling schemas (should properly regenerate the schema when a downstream service restarts)', async (t) => {
-  const clock = FakeTimers.install({
-    shouldAdvanceTime: true,
-    advanceTimeDelta: 40
-  })
-  t.teardown(() => clock.uninstall())
   const oldSchema = `
     type Query {
       me: User
@@ -916,7 +903,9 @@ test('Polling schemas (should properly regenerate the schema when a downstream s
 
   await restartedUserService.listen({ port: userServicePort })
 
-  await clock.tickAsync(10000)
+  for (let i = 0; i < 100; i++) {
+    await t.context.clock.tickAsync(100)
+  }
 
   // We need the event loop to actually spin twice to
   // be able to propagate the change
@@ -982,12 +971,6 @@ test('Polling schemas (should properly regenerate the schema when a downstream s
 test('Polling schemas (subscriptions should be handled)', async (t) => {
   t.plan(12)
 
-  const clock = FakeTimers.install({
-    shouldAdvanceTime: true,
-    advanceTimeDelta: 40
-  })
-  t.teardown(() => clock.uninstall())
-
   const user = {
     id: 'u1',
     name: 'John',
@@ -1172,7 +1155,7 @@ test('Polling schemas (subscriptions should be handled)', async (t) => {
 
   userService.graphql.defineResolvers(resolvers)
 
-  await clock.tickAsync(10000)
+  await t.context.clock.tickAsync(10000)
 
   // We need the event loop to actually spin twice to
   // be able to propagate the change
diff --git a/test/gateway/retry-failed-services.js b/test/gateway/retry-failed-services.js
index db168db..b8ba440 100644
--- a/test/gateway/retry-failed-services.js
+++ b/test/gateway/retry-failed-services.js
@@ -1,6 +1,6 @@
 'use strict'
 
-const { test } = require('tap')
+const { test, t } = require('tap')
 const Fastify = require('fastify')
 const { GraphQLSchema } = require('graphql')
 const GQL = require('../..')
@@ -112,18 +112,25 @@ const postService = {
   }
 }
 
-test('gateway - retry mandatory failed services on startup', async (t) => {
-  t.plan(5)
-  const clock = FakeTimers.install({
+t.beforeEach(({ context }) => {
+  context.clock = FakeTimers.install({
+    shouldClearNativeTimers: true,
     shouldAdvanceTime: true,
     advanceTimeDelta: 100
   })
+})
 
-  const service1 = await createTestService(5001, userService.schema, userService.resolvers)
+t.afterEach(({ context }) => {
+  context.clock.uninstall()
+})
+
+test('gateway - retry mandatory failed services on startup', async (t) => {
+  t.plan(5)
+  const service1 = await createTestService(0, userService.schema, userService.resolvers)
 
   let service2 = null
-  setTimeout(async () => {
-    service2 = await createTestService(5002, postService.schema, postService.resolvers)
+  t.context.clock.setTimeout(async () => {
+    service2 = await createTestService(5113, postService.schema, postService.resolvers)
   }, 5000)
 
   const app = Fastify()
@@ -131,7 +138,6 @@ test('gateway - retry mandatory failed services on startup', async (t) => {
     await app.close()
     await service1.close()
     await service2.close()
-    clock.uninstall()
   })
 
   await app.register(GQL, {
@@ -140,12 +146,12 @@ test('gateway - retry mandatory failed services on startup', async (t) => {
       services: [
         {
           name: 'user',
-          url: 'http://localhost:5001/graphql',
+          url: `http://localhost:${service1.server.address().port}/graphql`,
           mandatory: false
         },
         {
           name: 'post',
-          url: 'http://localhost:5002/graphql',
+          url: 'http://localhost:5113/graphql',
           mandatory: true
         }
       ]
@@ -178,7 +184,7 @@ test('gateway - retry mandatory failed services on startup', async (t) => {
     body: JSON.stringify({ query })
   })
 
-  t.same(JSON.parse(res.body), {
+  t.same(await res.json(), {
     errors: [
       {
         message: 'Cannot query field "posts" on type "User".',
@@ -188,8 +194,8 @@ test('gateway - retry mandatory failed services on startup', async (t) => {
     data: null
   })
 
-  for (let i = 0; i < 10; i++) {
-    await clock.tickAsync(2000)
+  for (let i = 0; i < 100; i++) {
+    await t.context.clock.tickAsync(100)
   }
 
   const res1 = await app.inject({
@@ -216,24 +222,18 @@ test('gateway - retry mandatory failed services on startup', async (t) => {
 
 test('gateway - should not call onGatewayReplaceSchemaHandler if the hook is not specified', async (t) => {
   t.plan(2)
-  const clock = FakeTimers.install({
-    shouldAdvanceTime: true,
-    advanceTimeDelta: 100
-  })
-
-  const service1 = await createTestService(5001, userService.schema, userService.resolvers)
+  const service1 = await createTestService(0, userService.schema, userService.resolvers)
 
   let service2 = null
-  setTimeout(async () => {
-    service2 = await createTestService(5002, postService.schema, postService.resolvers)
-  }, 5000)
+  t.context.clock.setTimeout(async () => {
+    service2 = await createTestService(5111, postService.schema, postService.resolvers)
+  }, 2000)
 
   const app = Fastify()
   t.teardown(async () => {
     await app.close()
     await service1.close()
     await service2.close()
-    clock.uninstall()
   })
 
   await app.register(GQL, {
@@ -242,12 +242,12 @@ test('gateway - should not call onGatewayReplaceSchemaHandler if the hook is not
       services: [
         {
           name: 'user',
-          url: 'http://localhost:5001/graphql',
+          url: `http://localhost:${service1.server.address().port}/graphql`,
           mandatory: false
         },
         {
           name: 'post',
-          url: 'http://localhost:5002/graphql',
+          url: 'http://localhost:5111/graphql',
           mandatory: true
         }
       ],
@@ -276,7 +276,7 @@ test('gateway - should not call onGatewayReplaceSchemaHandler if the hook is not
     body: JSON.stringify({ query })
   })
 
-  t.same(JSON.parse(res.body), {
+  t.same(await res.json(), {
     errors: [
       {
         message: 'Cannot query field "posts" on type "User".',
@@ -286,8 +286,8 @@ test('gateway - should not call onGatewayReplaceSchemaHandler if the hook is not
     data: null
   })
 
-  for (let i = 0; i < 10; i++) {
-    await clock.tickAsync(1000)
+  for (let i = 0; i < 100; i++) {
+    await t.context.clock.tickAsync(100)
   }
 
   const res1 = await app.inject({
@@ -314,18 +314,13 @@ test('gateway - should not call onGatewayReplaceSchemaHandler if the hook is not
 
 test('gateway - dont retry non-mandatory failed services on startup', async (t) => {
   t.plan(2)
-  const clock = FakeTimers.install({
-    shouldAdvanceTime: true,
-    advanceTimeDelta: 100
-  })
-
-  const service1 = await createTestService(5001, userService.schema, userService.resolvers)
+  const service1 = await createTestService(0, userService.schema, userService.resolvers)
 
   const app = Fastify()
   t.teardown(async () => {
     await app.close()
     await service1.close()
-    clock.uninstall()
+    t.context.clock.uninstall()
   })
 
   app.register(GQL, {
@@ -334,12 +329,12 @@ test('gateway - dont retry non-mandatory failed services on startup', async (t)
       services: [
         {
           name: 'user',
-          url: 'http://localhost:5001/graphql',
+          url: `http://localhost:${service1.server.address().port}/graphql`,
           mandatory: false
         },
         {
           name: 'post',
-          url: 'http://localhost:5002/graphql',
+          url: 'http://localhost:5112/graphql',
           mandatory: false
         }
       ],
@@ -367,7 +362,7 @@ test('gateway - dont retry non-mandatory failed services on startup', async (t)
     body: JSON.stringify({ query })
   })
 
-  t.same(JSON.parse(res.body), {
+  t.same(await res.json(), {
     errors: [
       {
         message: 'Cannot query field "posts" on type "User".',
@@ -377,8 +372,8 @@ test('gateway - dont retry non-mandatory failed services on startup', async (t)
     data: null
   })
 
-  for (let i = 0; i < 10; i++) {
-    await clock.tickAsync(1500)
+  for (let i = 0; i < 100; i++) {
+    await t.context.clock.tickAsync(150)
   }
 
   const res1 = await app.inject({
@@ -388,7 +383,7 @@ test('gateway - dont retry non-mandatory failed services on startup', async (t)
     body: JSON.stringify({ query })
   })
 
-  t.same(JSON.parse(res1.body), {
+  t.same(await res1.json(), {
     errors: [
       {
         message: 'Cannot query field "posts" on type "User".',
@@ -401,16 +396,11 @@ test('gateway - dont retry non-mandatory failed services on startup', async (t)
 
 test('gateway - should log error if retry throws', async (t) => {
   t.plan(1)
-  const clock = FakeTimers.install({
-    shouldAdvanceTime: true,
-    advanceTimeDelta: 100
-  })
-
-  const service1 = await createTestService(5001, userService.schema, userService.resolvers)
+  const service1 = await createTestService(0, userService.schema, userService.resolvers)
 
   let service2 = null
-  setTimeout(async () => {
-    service2 = await createTestService(5002, postService.schema, postService.resolvers)
+  t.context.clock.setTimeout(async () => {
+    service2 = await createTestService(5114, postService.schema, postService.resolvers)
   }, 2000)
 
   const app = Fastify()
@@ -427,7 +417,6 @@ test('gateway - should log error if retry throws', async (t) => {
     await app.close()
     await service1.close()
     await service2.close()
-    clock.uninstall()
   })
 
   await app.register(GQL, {
@@ -436,17 +425,17 @@ test('gateway - should log error if retry throws', async (t) => {
       services: [
         {
           name: 'user',
-          url: 'http://localhost:5001/graphql',
+          url: `http://localhost:${service1.server.address().port}/graphql`,
           mandatory: false
         },
         {
           name: 'post',
-          url: 'http://localhost:5002/graphql',
+          url: 'http://localhost:5114/graphql',
           mandatory: true
         }
       ],
       retryServicesCount: 1,
-      retryServicesInterval: 2000
+      retryServicesInterval: 3000
     }
   })
 
@@ -456,18 +445,13 @@ test('gateway - should log error if retry throws', async (t) => {
 
   await app.ready()
 
-  for (let i = 0; i < 10; i++) {
-    await clock.tickAsync(1000)
+  for (let i = 0; i < 200; i++) {
+    await t.context.clock.tickAsync(50)
   }
 })
 
 test('gateway - stop retrying after no. of retries exceeded', async (t) => {
   t.plan(1)
-  const clock = FakeTimers.install({
-    shouldAdvanceTime: true,
-    advanceTimeDelta: 100
-  })
-
   const service1 = await createTestService(0, userService.schema, userService.resolvers)
 
   const app = Fastify()
@@ -483,7 +467,6 @@ test('gateway - stop retrying after no. of retries exceeded', async (t) => {
   t.teardown(async () => {
     await app.close()
     await service1.close()
-    clock.uninstall()
   })
 
   await app.register(GQL, {
@@ -497,7 +480,7 @@ test('gateway - stop retrying after no. of retries exceeded', async (t) => {
         },
         {
           name: 'post',
-          url: 'http://localhost:5002/graphql',
+          url: 'http://localhost:5115/graphql',
           mandatory: true
         }
       ],
@@ -508,7 +491,7 @@ test('gateway - stop retrying after no. of retries exceeded', async (t) => {
 
   await app.ready()
 
-  for (let i = 0; i < 10; i++) {
-    await clock.tickAsync(1500)
+  for (let i = 0; i < 100; i++) {
+    await t.context.clock.tickAsync(150)
   }
 })
diff --git a/test/gateway/value-types.js b/test/gateway/value-types.js
index c7e9bc4..8ffce76 100644
--- a/test/gateway/value-types.js
+++ b/test/gateway/value-types.js
@@ -1,6 +1,6 @@
 'use strict'
 
-const { test } = require('tap')
+const { test, t } = require('tap')
 const Fastify = require('fastify')
 const FakeTimers = require('@sinonjs/fake-timers')
 const { promisify } = require('util')
@@ -485,6 +485,18 @@ const commentMutation = `
   }  
 `
 
+t.beforeEach(({ context }) => {
+  context.clock = FakeTimers.install({
+    shouldClearNativeTimers: true,
+    shouldAdvanceTime: true,
+    advanceTimeDelta: 100
+  })
+})
+
+t.afterEach(({ context }) => {
+  context.clock.uninstall()
+})
+
 test('Should be able to query with value types', async (t) => {
   const [userService, userServicePort] = await createService(t, userSchema, userResolvers)
   const [postService, postServicePort] = await createService(t, postSchema, postResolvers)
@@ -705,12 +717,6 @@ test('Should be able to query top-level with value types', async (t) => {
 })
 
 test('Should be able to query with value types and polling', async (t) => {
-  const clock = FakeTimers.install({
-    shouldAdvanceTime: true,
-    advanceTimeDelta: 40
-  })
-  t.teardown(() => clock.uninstall())
-
   const [userService, userServicePort] = await createService(t, userSchema, userResolvers)
   const [postService, postServicePort] = await createService(t, postSchema, postResolvers)
   const [commentService, commentServicePort] = await createService(t, commentSchema, commentResolvers)
@@ -797,7 +803,7 @@ test('Should be able to query with value types and polling', async (t) => {
     }
   })
 
-  await clock.tickAsync(2000)
+  await t.context.clock.tickAsync(2000)
 
   // We need the event loop to actually spin twice to
   // be able to propagate the change
diff --git a/test/options.js b/test/options.js
index bf294f6..a4eccf0 100644
--- a/test/options.js
+++ b/test/options.js
@@ -30,8 +30,6 @@ const resolvers = {
 }
 
 test('call compileQuery with correct options if compilerOptions specified', async t => {
-  t.plan(1)
-
   const app = Fastify()
   t.teardown(() => app.close())
 
@@ -82,5 +80,5 @@ test('call compileQuery with correct options if compilerOptions specified', asyn
     body: JSON.stringify({ query })
   })
 
-  t.ok(compileQueryStub.calledOnceWith(sinon.match.any, sinon.match.any, sinon.match.any, { customJSONSerializer: true }))
+  sinon.assert.calledOnceWithExactly(compileQueryStub, sinon.match.any, sinon.match.any, sinon.match.any, { customJSONSerializer: true })
 })
diff --git a/test/subscription-client.js b/test/subscription-client.js
index 8071851..0717fb9 100644
--- a/test/subscription-client.js
+++ b/test/subscription-client.js
@@ -379,7 +379,14 @@ test('subscription client sends GQL_CONNECTION_KEEP_ALIVE when the keep alive op
         ws.send(JSON.stringify({ id: '1', type: 'connection_ack' }))
       } else if (data.type === 'start') {
         ws.send(JSON.stringify({ id: '2', type: 'complete' }))
-      } else if (data.type === 'ka') {
+      } else if (data.type === 'ping') {
+        // this is a client sent ping, we reply with our pong
+        ws.send(JSON.stringify({ id: '3', type: 'pong' }))
+
+        // send a ping to the client so that it replies with its own pong
+        // pings and pongs are bidirectional
+        ws.send(JSON.stringify({ id: '4', type: 'ping' }))
+      } else if (data.type === 'pong') {
         client.close()
         server.close()
         t.end()
@@ -414,7 +421,7 @@ test('subscription client not throwing error on GQL_CONNECTION_KEEP_ALIVE type p
     ws.send(JSON.stringify({ id: undefined, type: 'connection_ack' }))
 
     clock.setInterval(() => {
-      ws.send(JSON.stringify({ type: 'ka' }))
+      ws.send(JSON.stringify({ type: 'ping' }))
     }, 200)
   })
 
diff --git a/test/subscription-connection.js b/test/subscription-connection.js
index 323543d..2624421 100644
--- a/test/subscription-connection.js
+++ b/test/subscription-connection.js
@@ -341,6 +341,69 @@ test('subscription connection handles when GQL_START is called before GQL_INIT',
   }))
 })
 
+test('subscription connection replies to GQL_CONNECTION_KEEP_ALIVE message with GQL_CONNECTION_KEEP_ALIVE_ACK', async (t) => {
+  t.plan(1)
+
+  const sc = new SubscriptionConnection(
+    {
+      on () {},
+      close () {},
+      send (message) {
+        t.equal(
+          JSON.stringify({
+            type: 'pong',
+            id: 1
+          }),
+          message
+        )
+      },
+      protocol: GRAPHQL_TRANSPORT_WS
+    },
+    {}
+  )
+
+  await sc.handleMessage(
+    JSON.stringify({
+      id: 1,
+      type: 'ping',
+      payload: {}
+    })
+  )
+})
+
+test('subscription connection does not error if client sends GQL_CONNECTION_KEEP_ALIVE_ACK', async (t) => {
+  t.plan(1)
+
+  const sc = new SubscriptionConnection(
+    {
+      on () {},
+      close () {},
+      send (message) {
+        t.fail()
+      },
+      protocol: GRAPHQL_TRANSPORT_WS
+    },
+    {}
+  )
+
+  await sc.handleMessage(
+    JSON.stringify({
+      id: 1,
+      type: 'pong',
+      payload: {}
+    })
+  )
+
+  await sc.handleMessage(
+    JSON.stringify({
+      id: 1,
+      type: 'complete'
+    })
+  )
+
+  t.equal(sc.subscriptionContexts.size, 0)
+})
+
 test('subscription connection extends context with onConnect return value', async (t) => {
   t.plan(3)
 
@@ -574,24 +637,17 @@ test('subscription connection handles query when fullWsTransport: true', async (
     })
   )
 
-  t.ok(
-    send.withArgs(
-      JSON.stringify({
-        type: 'next',
-        id: 1,
-        payload: {}
-      })
-    ).calledOnce
-  )
-  t.ok(
-    send.withArgs(
-      JSON.stringify({
-        type: 'complete',
-        id: 1,
-        payload: null
-      })
-    ).calledOnce
-  )
+  sinon.assert.calledTwice(send)
+  sinon.assert.calledWith(send, JSON.stringify({
+    type: 'next',
+    id: 1,
+    payload: {}
+  }))
+  sinon.assert.calledWith(send, JSON.stringify({
+    type: 'complete',
+    id: 1,
+    payload: null
+  }))
 })
 
 test('subscription connection handles mutation when fullWsTransport: true', async (t) => {
@@ -625,24 +681,17 @@ test('subscription connection handles mutation when fullWsTransport: true', asyn
     })
   )
 
-  t.ok(
-    send.withArgs(
-      JSON.stringify({
-        type: 'next',
-        id: 1,
-        payload: {}
-      })
-    ).calledOnce
-  )
-  t.ok(
-    send.withArgs(
-      JSON.stringify({
-        type: 'complete',
-        id: 1,
-        payload: null
-      })
-    ).calledOnce
-  )
+  sinon.assert.calledTwice(send)
+  sinon.assert.calledWith(send, JSON.stringify({
+    type: 'next',
+    id: 1,
+    payload: {}
+  }))
+  sinon.assert.calledWith(send, JSON.stringify({
+    type: 'complete',
+    id: 1,
+    payload: null
+  }))
 })
 
 test('subscription data is released right after it ends', async (t) => {
diff --git a/test/subscription-protocol.js b/test/subscription-protocol.js
index 4bc8c26..4a3b6a8 100644
--- a/test/subscription-protocol.js
+++ b/test/subscription-protocol.js
@@ -20,7 +20,8 @@ test('getProtocolByName returns correct protocol message types', t => {
     GQL_CONNECTION_INIT: 'connection_init',
     GQL_CONNECTION_ACK: 'connection_ack',
     GQL_CONNECTION_ERROR: 'connection_error',
-    GQL_CONNECTION_KEEP_ALIVE: 'ka',
+    GQL_CONNECTION_KEEP_ALIVE: 'ping',
+    GQL_CONNECTION_KEEP_ALIVE_ACK: 'pong',
     GQL_CONNECTION_TERMINATE: 'connection_terminate',
     GQL_START: 'subscribe',
     GQL_DATA: 'next',
diff --git a/test/subscription.js b/test/subscription.js
index 3067464..143e1ff 100644
--- a/test/subscription.js
+++ b/test/subscription.js
@@ -1,6 +1,6 @@
 'use strict'
 
-const { test } = require('tap')
+const { test, t } = require('tap')
 const Fastify = require('fastify')
 const WebSocket = require('ws')
 const mq = require('mqemitter')
@@ -10,6 +10,18 @@ const GQL = require('..')
 
 const FakeTimers = require('@sinonjs/fake-timers')
 
+t.beforeEach(({ context }) => {
+  context.clock = FakeTimers.install({
+    shouldClearNativeTimers: true,
+    shouldAdvanceTime: true,
+    advanceTimeDelta: 40
+  })
+})
+
+t.afterEach(({ context }) => {
+  context.clock.uninstall()
+})
+
 test('subscription server replies with connection_ack', t => {
   const app = Fastify()
   t.teardown(() => app.close())
@@ -730,15 +742,9 @@ test('subscription connection is closed if context function throws', t => {
 })
 
 test('subscription server sends update to subscriptions with custom async context', t => {
-  const clock = FakeTimers.install({
-    shouldAdvanceTime: true,
-    advanceTimeDelta: 40
-  })
-
   const app = Fastify()
   t.teardown(async () => {
     await app.close()
-    clock.uninstall()
   })
 
   const sendTestQuery = () => {
@@ -837,7 +843,7 @@ test('subscription server sends update to subscriptions with custom async contex
     subscription: {
       emitter,
       context: async () => {
-        await clock.tickAsync(200)
+        await t.context.clock.tickAsync(200)
         return { topic: 'NOTIFICATION_ADDED' }
       }
     }
@@ -917,15 +923,9 @@ test('subscription server sends update to subscriptions with custom async contex
 })
 
 test('subscription connection is closed if async context function throws', t => {
-  const clock = FakeTimers.install({
-    shouldAdvanceTime: true,
-    advanceTimeDelta: 40
-  })
-
   const app = Fastify()
   t.teardown(async () => {
     await app.close()
-    clock.uninstall()
   })
 
   const schema = `
@@ -945,7 +945,7 @@ test('subscription connection is closed if async context function throws', t =>
     resolvers,
     subscription: {
       context: async function () {
-        await clock.tickAsync(200)
+        await t.context.clock.tickAsync(200)
         throw new Error('kaboom')
       }
     }
diff --git a/test/types/index.ts b/test/types/index.ts
index de78815..92da02f 100644
--- a/test/types/index.ts
+++ b/test/types/index.ts
@@ -1,4 +1,4 @@
-import { expectAssignable, expectError } from 'tsd'
+import { expectAssignable, expectError, expectType } from 'tsd'
 /* eslint-disable no-unused-expressions */
 import { EventEmitter } from 'events'
 // eslint-disable-next-line no-unused-vars
@@ -780,3 +780,10 @@ expectError(() => {
 expectError(() => {
   return mercurius.defaultErrorFormatter({}, undefined)
 })
+
+// Context contains correct information about (batched) query identity
+app.graphql.addHook('onResolution', async function (_execution, context) {
+  expectType<number | undefined>(context.operationId)
+  expectType<number | undefined>(context.operationsCount)
+  expectType<string>(context.__currentQuery)
+})