Repository: mercurius-js/mercurius
Pull request #395 (2021-02-14 23:22:13)

Fix intermittent test failures
Fixes: https://github.com/mercurius-js/mercurius/issues/394

mcollina: is this ready to go?
jonnydgreen: The only thing this PR is missing is the failing test in the CI I just spotted now: https://github.com/mercurius-js/mercurius/issues/394#issuecomment-778858799

If you are happy for me to submit a solution for this test failure in a separate PR (I'm not sure how long it'll take to find the cause of the issue as it only seems to occur in the CI), this PR is good to go! :)
jonnydgreen: Looking into it just now, I think it's due to not waiting enough for the pollingInterval changes to propagate (based on your solution in the same file for `Polling schemas (subscriptions should be handled)`):

Currently:

```js
await clock.tickAsync(2000)
```

Should be:

```js
await clock.tickAsync(10000)

// We need the event loop to actually spin twice to
// be able to propagate the change
await immediate()
await immediate()
```

I'll test this theory now/tomorrow!

Patch:
diff --git a/docs/plugins.md b/docs/plugins.md
index 6e3ac8f..09cc56f 100644
--- a/docs/plugins.md
+++ b/docs/plugins.md
@@ -41,3 +41,37 @@ app.listen(3000)
 And it will be available at `http://localhost:3000/altair` 🎉
 
 Check [here](https://github.com/imolorhe/altair/tree/staging/packages/altair-fastify-plugin) for more information.
+
+## mercurius-apollo-registry
+
+A Mercurius plugin for schema reporting to Apollo Studio.
+
+Check [https://github.com/nearform/mercurius-apollo-registry](https://github.com/nearform/mercurius-apollo-registry) for usage and readme.
+
+```bash
+npm install mercurius-apollo-registry
+```
+
+```js
+const app = Fastify()
+const mercurius = require('mercurius')
+const mercuriusApolloRegistry = require('mercurius-apollo-registry')
+
+const schema = `define schema here`
+const resolvers = { 
+  // ... 
+}
+
+app.register(mercurius, {
+  schema,
+  resolvers,
+  graphiql: true
+})
+
+app.register(mercuriusApolloRegistry, {
+  schema,
+  apiKey: 'REPLACE-THIS-VALUE-WITH-APOLLO-API-KEY'
+})
+
+app.listen(3000)
+```
diff --git a/lib/subscription-connection.js b/lib/subscription-connection.js
index 51740bf..09cf1a7 100644
--- a/lib/subscription-connection.js
+++ b/lib/subscription-connection.js
@@ -37,6 +37,7 @@ module.exports = class SubscriptionConnection {
     this.onConnect = onConnect
     this.onDisconnect = onDisconnect
     this.subscriptionContexts = new Map()
+    this.subscriptionIters = new Map()
     this.context = context
     this.isReady = false
     this.resolveContext = resolveContext
@@ -165,12 +166,12 @@ module.exports = class SubscriptionConnection {
       fastify: this.fastify,
       pubsub: this.subscriber
     })
-    this.subscriptionContexts.set(id, sc)
 
     const document = typeof query !== 'string' ? query : parse(query)
     const schema = this.fastify ? this.fastify.graphql.schema : undefined
+    this.subscriptionContexts.set(id, sc)
 
-    const result = await subscribe(
+    const subIter = await subscribe(
       schema,
       document,
       {}, // rootValue
@@ -186,14 +187,15 @@ module.exports = class SubscriptionConnection {
       variables,
       operationName
     )
+    this.subscriptionIters.set(id, subIter)
 
-    if (result.errors) {
-      this.fastify.log.error(result.errors)
-      throw result.errors[0]
+    if (subIter.errors) {
+      this.fastify.log.error(subIter.errors)
+      throw subIter.errors[0]
     }
 
     // TODO implement backpressure
-    for await (const value of result) {
+    for await (const value of subIter) {
       this.sendMessage(GQL_DATA, data.id, value)
     }
 
@@ -201,20 +203,25 @@ module.exports = class SubscriptionConnection {
   }
 
   handleGQLStop (data) {
-    const subscriptionContextToRemove = this.subscriptionContexts.get(data.id)
-    if (!subscriptionContextToRemove) {
-      return
+    const sc = this.subscriptionContexts.get(data.id)
+    if (sc) {
+      sc.close && sc.close()
+      this.subscriptionContexts.delete(data.id)
+    }
+    const subIter = this.subscriptionIters.get(data.id)
+    if (subIter) {
+      subIter.return && subIter.return()
+      this.subscriptionIters.delete(data.id)
     }
-
-    subscriptionContextToRemove.close()
-    this.subscriptionContexts.delete(data.id)
   }
 
   handleConnectionClose () {
     Array
       .from(this.subscriptionContexts.values())
-      .map(subscriptionContext =>
-        subscriptionContext.close())
+      .map((sc) => sc.close())
+    Array
+      .from(this.subscriptionIters.values())
+      .map((subIter) => subIter.return && subIter.return())
     this.socket.close()
 
     if (typeof this.onDisconnect === 'function') {


Test patch:
diff --git a/test/app-decorator.js b/test/app-decorator.js
index 03883bf..f537ee6 100644
--- a/test/app-decorator.js
+++ b/test/app-decorator.js
@@ -1118,9 +1118,7 @@ test('defineResolvers should throw if field is not defined in schema', async (t)
 
 test('calling extendSchema throws an error if federationMetadata is enabled', async (t) => {
   const service = Fastify()
-  t.tearDown(() => {
-    service.close()
-  })
+  t.tearDown(() => service.close())
   service.register(GQL, {
     schema: `
       extend type Query {
diff --git a/test/gateway/connectionInit-extension.js b/test/gateway/connectionInit-extension.js
index c9e28b9..7ddff2e 100644
--- a/test/gateway/connectionInit-extension.js
+++ b/test/gateway/connectionInit-extension.js
@@ -28,10 +28,10 @@ test('connectionInit extension e2e testing', t => {
   const userService = Fastify()
   const notificationService = Fastify()
   const gateway = Fastify()
-  t.tearDown(() => {
-    userService.close()
-    notificationService.close()
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await userService.close()
+    await notificationService.close()
   })
 
   let userId = 1
diff --git a/test/gateway/extends-directive.js b/test/gateway/extends-directive.js
index f5863b1..7a09b43 100644
--- a/test/gateway/extends-directive.js
+++ b/test/gateway/extends-directive.js
@@ -6,9 +6,6 @@ const GQL = require('../..')
 
 async function createService (t, schema, resolvers = {}) {
   const service = Fastify()
-  t.tearDown(() => {
-    service.close()
-  })
   service.register(GQL, {
     schema,
     resolvers,
@@ -16,7 +13,7 @@ async function createService (t, schema, resolvers = {}) {
   })
   await service.listen(0)
 
-  return service.server.address().port
+  return [service, service.server.address().port]
 }
 
 const users = {
@@ -46,7 +43,7 @@ const posts = {
 }
 
 test('gateway handles @extends directive correctly', async (t) => {
-  const userServicePort = await createService(t, `
+  const [userService, userServicePort] = await createService(t, `
     type Query @extends {
       me: User
     }
@@ -68,7 +65,7 @@ test('gateway handles @extends directive correctly', async (t) => {
     }
   })
 
-  const postServicePort = await createService(t, `
+  const [postService, postServicePort] = await createService(t, `
     type Post @key(fields: "pid") {
       pid: ID!
       title: String
@@ -111,8 +108,10 @@ test('gateway handles @extends directive correctly', async (t) => {
   })
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await postService.close()
+    await userService.close()
   })
   gateway.register(GQL, {
     gateway: {
@@ -188,7 +187,7 @@ test('gateway passes field arguments through to types labeled by @extends direct
     }
   }
 
-  const userServicePort = await createService(t, `
+  const [userService, userServicePort] = await createService(t, `
     type Query @extends {
       me: User
     }
@@ -210,7 +209,7 @@ test('gateway passes field arguments through to types labeled by @extends direct
     }
   })
 
-  const postServicePort = await createService(t, `
+  const [postService, postServicePort] = await createService(t, `
     type Post @key(fields: "pid") {
       pid: ID!
       author: User
@@ -240,8 +239,10 @@ test('gateway passes field arguments through to types labeled by @extends direct
   })
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await postService.close()
+    await userService.close()
   })
   gateway.register(GQL, {
     gateway: {
@@ -335,7 +336,7 @@ test('gateway distributes query correctly to services when querying with inline
     }
   }
 
-  const userServicePort = await createService(t, `
+  const [userService, userServicePort] = await createService(t, `
     type Query @extends {
       me: UserUnion
     }
@@ -363,7 +364,7 @@ test('gateway distributes query correctly to services when querying with inline
     }
   })
 
-  const postServicePort = await createService(t, `
+  const [postService, postServicePort] = await createService(t, `
     type Post @key(fields: "pid") {
       pid: ID!
       author: User
@@ -393,8 +394,10 @@ test('gateway distributes query correctly to services when querying with inline
   })
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await postService.close()
+    await userService.close()
   })
   gateway.register(GQL, {
     gateway: {
diff --git a/test/gateway/hooks.js b/test/gateway/hooks.js
index 0e7fb0b..2cc3108 100644
--- a/test/gateway/hooks.js
+++ b/test/gateway/hooks.js
@@ -10,17 +10,13 @@ const immediate = promisify(setImmediate)
 
 async function createTestService (t, schema, resolvers = {}) {
   const service = Fastify()
-  t.tearDown(() => {
-    service.close()
-  })
   service.register(GQL, {
     schema,
     resolvers,
     federationMetadata: true
   })
   await service.listen(0)
-
-  return service.server.address().port
+  return [service, service.server.address().port]
 }
 
 const users = {
@@ -102,7 +98,7 @@ async function createTestGatewayServer (t) {
       }
     }
   }
-  const userServicePort = await createTestService(t, userServiceSchema, userServiceResolvers)
+  const [userService, userServicePort] = await createTestService(t, userServiceSchema, userServiceResolvers)
 
   // Post service
   const postServiceSchema = `
@@ -140,11 +136,13 @@ async function createTestGatewayServer (t) {
       topPosts: (root, { count = 2 }) => Object.values(posts).slice(0, count)
     }
   }
-  const postServicePort = await createTestService(t, postServiceSchema, postServiceResolvers)
+  const [postService, postServicePort] = await createTestService(t, postServiceSchema, postServiceResolvers)
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await userService.close()
+    await postService.close()
   })
   gateway.register(GQL, {
     gateway: {
@@ -199,7 +197,7 @@ test('gateway - hooks', async (t) => {
   app.graphql.addHook('preGatewayExecution', async function (schema, document, context) {
     await immediate()
     t.type(schema, GraphQLSchema)
-    t.true(document, 'object')
+    t.type(document, 'object')
     t.type(context, 'object')
     t.ok('preGatewayExecution called')
   })
diff --git a/test/gateway/methods.js b/test/gateway/methods.js
index 4699a5e..6b6f17d 100644
--- a/test/gateway/methods.js
+++ b/test/gateway/methods.js
@@ -6,9 +6,6 @@ const GQL = require('../..')
 
 async function createService (t, schema, resolvers = {}) {
   const service = Fastify()
-  t.tearDown(() => {
-    service.close()
-  })
   service.register(GQL, {
     schema,
     resolvers,
@@ -16,11 +13,11 @@ async function createService (t, schema, resolvers = {}) {
   })
   await service.listen(0)
 
-  return service.server.address().port
+  return [service, service.server.address().port]
 }
 
 test('calling defineLoaders throws an error in gateway mode', async (t) => {
-  const port = await createService(t, `
+  const [service, port] = await createService(t, `
     extend type Query {
       me: User
     }
@@ -32,8 +29,9 @@ test('calling defineLoaders throws an error in gateway mode', async (t) => {
   `)
 
   const app = Fastify()
-  t.tearDown(() => {
-    app.close()
+  t.tearDown(async () => {
+    await app.close()
+    await service.close()
   })
 
   app.register(GQL, {
@@ -59,7 +57,7 @@ test('calling defineLoaders throws an error in gateway mode', async (t) => {
 })
 
 test('calling defineResolvers throws an error in gateway mode', async (t) => {
-  const port = await createService(t, `
+  const [service, port] = await createService(t, `
     extend type Query {
       me: User
     }
@@ -71,8 +69,9 @@ test('calling defineResolvers throws an error in gateway mode', async (t) => {
   `)
 
   const app = Fastify()
-  t.tearDown(() => {
-    app.close()
+  t.tearDown(async () => {
+    await app.close()
+    await service.close()
   })
 
   app.register(GQL, {
@@ -98,7 +97,7 @@ test('calling defineResolvers throws an error in gateway mode', async (t) => {
 })
 
 test('calling extendSchema throws an error in gateway mode', async (t) => {
-  const port = await createService(t, `
+  const [service, port] = await createService(t, `
     extend type Query {
       me: User
     }
@@ -110,8 +109,9 @@ test('calling extendSchema throws an error in gateway mode', async (t) => {
   `)
 
   const app = Fastify()
-  t.tearDown(() => {
-    app.close()
+  t.tearDown(async () => {
+    await app.close()
+    await service.close()
   })
 
   app.register(GQL, {
diff --git a/test/gateway/pollingInterval.js b/test/gateway/pollingInterval.js
index df004f7..7e3709e 100644
--- a/test/gateway/pollingInterval.js
+++ b/test/gateway/pollingInterval.js
@@ -18,6 +18,7 @@ test('Polling schemas', async (t) => {
     shouldAdvanceTime: true,
     advanceTimeDelta: 40
   })
+  t.tearDown(() => clock.uninstall())
 
   const resolvers = {
     Query: {
@@ -36,6 +37,10 @@ test('Polling schemas', async (t) => {
 
   const userService = Fastify()
   const gateway = Fastify()
+  t.tearDown(async () => {
+    await gateway.close()
+    await userService.close()
+  })
 
   userService.register(GQL, {
     schema: `
@@ -68,8 +73,6 @@ test('Polling schemas', async (t) => {
     }
   })
 
-  await gateway.listen(0)
-
   const res = await gateway.inject({
     method: 'POST',
     headers: {
@@ -143,7 +146,11 @@ test('Polling schemas', async (t) => {
   userService.graphql.defineResolvers(resolvers)
 
   await clock.tickAsync(2000)
-  await clock.tickAsync()
+
+  // We need the event loop to actually spin twice to
+  // be able to propagate the change
+  await immediate()
+  await immediate()
 
   const res3 = await gateway.inject({
     method: 'POST',
@@ -173,10 +180,6 @@ test('Polling schemas', async (t) => {
       }
     }
   })
-
-  await gateway.close()
-  await userService.close()
-  clock.uninstall()
 })
 
 test('Polling schemas (gateway.polling interval is not a number)', async (t) => {
@@ -248,6 +251,7 @@ test("Polling schemas (if service is down, schema shouldn't be changed)", async
     shouldAdvanceTime: true,
     advanceTimeDelta: 40
   })
+  t.tearDown(() => clock.uninstall())
 
   const resolvers = {
     Query: {
@@ -304,7 +308,6 @@ test("Polling schemas (if service is down, schema shouldn't be changed)", async
     }
   })
 
-  await gateway.listen(0)
   await clock.tickAsync()
 
   {
@@ -404,8 +407,6 @@ test("Polling schemas (if service is down, schema shouldn't be changed)", async
       data: null
     })
   }
-
-  clock.uninstall()
 })
 
 test('Polling schemas (if service is mandatory, exception should be thrown)', async (t) => {
@@ -426,6 +427,10 @@ test('Polling schemas (if service is mandatory, exception should be thrown)', as
 
   const userService = Fastify()
   const gateway = Fastify()
+  t.tearDown(async () => {
+    await gateway.close()
+    await userService.close()
+  })
 
   userService.register(GQL, {
     schema: `
@@ -458,8 +463,6 @@ test('Polling schemas (if service is mandatory, exception should be thrown)', as
     }
   })
 
-  await gateway.listen(0)
-
   {
     const { body } = await gateway.inject({
       method: 'POST',
@@ -527,8 +530,6 @@ test('Polling schemas (if service is mandatory, exception should be thrown)', as
   t.rejects(async () => {
     await gateway.graphql.gateway.refresh()
   })
-
-  await gateway.close()
 })
 
 test('Polling schemas (cache should be cleared)', async (t) => {
@@ -536,6 +537,7 @@ test('Polling schemas (cache should be cleared)', async (t) => {
     shouldAdvanceTime: true,
     advanceTimeDelta: 40
   })
+  t.tearDown(() => clock.uninstall())
 
   const user = {
     id: 'u1',
@@ -545,6 +547,10 @@ test('Polling schemas (cache should be cleared)', async (t) => {
 
   const userService = Fastify()
   const gateway = Fastify()
+  t.tearDown(async () => {
+    await gateway.close()
+    await userService.close()
+  })
 
   userService.register(GQL, {
     schema: `
@@ -634,7 +640,12 @@ test('Polling schemas (cache should be cleared)', async (t) => {
     }
   })
 
-  await clock.tickAsync(2000)
+  await clock.tickAsync(10000)
+
+  // We need the event loop to actually spin twice to
+  // be able to propagate the change
+  await immediate()
+  await immediate()
 
   const res2 = await gateway.inject({
     method: 'POST',
@@ -690,10 +701,213 @@ 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 = `
+    directive @extends on INTERFACE | OBJECT
+
+    directive @external on FIELD_DEFINITION | OBJECT
+
+    directive @key(fields: String!) on INTERFACE | OBJECT
+
+    directive @provides(fields: String!) on FIELD_DEFINITION
+
+    directive @requires(fields: String!) on FIELD_DEFINITION
+
+    type Query {
+      me: User
+    }
+
+    type User @key(fields: "id") {
+      id: ID!
+      name: String!
+    }
+  `
+  const user = {
+    id: 'u1',
+    name: 'John',
+    lastName: 'Doe'
+  }
+
+  const userService = Fastify()
+  const gateway = Fastify()
+
+  userService.register(GQL, {
+    schema: oldSchema,
+    resolvers: {
+      Query: {
+        me: (root, args, context, info) => user
+      },
+      User: {
+        __resolveReference: (user, args, context, info) => user
+      }
+    },
+    federationMetadata: true
+  })
+
+  await userService.listen(0)
+
+  const userServicePort = userService.server.address().port
+
+  gateway.register(GQL, {
+    gateway: {
+      services: [
+        {
+          name: 'user',
+          url: `http://localhost:${userServicePort}/graphql`
+        }
+      ],
+      pollingInterval: 2000
+    }
+  })
+
+  const res = await gateway.inject({
+    method: 'POST',
+    headers: {
+      'content-type': 'application/json'
+    },
+    url: '/graphql',
+    body: JSON.stringify({
+      query: `
+        query MainQuery {
+          me {
+            id
+            name
+          }
+        }
+      `
+    })
+  })
+
+  t.deepEqual(JSON.parse(res.body), {
+    data: {
+      me: {
+        id: 'u1',
+        name: 'John'
+      }
+    }
+  })
 
-  await gateway.close()
   await userService.close()
-  clock.uninstall()
+
+  const restartedUserService = Fastify()
+  t.tearDown(async () => {
+    await gateway.close()
+    await userService.close()
+    await restartedUserService.close()
+  })
+
+  const refreshedSchema = `
+    directive @extends on INTERFACE | OBJECT
+
+    directive @external on FIELD_DEFINITION | OBJECT
+
+    directive @key(fields: String!) on INTERFACE | OBJECT
+
+    directive @provides(fields: String!) on FIELD_DEFINITION
+    
+    directive @requires(fields: String!) on FIELD_DEFINITION
+
+    type User @key(fields: "id") {
+      id: ID!
+      lastName: String!
+      name: String!
+    }
+
+    type Mutation {
+      create: User!
+    }
+
+    type Query {
+      me2: User
+    }
+  `
+
+  restartedUserService.register(GQL, {
+    schema: refreshedSchema,
+    resolvers: {
+      Query: {
+        me2: (root, args, context, info) => user
+      },
+      Mutation: {
+        create: (root, args, context, info) => user
+      },
+      User: {
+        __resolveReference: (user, args, context, info) => user
+      }
+    },
+    federationMetadata: true
+  })
+
+  await restartedUserService.listen(userServicePort)
+
+  await clock.tickAsync(10000)
+
+  // We need the event loop to actually spin twice to
+  // be able to propagate the change
+  await immediate()
+  await immediate()
+
+  const res2 = await gateway.inject({
+    method: 'POST',
+    headers: {
+      'content-type': 'application/json'
+    },
+    url: '/graphql',
+    body: JSON.stringify({
+      query: `
+        query MainQuery {
+          me {
+            id
+            name
+          }
+        }
+      `
+    })
+  })
+
+  t.deepEqual(JSON.parse(res2.body), {
+    errors: [
+      {
+        message: 'Cannot query field "me" on type "Query". Did you mean "me2"?',
+        locations: [{ line: 3, column: 11 }]
+      }
+    ],
+    data: null
+  })
+
+  const res3 = await gateway.inject({
+    method: 'POST',
+    headers: {
+      'content-type': 'application/json'
+    },
+    url: '/graphql',
+    body: JSON.stringify({
+      query: `
+        mutation NewMutation {
+          create {
+            id
+            name
+          }
+        }
+      `
+    })
+  })
+
+  t.deepEqual(JSON.parse(res3.body), {
+    data: {
+      create: {
+        id: 'u1',
+        name: 'John'
+      }
+    }
+  })
 })
 
 test('Polling schemas (subscriptions should be handled)', async (t) => {
@@ -701,6 +915,7 @@ test('Polling schemas (subscriptions should be handled)', async (t) => {
     shouldAdvanceTime: true,
     advanceTimeDelta: 40
   })
+  t.teardown(() => clock.uninstall())
 
   const user = {
     id: 'u1',
@@ -738,9 +953,9 @@ test('Polling schemas (subscriptions should be handled)', async (t) => {
   const userService = Fastify()
   const gateway = Fastify()
 
-  t.tearDown(() => {
-    userService.close()
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await userService.close()
   })
 
   userService.register(GQL, {
@@ -798,57 +1013,58 @@ test('Polling schemas (subscriptions should be handled)', async (t) => {
     encoding: 'utf8',
     objectMode: true
   })
-
-  t.tearDown(() => {
-    client.destroy()
-  })
-
+  t.tearDown(client.destroy.bind(client))
   client.setEncoding('utf8')
 
-  client.write(
-    JSON.stringify({
-      type: 'connection_init'
-    })
-  )
-
-  client.write(
-    JSON.stringify({
-      id: 1,
-      type: 'start',
-      payload: {
-        query: `
-          subscription {
-            updatedUser {
-              id
-              name
+  process.nextTick(() => {
+    client.write(
+      JSON.stringify({
+        type: 'connection_init'
+      })
+    )
+
+    client.write(
+      JSON.stringify({
+        id: 1,
+        type: 'start',
+        payload: {
+          query: `
+            subscription {
+              updatedUser {
+                id
+                name
+              }
             }
-          }
-        `
-      }
-    })
-  )
+          `
+        }
+      })
+    )
+  })
 
   {
     const [chunk] = await once(client, 'data')
     const data = JSON.parse(chunk)
     t.equal(data.type, 'connection_ack')
 
-    await gateway.inject({
-      method: 'POST',
-      url: '/graphql',
-      body: {
-        query: `
-          mutation {
-            triggerUser
-          }
-        `
-      }
+    process.nextTick(() => {
+      gateway.inject({
+        method: 'POST',
+        url: '/graphql',
+        body: {
+          query: `
+            mutation {
+              triggerUser
+            }
+          `
+        }
+      })
     })
   }
 
   {
     const [chunk] = await once(client, 'data')
     const data = JSON.parse(chunk)
+    client.end()
     t.equal(data.type, 'data')
     t.equal(data.id, 1)
 
@@ -910,58 +1126,59 @@ test('Polling schemas (subscriptions should be handled)', async (t) => {
     encoding: 'utf8',
     objectMode: true
   })
-
-  t.tearDown(() => {
-    client2.destroy()
-  })
-
+  t.tearDown(client2.destroy.bind(client2))
   client2.setEncoding('utf8')
 
-  client2.write(
-    JSON.stringify({
-      type: 'connection_init'
-    })
-  )
-
-  client2.write(
-    JSON.stringify({
-      id: 2,
-      type: 'start',
-      payload: {
-        query: `
-          subscription {
-            updatedUser {
-              id
-              name
-              lastName
+  process.nextTick(() => {
+    client2.write(
+      JSON.stringify({
+        type: 'connection_init'
+      })
+    )
+
+    client2.write(
+      JSON.stringify({
+        id: 2,
+        type: 'start',
+        payload: {
+          query: `
+            subscription {
+              updatedUser {
+                id
+                name
+                lastName
+              }
             }
-          }
-        `
-      }
-    })
-  )
+          `
+        }
+      })
+    )
+  })
 
   {
     const [chunk] = await once(client2, 'data')
     const data = JSON.parse(chunk)
     t.equal(data.type, 'connection_ack')
 
-    await gateway.inject({
-      method: 'POST',
-      url: '/graphql',
-      body: {
-        query: `
-          mutation {
-            triggerUser
-          }
-        `
-      }
+    process.nextTick(() => {
+      gateway.inject({
+        method: 'POST',
+        url: '/graphql',
+        body: {
+          query: `
+            mutation {
+              triggerUser
+            }
+          `
+        }
+      })
     })
   }
 
   {
     const [chunk] = await once(client2, 'data')
     const data = JSON.parse(chunk)
+    client2.end()
     t.equal(data.type, 'data')
     t.equal(data.id, 2)
 
@@ -976,207 +1193,6 @@ test('Polling schemas (subscriptions should be handled)', async (t) => {
 
   t.equal(ws2.readyState, WebSocket.OPEN)
 
-  clock.uninstall()
-})
-
-test('Polling schemas (should properly regenerate the schema when a downstream service restarts)', async (t) => {
-  const clock = FakeTimers.install({
-    shouldAdvanceTime: true,
-    advanceTimeDelta: 40
-  })
-  const oldSchema = `
-    directive @extends on INTERFACE | OBJECT
-
-    directive @external on FIELD_DEFINITION | OBJECT
-
-    directive @key(fields: String!) on INTERFACE | OBJECT
-
-    directive @provides(fields: String!) on FIELD_DEFINITION
-
-    directive @requires(fields: String!) on FIELD_DEFINITION
-
-    type Query {
-      me: User
-    }
-
-    type User @key(fields: "id") {
-      id: ID!
-      name: String!
-    }
-  `
-  const user = {
-    id: 'u1',
-    name: 'John',
-    lastName: 'Doe'
-  }
-
-  const userService = Fastify()
-  const gateway = Fastify()
-
-  userService.register(GQL, {
-    schema: oldSchema,
-    resolvers: {
-      Query: {
-        me: (root, args, context, info) => user
-      },
-      User: {
-        __resolveReference: (user, args, context, info) => user
-      }
-    },
-    federationMetadata: true
-  })
-
-  await userService.listen(0)
-
-  const userServicePort = userService.server.address().port
-
-  gateway.register(GQL, {
-    gateway: {
-      services: [
-        {
-          name: 'user',
-          url: `http://localhost:${userServicePort}/graphql`
-        }
-      ],
-      pollingInterval: 2000
-    }
-  })
-
-  await gateway.listen(0)
-
-  const res = await gateway.inject({
-    method: 'POST',
-    headers: {
-      'content-type': 'application/json'
-    },
-    url: '/graphql',
-    body: JSON.stringify({
-      query: `
-        query MainQuery {
-          me {
-            id
-            name
-          }
-        }
-      `
-    })
-  })
-
-  t.deepEqual(JSON.parse(res.body), {
-    data: {
-      me: {
-        id: 'u1',
-        name: 'John'
-      }
-    }
-  })
-
-  await userService.close()
-
-  const restartedUserService = Fastify()
-
-  const refreshedSchema = `
-    directive @extends on INTERFACE | OBJECT
-
-    directive @external on FIELD_DEFINITION | OBJECT
-
-    directive @key(fields: String!) on INTERFACE | OBJECT
-
-    directive @provides(fields: String!) on FIELD_DEFINITION
-    
-    directive @requires(fields: String!) on FIELD_DEFINITION
-
-    type User @key(fields: "id") {
-      id: ID!
-      lastName: String!
-      name: String!
-    }
-
-    type Mutation {
-      create: User!
-    }
-
-    type Query {
-      me2: User
-    }
-  `
-
-  restartedUserService.register(GQL, {
-    schema: refreshedSchema,
-    resolvers: {
-      Query: {
-        me2: (root, args, context, info) => user
-      },
-      Mutation: {
-        create: (root, args, context, info) => user
-      },
-      User: {
-        __resolveReference: (user, args, context, info) => user
-      }
-    },
-    federationMetadata: true
-  })
-
-  await restartedUserService.listen(userServicePort)
-
-  await clock.tickAsync(4000)
-
-  const res2 = await gateway.inject({
-    method: 'POST',
-    headers: {
-      'content-type': 'application/json'
-    },
-    url: '/graphql',
-    body: JSON.stringify({
-      query: `
-        query MainQuery {
-          me {
-            id
-            name
-          }
-        }
-      `
-    })
-  })
-
-  t.deepEqual(JSON.parse(res2.body), {
-    errors: [
-      {
-        message: 'Cannot query field "me" on type "Query". Did you mean "me2"?',
-        locations: [{ line: 3, column: 11 }]
-      }
-    ],
-    data: null
-  })
-
-  const res3 = await gateway.inject({
-    method: 'POST',
-    headers: {
-      'content-type': 'application/json'
-    },
-    url: '/graphql',
-    body: JSON.stringify({
-      query: `
-        mutation NewMutation {
-          create {
-            id
-            name
-          }
-        }
-      `
-    })
-  })
-
-  t.deepEqual(JSON.parse(res3.body), {
-    data: {
-      create: {
-        id: 'u1',
-        name: 'John'
-      }
-    }
-  })
-
   await gateway.close()
-  await restartedUserService.close()
-  clock.uninstall()
+  await userService.close()
 })
diff --git a/test/gateway/request.js b/test/gateway/request.js
index 6838ce1..b5c55e4 100644
--- a/test/gateway/request.js
+++ b/test/gateway/request.js
@@ -32,7 +32,7 @@ test('sendRequest method rejects when response is not valid json', async (t) =>
   const { request, close } = buildRequest({ url })
   t.tearDown(() => {
     close()
-    app.close()
+    return app.close()
   })
   t.rejects(sendRequest(request, url)({
     method: 'POST',
@@ -62,7 +62,7 @@ test('sendRequest method rejects when response contains errors', async (t) => {
   const { request, close } = buildRequest({ url })
   t.tearDown(() => {
     close()
-    app.close()
+    return app.close()
   })
   t.rejects(
     sendRequest(
diff --git a/test/gateway/requires-directive.js b/test/gateway/requires-directive.js
index 3a3a2fb..a9d5690 100644
--- a/test/gateway/requires-directive.js
+++ b/test/gateway/requires-directive.js
@@ -6,9 +6,6 @@ const GQL = require('../..')
 
 async function createService (t, schema, resolvers = {}) {
   const service = Fastify()
-  t.tearDown(() => {
-    service.close()
-  })
   service.register(GQL, {
     schema,
     resolvers,
@@ -16,7 +13,7 @@ async function createService (t, schema, resolvers = {}) {
   })
   await service.listen(0)
 
-  return service.server.address().port
+  return [service, service.server.address().port]
 }
 
 const users = {
@@ -52,7 +49,7 @@ const posts = {
 }
 
 test('gateway handles @requires directive correctly', async (t) => {
-  const userServicePort = await createService(t, `
+  const [userService, userServicePort] = await createService(t, `
     extend type Query {
       me: User
     }
@@ -86,7 +83,7 @@ test('gateway handles @requires directive correctly', async (t) => {
     }
   })
 
-  const biographyServicePort = await createService(t, `
+  const [biographyService, biographyServicePort] = await createService(t, `
     type User @key(fields: "id") @extends {
       id: ID! @external
       name: String @external
@@ -102,8 +99,10 @@ test('gateway handles @requires directive correctly', async (t) => {
   })
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await biographyService.close()
+    await userService.close()
   })
   gateway.register(GQL, {
     gateway: {
diff --git a/test/gateway/schema.js b/test/gateway/schema.js
index dd1aaf4..517d3b8 100644
--- a/test/gateway/schema.js
+++ b/test/gateway/schema.js
@@ -6,9 +6,6 @@ const GQL = require('../..')
 
 async function createService (t, schema, resolvers = {}) {
   const service = Fastify()
-  t.tearDown(() => {
-    service.close()
-  })
   service.register(GQL, {
     schema,
     resolvers,
@@ -16,7 +13,7 @@ async function createService (t, schema, resolvers = {}) {
   })
   await service.listen(0)
 
-  return service.server.address().port
+  return [service, service.server.address().port]
 }
 
 test('It builds the gateway schema correctly', async (t) => {
@@ -62,7 +59,7 @@ test('It builds the gateway schema correctly', async (t) => {
     }
   }
 
-  const userServicePort = await createService(t, `
+  const [userService, userServicePort] = await createService(t, `
     directive @customDirective on FIELD_DEFINITION
 
     extend type Query {
@@ -98,7 +95,7 @@ test('It builds the gateway schema correctly', async (t) => {
     }
   })
 
-  const postServicePort = await createService(t, `
+  const [postService, postServicePort] = await createService(t, `
     type Post @key(fields: "pid") {
       pid: ID!
       title: String
@@ -142,8 +139,10 @@ test('It builds the gateway schema correctly', async (t) => {
   })
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await postService.close()
+    await userService.close()
   })
   gateway.register(GQL, {
     gateway: {
@@ -312,7 +311,7 @@ test('It support variable inside nested arguments', async (t) => {
     }
   }
 
-  const userServicePort = await createService(t, `
+  const [userService, userServicePort] = await createService(t, `
     directive @customDirective on FIELD_DEFINITION
 
     extend type Query {
@@ -347,8 +346,9 @@ test('It support variable inside nested arguments', async (t) => {
   })
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await userService.close()
   })
   gateway.register(GQL, {
     gateway: {
@@ -427,7 +427,7 @@ test('Should not throw on nullable reference', async (t) => {
     }
   ]
 
-  const postServicePort = await createService(t, `
+  const [postService, postServicePort] = await createService(t, `
     extend type Query {
       topPosts: [Post]
     }
@@ -457,7 +457,7 @@ test('Should not throw on nullable reference', async (t) => {
     }
   })
 
-  const userServicePort = await createService(t, `
+  const [userService, userServicePort] = await createService(t, `
     type User @key(fields: "id") {
       id: ID!
       name: String
@@ -473,8 +473,10 @@ test('Should not throw on nullable reference', async (t) => {
   })
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await postService.close()
+    await userService.close()
   })
   gateway.register(GQL, {
     gateway: {
@@ -550,7 +552,7 @@ test('Should handle InlineFragment', async (t) => {
     }
   ]
 
-  const productServicePort = await createService(t, `
+  const [productService, productServicePort] = await createService(t, `
     extend type Query {
       products: [Product]
     }
@@ -582,8 +584,9 @@ test('Should handle InlineFragment', async (t) => {
   })
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await productService.close()
   })
   gateway.register(GQL, {
     gateway: {
@@ -665,7 +668,7 @@ test('Should support array references with _entities query', async (t) => {
     }
   ]
 
-  const postServicePort = await createService(t, `
+  const [postService, postServicePort] = await createService(t, `
     extend type Query {
       topPosts: [Post]
     }
@@ -695,7 +698,7 @@ test('Should support array references with _entities query', async (t) => {
     }
   })
 
-  const userServicePort = await createService(t, `
+  const [userService, userServicePort] = await createService(t, `
     type User @key(fields: "id") {
       id: ID!
       name: String
@@ -711,8 +714,10 @@ test('Should support array references with _entities query', async (t) => {
   })
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await postService.close()
+    await userService.close()
   })
   gateway.register(GQL, {
     gateway: {
@@ -786,7 +791,7 @@ test('Should support array references with _entities query', async (t) => {
 })
 
 test('Should support multiple `extends` of the same type in the service SDL', async (t) => {
-  const productServicePort = await createService(t, `
+  const [productService, productServicePort] = await createService(t, `
     extend type Query {
       ping: Int
     }
@@ -805,8 +810,9 @@ test('Should support multiple `extends` of the same type in the service SDL', as
   })
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await productService.close()
   })
   gateway.register(GQL, {
     gateway: {
@@ -885,7 +891,7 @@ test('Should support array references with _entities query and empty response',
     }
   ]
 
-  const postServicePort = await createService(t, `
+  const [postService, postServicePort] = await createService(t, `
     extend type Query {
       topPosts: [Post]
     }
@@ -915,7 +921,7 @@ test('Should support array references with _entities query and empty response',
     }
   })
 
-  const userServicePort = await createService(t, `
+  const [userService, userServicePort] = await createService(t, `
     type User @key(fields: "id") {
       id: ID!
       name: String
@@ -931,8 +937,10 @@ test('Should support array references with _entities query and empty response',
   })
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await postService.close()
+    await userService.close()
   })
   gateway.register(GQL, {
     gateway: {
@@ -1022,7 +1030,7 @@ test('Should support array references with _entities query and empty response an
     }
   ]
 
-  const postServicePort = await createService(t, `
+  const [postService, postServicePort] = await createService(t, `
     extend type Query {
       topPosts: [Post]
     }
@@ -1052,7 +1060,7 @@ test('Should support array references with _entities query and empty response an
     }
   })
 
-  const userServicePort = await createService(t, `
+  const [userService, userServicePort] = await createService(t, `
     type User @key(fields: "id") {
       id: ID!
       name: String
@@ -1068,8 +1076,10 @@ test('Should support array references with _entities query and empty response an
   })
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await postService.close()
+    await userService.close()
   })
   gateway.register(GQL, {
     gateway: {
@@ -1142,7 +1152,7 @@ test('Should handle union with InlineFragment', async (t) => {
     }
   ]
 
-  const productServicePort = await createService(t, `
+  const [productService, productServicePort] = await createService(t, `
     extend type Query {
       products: [Product]
       shelve: Shelve
@@ -1180,8 +1190,9 @@ test('Should handle union with InlineFragment', async (t) => {
   })
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await productService.close()
   })
   gateway.register(GQL, {
     gateway: {
@@ -1245,9 +1256,6 @@ test('Should handle union with InlineFragment', async (t) => {
 test('Gateway sends initHeaders with _service sdl query', async (t) => {
   t.plan(1)
   const service = Fastify()
-  t.tearDown(() => {
-    service.close()
-  })
   service.register(GQL, {
     schema: `
       extend type Query {
@@ -1271,8 +1279,9 @@ test('Gateway sends initHeaders with _service sdl query', async (t) => {
   await service.listen(0)
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await service.close()
   })
   gateway.register(GQL, {
     gateway: {
@@ -1294,9 +1303,6 @@ test('Gateway sends initHeaders with _service sdl query', async (t) => {
 test('Gateway sends initHeaders function result with _service sdl query', async (t) => {
   t.plan(1)
   const service = Fastify()
-  t.tearDown(() => {
-    service.close()
-  })
   service.register(GQL, {
     schema: `
       extend type Query {
@@ -1320,8 +1326,9 @@ test('Gateway sends initHeaders function result with _service sdl query', async
   await service.listen(0)
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await service.close()
   })
   gateway.register(GQL, {
     gateway: {
@@ -1356,7 +1363,7 @@ test('Should handle interface', async (t) => {
     }
   ]
 
-  const productServicePort = await createService(t, `
+  const [productService, productServicePort] = await createService(t, `
     extend type Query {
       products: [Product]
       shelve: Shelve
@@ -1400,8 +1407,9 @@ test('Should handle interface', async (t) => {
   })
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await productService.close()
   })
   gateway.register(GQL, {
     gateway: {
@@ -1493,7 +1501,7 @@ test('Should handle interface referenced multiple times in different services',
     }
   ]
 
-  const bookServicePort = await createService(t, `
+  const [bookService, bookServicePort] = await createService(t, `
     extend type Query {
       books: [Book]
     }
@@ -1525,7 +1533,7 @@ test('Should handle interface referenced multiple times in different services',
       }
     }
   })
-  const dictionariesServicePort = await createService(t, `
+  const [dictionariesService, dictionariesServicePort] = await createService(t, `
     extend type Query {
       dictionaries: [Dictionary]
     }
@@ -1559,8 +1567,10 @@ test('Should handle interface referenced multiple times in different services',
   })
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await dictionariesService.close()
+    await bookService.close()
   })
   gateway.register(GQL, {
     gateway: {
@@ -1736,7 +1746,7 @@ test('Should handle complex and nested interfaces with external types', async (t
     }
   `
 
-  const userServicePort = await createService(t, `
+  const [userService, userServicePort] = await createService(t, `
     type User @key(fields: "id") {
       id: ID!
       name: String!
@@ -1757,7 +1767,7 @@ test('Should handle complex and nested interfaces with external types', async (t
       }
     }
   })
-  const configABServicePort = await createService(t, `
+  const [configABService, configABServicePort] = await createService(t, `
     ${configInterface}
     type ConfigA implements ConfigInterface {
       type: EConfig!
@@ -1801,7 +1811,7 @@ test('Should handle complex and nested interfaces with external types', async (t
       }
     }
   })
-  const configCServicePort = await createService(t, `
+  const [configCService, configCServicePort] = await createService(t, `
     ${configInterface}
     type ConfigC implements ConfigInterface {
       type: EConfig!
@@ -1844,8 +1854,11 @@ test('Should handle complex and nested interfaces with external types', async (t
   })
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await configCService.close()
+    await configABService.close()
+    await userService.close()
   })
   gateway.register(GQL, {
     gateway: {
@@ -2013,7 +2026,7 @@ test('Uses the supplied schema for federation rather than fetching it remotely',
     }
   }
 
-  const userServicePort = await createService(t, `
+  const [userService, userServicePort] = await createService(t, `
     directive @customDirective on FIELD_DEFINITION
 
     extend type Query {
@@ -2070,7 +2083,7 @@ test('Uses the supplied schema for federation rather than fetching it remotely',
     }
   `
 
-  const postServicePort = await createService(t, postServiceSdl, {
+  const [postService, postServicePort] = await createService(t, postServiceSdl, {
     Post: {
       __resolveReference: (post, args, context, info) => {
         return posts[post.pid]
@@ -2097,8 +2110,10 @@ test('Uses the supplied schema for federation rather than fetching it remotely',
   })
 
   const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
+  t.tearDown(async () => {
+    await gateway.close()
+    await postService.close()
+    await userService.close()
   })
   gateway.register(GQL, {
     gateway: {
@@ -2252,11 +2267,7 @@ test('Uses the supplied schema for federation rather than fetching it remotely',
 })
 
 test('Non mandatory gateway failure wont stop gateway creation', async (t) => {
-  const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
-  })
-  const brokenServicePort = await createService(t, `
+  const [brokenService, brokenServicePort] = await createService(t, `
     extend type Query {
       _service: String
     }
@@ -2268,7 +2279,7 @@ test('Non mandatory gateway failure wont stop gateway creation', async (t) => {
     }
   })
 
-  const workingServicePort = await createService(t, `
+  const [workingService, workingServicePort] = await createService(t, `
     extend type Query {
       hello: String!
     }
@@ -2278,6 +2289,12 @@ test('Non mandatory gateway failure wont stop gateway creation', async (t) => {
     }
   })
 
+  const gateway = Fastify()
+  t.tearDown(async () => {
+    await gateway.close()
+    await brokenService.close()
+    await workingService.close()
+  })
   gateway.register(GQL, {
     gateway: {
       services: [{
@@ -2311,11 +2328,6 @@ test('Non mandatory gateway failure wont stop gateway creation', async (t) => {
 })
 
 test('Update the schema', async (t) => {
-  const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
-  })
-
   const schema = `
     extend type Query {
       hello: String!
@@ -2329,13 +2341,18 @@ test('Update the schema', async (t) => {
     }
   `
 
-  const servicePort = await createService(t, fullSchema, {
+  const [service, servicePort] = await createService(t, fullSchema, {
     Query: {
       hello: () => 'world',
       world: () => 'hello'
     }
   })
 
+  const gateway = Fastify()
+  t.tearDown(async () => {
+    await gateway.close()
+    await service.close()
+  })
   gateway.register(GQL, {
     gateway: {
       services: [{
@@ -2388,23 +2405,23 @@ test('Update the schema', async (t) => {
 })
 
 test('Update the schema without any changes', async (t) => {
-  const gateway = Fastify()
-  t.tearDown(() => {
-    gateway.close()
-  })
-
   const schema = `
     extend type Query {
       hello: String!
     }
   `
 
-  const servicePort = await createService(t, schema, {
+  const [service, servicePort] = await createService(t, schema, {
     Query: {
       hello: () => 'world'
     }
   })
 
+  const gateway = Fastify()
+  t.tearDown(async () => {
+    await gateway.close()
+    await service.close()
+  })
   gateway.register(GQL, {
     gateway: {
       services: [{
diff --git a/test/gateway/subscription.js b/test/gateway/subscription.js
index 0371601..89f72b7 100644
--- a/test/gateway/subscription.js
+++ b/test/gateway/subscription.js
@@ -176,11 +176,11 @@ test('gateway subscription handling works correctly', t => {
   function runSubscription () {
     const ws = new WebSocket(`ws://localhost:${(gateway.server.address()).port}/graphql`, 'graphql-ws')
     const client = WebSocket.createWebSocketStream(ws, { encoding: 'utf8', objectMode: true })
-    t.tearDown(() => {
+    t.tearDown(async () => {
       client.destroy()
-      messageService.close()
-      userService.close()
-      gateway.close()
+      await gateway.close()
+      await messageService.close()
+      await userService.close()
     })
     client.setEncoding('utf8')
 
@@ -295,9 +295,6 @@ test('gateway wsConnectionParams object is passed to SubscriptionClient', t => {
     hello: 'world'
   }
   const testService = Fastify()
-  t.tearDown(() => {
-    testService.close()
-  })
 
   testService.register(GQL, {
     schema: `
@@ -314,8 +311,9 @@ test('gateway wsConnectionParams object is passed to SubscriptionClient', t => {
     const testServicePort = testService.server.address().port
 
     const gateway = Fastify()
-    t.tearDown(() => {
-      gateway.close()
+    t.tearDown(async () => {
+      await gateway.close()
+      await testService.close()
     })
     gateway.register(GQL, {
       subscription: true,
@@ -344,9 +342,6 @@ test('gateway wsConnectionParams function is passed to SubscriptionClient', t =>
     hello: 'world'
   }
   const testService = Fastify()
-  t.tearDown(() => {
-    testService.close()
-  })
 
   testService.register(GQL, {
     schema: `
@@ -363,8 +358,9 @@ test('gateway wsConnectionParams function is passed to SubscriptionClient', t =>
     const testServicePort = testService.server.address().port
 
     const gateway = Fastify()
-    t.tearDown(() => {
-      gateway.close()
+    t.tearDown(async () => {
+      await gateway.close()
+      await testService.close()
     })
     gateway.register(GQL, {
       subscription: true,
@@ -385,7 +381,7 @@ test('gateway wsConnectionParams function is passed to SubscriptionClient', t =>
   })
 })
 
-test('gateway forwards the connectin_init payload to the federated service on gql_start using the connectionInit extension', t => {
+test('gateway forwards the connection_init payload to the federated service on gql_start using the connectionInit extension', t => {
   t.plan(3)
   function onConnect (data) {
     if (data && data.payload && Object.entries(data.payload).length) {
@@ -399,9 +395,6 @@ test('gateway forwards the connectin_init payload to the federated service on gq
     hello: 'world'
   }
   const testService = Fastify()
-  t.tearDown(() => {
-    testService.close()
-  })
 
   testService.register(GQL, {
     schema: `
@@ -440,7 +433,10 @@ test('gateway forwards the connectin_init payload to the federated service on gq
     const testServicePort = testService.server.address().port
 
     const gateway = Fastify()
-    t.tearDown(() => { gateway.close() })
+    t.tearDown(async () => {
+      await gateway.close()
+      await testService.close()
+    })
     gateway.register(GQL, {
       subscription: true,
       gateway: {
@@ -456,9 +452,7 @@ test('gateway forwards the connectin_init payload to the federated service on gq
       t.error(err)
       const ws = new WebSocket(`ws://localhost:${(gateway.server.address()).port}/graphql`, 'graphql-ws')
       const client = WebSocket.createWebSocketStream(ws, { encoding: 'utf8', objectMode: true })
-      t.tearDown(() => {
-        client.destroy()
-      })
+      t.tearDown(client.destroy.bind(client))
       client.setEncoding('utf8')
 
       client.write(JSON.stringify({
diff --git a/test/subscription-connection.js b/test/subscription-connection.js
index 2c493b3..01a5811 100644
--- a/test/subscription-connection.js
+++ b/test/subscription-connection.js
@@ -92,7 +92,7 @@ test('subscription connection handles GQL_CONNECTION_TERMINATE message correctly
   }))
 })
 
-test('subscription connection handles GQL_STOP message correctly', async (t) => {
+test('subscription connection closes context on GQL_STOP message correctly', async (t) => {
   t.plan(2)
   const sc = new SubscriptionConnection({
     on () {},
@@ -115,6 +115,29 @@ test('subscription connection handles GQL_STOP message correctly', async (t) =>
   t.equal(sc.subscriptionContexts.size, 0)
 })
 
+test('subscription connection completes resolver iterator on GQL_STOP message correctly', async (t) => {
+  t.plan(2)
+  const sc = new SubscriptionConnection({
+    on () {},
+    close () {},
+    send (message) {}
+  }, {})
+
+  sc.subscriptionIters = new Map()
+  sc.subscriptionIters.set(1, {
+    return () {
+      t.pass()
+    }
+  })
+
+  await sc.handleMessage(JSON.stringify({
+    id: 1,
+    type: 'stop'
+  }))
+
+  t.equal(sc.subscriptionIters.size, 0)
+})
+
 test('handles error in send and closes connection', async t => {
   t.plan(1)
 
diff --git a/test/subscription.js b/test/subscription.js
index e30436c..9ce7a1f 100644
--- a/test/subscription.js
+++ b/test/subscription.js
@@ -621,6 +621,7 @@ test('subscription socket protocol different than graphql-ws, protocol = foobar'
     const url = 'ws://localhost:' + (app.server.address()).port + '/graphql'
     const ws = new WebSocket(url, 'foobar')
     const client = WebSocket.createWebSocketStream(ws, { encoding: 'utf8', objectMode: true })
+    t.tearDown(client.destroy.bind(client))
     client.setEncoding('utf8')
     ws.on('close', () => {
       client.end()
@@ -661,6 +662,7 @@ test('subscription connection is closed if context function throws', t => {
     const url = 'ws://localhost:' + (app.server.address()).port + '/graphql'
     const ws = new WebSocket(url, 'graphql-ws')
     const client = WebSocket.createWebSocketStream(ws, { encoding: 'utf8', objectMode: true })
+    t.tearDown(client.destroy.bind(client))
 
     client.write(JSON.stringify({
       type: 'connection_init'
@@ -683,8 +685,8 @@ test('subscription server sends update to subscriptions with custom async contex
   })
 
   const app = Fastify()
-  t.tearDown(() => {
-    app.close()
+  t.tearDown(async () => {
+    await app.close()
     clock.uninstall()
   })
 
@@ -870,8 +872,8 @@ test('subscription connection is closed if async context function throws', t =>
   })
 
   const app = Fastify()
-  t.tearDown(() => {
-    app.close()
+  t.tearDown(async () => {
+    await app.close()
     clock.uninstall()
   })
 
@@ -1825,27 +1827,6 @@ test('subscription server works with fastify-websocket', t => {
     t.tearDown(subscriptionClient.destroy.bind(subscriptionClient))
     subscriptionClient.setEncoding('utf8')
 
-    client.write('hi from client')
-
-    subscriptionClient.write(JSON.stringify({
-      type: 'connection_init'
-    }))
-
-    subscriptionClient.write(JSON.stringify({
-      id: 1,
-      type: 'start',
-      payload: {
-        query: `
-          subscription {
-            notificationAdded {
-              id
-              message
-            }
-          }
-        `
-      }
-    }))
-
     client.on('data', chunk => {
       t.equal(chunk, 'hi from server')
       client.end()
@@ -1868,10 +1849,30 @@ test('subscription server works with fastify-websocket', t => {
           }
         }))
         subscriptionClient.end()
-        t.end()
       } else {
         sendTestMutation()
       }
     })
+
+    client.write('hi from client')
+
+    subscriptionClient.write(JSON.stringify({
+      type: 'connection_init'
+    }))
+
+    subscriptionClient.write(JSON.stringify({
+      id: 1,
+      type: 'start',
+      payload: {
+        query: `
+          subscription {
+            notificationAdded {
+              id
+              message
+            }
+          }
+        `
+      }
+    }))
   })
 })