Repository: mercurius-js/mercurius
Pull request #638 (2021-11-09 12:09:04)

Fix incorrect message type for GQL_STOP when using graphql-transport-ws
Closes https://github.com/mercurius-js/mercurius/issues/629

Let me know if the unit-tests should be structured differently. Or if you want me to extend the PR and add tests for all the message types of `graphql-transport-ws`.


Patch:
diff --git a/README.md b/README.md
index 514e5a4..305a1bd 100644
--- a/README.md
+++ b/README.md
@@ -34,6 +34,7 @@ Features:
 - [Integrations](docs/integrations/)
 - [Related Plugins](docs/plugins.md)
 - [Protocol Extensions](/docs/protocol-extension.md)
+- [Faq](/docs/faq.md)
 - [Acknowledgements](#acknowledgements)
 - [License](#license)
 
diff --git a/docs/loaders.md b/docs/loaders.md
index c390e44..6834196 100644
--- a/docs/loaders.md
+++ b/docs/loaders.md
@@ -9,10 +9,12 @@ also cache the results, so that other parts of the GraphQL do not have
 to fetch the same data.
 
 Each loader function has the signature `loader(queries, context)`.
-`queries` is an array of objects defined as `{ obj, params }` where
-`obj` is the current object and `params` are the GraphQL params (those
-are the first two parameters of a normal resolver). The `context` is the
-GraphQL context, and it includes a `reply` object.
+`queries` is an array of objects defined as `{ obj, params, info }` where
+`obj` is the current object, `params` are the GraphQL params (those
+are the first two parameters of a normal resolver) and `info` contains
+additional information about the query and execution. `info` object is
+only available in the loader if the cache is set to `false`. The `context`
+is the GraphQL context, and it includes a `reply` object.
 
 Example:
 
diff --git a/docsify/sidebar.md b/docsify/sidebar.md
index 0df840a..fb0ca99 100644
--- a/docsify/sidebar.md
+++ b/docsify/sidebar.md
@@ -30,3 +30,4 @@
   * [mercurius-upload](/docs/plugins#mercurius-upload)
   * [altair-fastify-plugin](/docs/plugins#altair-fastify-plugin)
 * [Protocol Extensions](/docs/protocol-extension)
+* [Faq](/docs/faq)
diff --git a/index.js b/index.js
index e88b0c7..340add2 100644
--- a/index.js
+++ b/index.js
@@ -258,7 +258,7 @@ const plugin = fp(async function (app, opts) {
 
     context = Object.assign(context, { reply: this, app })
     if (app[kFactory]) {
-      this[kLoaders] = factory.create(context)
+      this[kLoaders] = app[kFactory].create(context)
     }
 
     return app.graphql(source, context, variables, operationName)
@@ -368,12 +368,14 @@ const plugin = fp(async function (app, opts) {
 
     function defineLoader (name) {
       // async needed because of throw
-      return async function (obj, params, { reply }) {
+      return async function (obj, params, { reply }, info) {
         if (!reply) {
           throw new MER_ERR_INVALID_OPTS('loaders only work via reply.graphql()')
         }
 
-        return reply[kLoaders][name]({ obj, params })
+        const query = opts.cache === false ? { obj, params, info } : { obj, params }
+
+        return reply[kLoaders][name](query)
       }
     }
 
diff --git a/lib/gateway/make-resolver.js b/lib/gateway/make-resolver.js
index ef25c72..118cf60 100644
--- a/lib/gateway/make-resolver.js
+++ b/lib/gateway/make-resolver.js
@@ -506,7 +506,9 @@ function makeResolver ({ service, createOperation, transformData, isQuery, isRef
       return transformData(response)
     }
 
-    const response = await reply[kEntityResolvers][`${service.name}Entity`]({
+    const entityResolvers = reply.entityResolversFactory ? reply.entityResolversFactory.create() : reply[kEntityResolvers]
+
+    const response = await entityResolvers[`${service.name}Entity`]({
       query,
       variables,
       originalRequestHeaders: reply.request.headers,
diff --git a/lib/subscription-connection.js b/lib/subscription-connection.js
index 7305b48..6911cef 100644
--- a/lib/subscription-connection.js
+++ b/lib/subscription-connection.js
@@ -1,7 +1,7 @@
 'use strict'
 
 const on = require('events.on')
-const { subscribe, parse } = require('graphql')
+const { subscribe, parse, print } = require('graphql')
 const { SubscriptionContext } = require('./subscriber')
 const { kEntityResolvers } = require('./gateway/make-resolver')
 const sJSON = require('secure-json-parse')
@@ -15,7 +15,7 @@ module.exports = class SubscriptionConnection {
     subscriber,
     fastify,
     lruGatewayResolvers,
-    entityResolvers,
+    entityResolversFactory,
     context = {},
     onConnect,
     onDisconnect,
@@ -24,7 +24,7 @@ module.exports = class SubscriptionConnection {
     this.fastify = fastify
     this.socket = socket
     this.lruGatewayResolvers = lruGatewayResolvers
-    this.entityResolvers = entityResolvers
+    this.entityResolversFactory = entityResolversFactory
     this.subscriber = subscriber
     this.onConnect = onConnect
     this.onDisconnect = onDisconnect
@@ -33,6 +33,8 @@ module.exports = class SubscriptionConnection {
     this.context = context
     this.isReady = false
     this.resolveContext = resolveContext
+    this.headers = {}
+
     this.protocolMessageTypes = getProtocolByName(socket.protocol)
     this.socket.on('error', this.handleConnectionClose.bind(this))
     this.handleConnection()
@@ -123,6 +125,10 @@ module.exports = class SubscriptionConnection {
 
     this.context._connectionInit = data.payload
 
+    if (data.payload && data.payload.headers) {
+      this.headers = data.payload.headers
+    }
+
     this.sendMessage(this.protocolMessageTypes.GQL_CONNECTION_ACK)
     this.isReady = true
   }
@@ -183,12 +189,16 @@ module.exports = class SubscriptionConnection {
       {}, // rootValue
       {
         ...context,
+        get __currentQuery () {
+          return print(document)
+        },
         pubsub: sc,
         lruGatewayResolvers: this.lruGatewayResolvers,
         reply: {
+          entityResolversFactory: this.entityResolversFactory,
+          request: { headers: this.headers },
           [kLoaders]: subscriptionLoaders && subscriptionLoaders.create(context),
-          [kEntityResolvers]: this.entityResolvers,
-          request: { headers: {} }
+          [kEntityResolvers]: this.entityResolvers
         }
       },
       variables,
diff --git a/lib/subscription-protocol.js b/lib/subscription-protocol.js
index a480555..79654b9 100644
--- a/lib/subscription-protocol.js
+++ b/lib/subscription-protocol.js
@@ -1,11 +1,14 @@
 'use strict'
 
 const GRAPHQL_WS = 'graphql-ws'
+const GRAPHQL_TRANSPORT_WS = 'graphql-transport-ws'
+
 module.exports.GRAPHQL_WS = GRAPHQL_WS
+module.exports.GRAPHQL_TRANSPORT_WS = GRAPHQL_TRANSPORT_WS
 
 module.exports.getProtocolByName = function (name) {
   switch (true) {
-    case (name.indexOf('graphql-transport-ws') !== -1):
+    case (name.indexOf(GRAPHQL_TRANSPORT_WS) !== -1):
       return {
         GQL_CONNECTION_INIT: 'connection_init', // Client -> Server
         GQL_CONNECTION_ACK: 'connection_ack', // Server -> Client
@@ -16,7 +19,7 @@ module.exports.getProtocolByName = function (name) {
         GQL_DATA: 'next', // Server -> Client
         GQL_ERROR: 'error', // Server -> Client
         GQL_COMPLETE: 'complete', // Server -> Client
-        GQL_STOP: 'stop' // Client -> Server
+        GQL_STOP: 'complete' // Client -> Server
       }
     case (name.indexOf(GRAPHQL_WS) !== -1):
       return {
diff --git a/lib/subscription.js b/lib/subscription.js
index b0026e9..009cc28 100644
--- a/lib/subscription.js
+++ b/lib/subscription.js
@@ -42,7 +42,7 @@ function createConnectionHandler ({ subscriber, fastify, onConnect, onDisconnect
       onConnect,
       onDisconnect,
       lruGatewayResolvers,
-      entityResolvers: entityResolversFactory && entityResolversFactory.create(),
+      entityResolversFactory,
       context,
       resolveContext
     })


Test patch:
diff --git a/tap-snapshots/test/reply-decorator.js.test.cjs b/tap-snapshots/test/reply-decorator.js.test.cjs
index f47f10d..6a458c6 100644
--- a/tap-snapshots/test/reply-decorator.js.test.cjs
+++ b/tap-snapshots/test/reply-decorator.js.test.cjs
@@ -8,3 +8,7 @@
 exports['test/reply-decorator.js TAP reply decorator set status code to 400 with bad query > must match snapshot 1'] = `
 {"errors":[{"message":"Syntax Error: Expected Name, found <EOF>.","locations":[{"line":1,"column":18}]}]}
 `
+
+exports['test/reply-decorator.js TAP reply decorator supports encapsulation when loaders are defined in parent object > must match snapshot 1'] = `
+{"data":{"multiply":25}}
+`
diff --git a/test/gateway/subscription.js b/test/gateway/subscription.js
index 97e79ae..4355806 100644
--- a/test/gateway/subscription.js
+++ b/test/gateway/subscription.js
@@ -760,3 +760,427 @@ test('subscriptions work with scalars', async t => {
   await createGateway()
   await runSubscription()
 })
+
+test('subscriptions work with different contexts', async (t) => {
+  let testService
+  let gateway
+
+  const schema = `
+  extend type Query {
+      ignored: Boolean!
+  }
+
+  extend type Mutation {
+      addTestEvent(value: Int!): Int!
+  }
+
+  type Event @key(fields: "value") {
+    value: Int! @external
+  }
+
+  extend type Subscription {
+      testEvent(value: Int!): Int!
+  }`
+
+  const resolvers = {
+    Query: {
+      ignored: () => true
+    },
+    Mutation: {
+      addTestEvent: async (_, { value }, { pubsub }) => {
+        await pubsub.publish({
+          topic: 'testEvent',
+          payload: { testEvent: value }
+        })
+
+        return value
+      }
+    },
+    Subscription: {
+      testEvent: {
+        subscribe: GQL.withFilter(
+          async (_, __, { pubsub }) => {
+            return await pubsub.subscribe('testEvent')
+          },
+          ({ testEvent }, { value }) => {
+            return testEvent === value
+          }
+        )
+      }
+    }
+  }
+
+  function createTestService () {
+    testService = Fastify()
+    testService.register(GQL, {
+      schema,
+      resolvers,
+      federationMetadata: true,
+      subscription: true
+    })
+
+    return testService.listen(0)
+  }
+
+  function createGateway () {
+    const testServicePort = testService.server.address().port
+
+    gateway = Fastify()
+    gateway.register(GQL, {
+      subscription: true,
+      gateway: {
+        services: [{
+          name: 'testService',
+          url: `http://localhost:${testServicePort}/graphql`,
+          wsUrl: `ws://localhost:${testServicePort}/graphql`
+        }]
+      }
+    })
+
+    return gateway.listen(0)
+  }
+
+  function runSubscription (id) {
+    const ws = new WebSocket(`ws://localhost:${(gateway.server.address()).port}/graphql`, 'graphql-ws')
+    const client = WebSocket.createWebSocketStream(ws, { encoding: 'utf8', objectMode: true })
+    t.teardown(async () => {
+      client.destroy()
+    })
+    client.setEncoding('utf8')
+
+    client.write(JSON.stringify({
+      type: 'connection_init'
+    }))
+
+    client.write(JSON.stringify({
+      id: 1,
+      type: 'start',
+      payload: {
+        query: `
+          subscription TestEvent($value: Int!) {
+            testEvent(value: $value)
+          }
+        `,
+        variables: { value: id }
+      }
+    }))
+
+    client.write(JSON.stringify({
+      id: 2,
+      type: 'start',
+      payload: {
+        query: `
+          subscription TestEvent($value: Int!) {
+            testEvent(value: $value)
+          }
+        `,
+        variables: { value: id }
+      }
+    }))
+
+    client.write(JSON.stringify({
+      id: 2,
+      type: 'stop'
+    }))
+
+    let end
+
+    const endPromise = new Promise(resolve => {
+      end = resolve
+    })
+
+    client.on('data', (chunk) => {
+      const data = JSON.parse(chunk)
+
+      if (data.id === 1 && data.type === 'data') {
+        t.equal(chunk, JSON.stringify({
+          type: 'data',
+          id: 1,
+          payload: {
+            data: {
+              testEvent: id
+            }
+          }
+        }))
+
+        client.end()
+        end()
+      } else if (data.id === 2 && data.type === 'complete') {
+        gateway.inject({
+          method: 'POST',
+          url: '/graphql',
+          body: {
+            query: `
+              mutation AddTestEvent($value: Int!) {
+                addTestEvent(value: $value)
+              }
+            `,
+            variables: { value: id }
+          }
+        })
+      }
+    })
+
+    return endPromise
+  }
+
+  await createTestService()
+  await createGateway()
+  const subscriptions = new Array(10).fill(null).map((_, i) => runSubscription(i))
+  await Promise.all(subscriptions)
+
+  t.teardown(async () => {
+    await gateway.close()
+    await testService.close()
+  })
+})
+
+test('connection_init headers available in federation event resolver', async (t) => {
+  let subscriptionService
+  let resolverService
+  let gateway
+
+  const onConnect = (data) => {
+    if (data.payload.gateway) {
+      return { headers: {} }
+    } else {
+      return {
+        headers: data.payload.headers
+      }
+    }
+  }
+
+  const wsConnectionParams = {
+    connectionInitPayload () {
+      return {
+        gateway: true
+      }
+    }
+  }
+
+  function createResolverService () {
+    const schema = `
+      extend type Query {
+        ignoredResolver: Boolean!
+      }
+    
+      extend type Event @key(fields: "value") {
+        id: ID! @external
+        userId: Int!
+      }  
+    `
+
+    const resolvers = {
+      Query: {
+        ignoredResolver: () => true
+      },
+      Event: {
+        userId: (root, args, ctx) => {
+          return parseInt(root.id)
+        }
+      }
+    }
+
+    resolverService = Fastify()
+    resolverService.register(GQL, {
+      schema,
+      resolvers,
+      federationMetadata: true,
+      subscription: { onConnect }
+    })
+
+    return resolverService.listen(0)
+  }
+
+  function createSubscriptionService () {
+    const schema = `
+      extend type Query {
+        ignored: Boolean!
+      }
+    
+      type Event @key(fields: "id") {
+        id: ID!
+      }
+
+      extend type Mutation {
+        addTestEvent(value: Int!): Int!
+      }
+      
+      extend type Subscription {
+        testEvent(value: Int!): Event!
+      }
+      `
+
+    const resolvers = {
+      Query: {
+        ignored: () => true
+      },
+      Mutation: {
+        addTestEvent: async (_, { value }, { pubsub }) => {
+          await pubsub.publish({
+            topic: 'testEvent',
+            payload: { testEvent: { id: value } }
+          })
+
+          return value
+        }
+      },
+      Subscription: {
+        testEvent: {
+          subscribe: GQL.withFilter(
+            async (_, __, { pubsub }) => {
+              return await pubsub.subscribe('testEvent')
+            },
+            (root, args, { headers }) => {
+              return headers.userId === root.testEvent.id
+            }
+          )
+        }
+      }
+    }
+
+    subscriptionService = Fastify()
+    subscriptionService.register(GQL, {
+      schema,
+      resolvers,
+      federationMetadata: true,
+      subscription: { onConnect }
+    })
+
+    return subscriptionService.listen(0)
+  }
+
+  function createGateway () {
+    const subscriptionServicePort = subscriptionService.server.address().port
+    const resolverServicePort = resolverService.server.address().port
+
+    gateway = Fastify()
+    gateway.register(GQL, {
+      subscription: true,
+      gateway: {
+        services: [
+          {
+            name: 'subscriptionService',
+            url: `http://localhost:${subscriptionServicePort}/graphql`,
+            wsUrl: `ws://localhost:${subscriptionServicePort}/graphql`,
+            wsConnectionParams
+          },
+          {
+            name: 'resolverService',
+            url: `http://localhost:${resolverServicePort}/graphql`,
+            wsUrl: `ws://localhost:${resolverServicePort}/graphql`,
+            wsConnectionParams
+          }
+        ]
+      }
+    })
+
+    return gateway.listen(0)
+  }
+
+  function runSubscription (id) {
+    const ws = new WebSocket(`ws://localhost:${(gateway.server.address()).port}/graphql`, 'graphql-ws')
+    const client = WebSocket.createWebSocketStream(ws, { encoding: 'utf8', objectMode: true })
+    t.teardown(async () => {
+      client.destroy()
+    })
+    client.setEncoding('utf8')
+
+    client.write(JSON.stringify({
+      type: 'connection_init',
+      payload: { headers: { userId: id } }
+    }))
+
+    client.write(JSON.stringify({
+      id: 1,
+      type: 'start',
+      payload: {
+        query: `
+          subscription TestEvent($value: Int!) {
+            testEvent(value: $value) {
+              id
+              userId
+            }
+          }
+        `,
+        variables: { value: id }
+      }
+    }))
+
+    client.write(JSON.stringify({
+      id: 2,
+      type: 'start',
+      payload: {
+        query: `
+          subscription TestEvent($value: Int!) {
+            testEvent(value: $value) {
+              id
+              userId
+            }
+          }
+        `,
+        variables: { value: id }
+      }
+    }))
+
+    client.write(JSON.stringify({
+      id: 2,
+      type: 'stop'
+    }))
+
+    let end
+
+    const endPromise = new Promise(resolve => {
+      end = resolve
+    })
+
+    client.on('data', (chunk) => {
+      const data = JSON.parse(chunk)
+
+      if (data.id === 1 && data.type === 'data') {
+        t.equal(chunk, JSON.stringify({
+          type: 'data',
+          id: 1,
+          payload: {
+            data: {
+              testEvent: {
+                id: String(id),
+                userId: id
+              }
+            }
+          }
+        }))
+
+        client.end()
+        end()
+      } else if (data.id === 2 && data.type === 'complete') {
+        gateway.inject({
+          method: 'POST',
+          url: '/graphql',
+          body: {
+            query: `
+              mutation AddTestEvent($value: Int!) {
+                addTestEvent(value: $value)
+              }
+            `,
+            variables: { value: id }
+          }
+        })
+      }
+    })
+
+    return endPromise
+  }
+
+  await createSubscriptionService()
+  await createResolverService()
+  await createGateway()
+  const subscriptions = new Array(10).fill(null).map((_, i) => runSubscription(i))
+  await Promise.all(subscriptions)
+
+  t.teardown(async () => {
+    await gateway.close()
+    await subscriptionService.close()
+    await resolverService.close()
+  })
+})
diff --git a/test/loaders.js b/test/loaders.js
index ed9523e..28cd24f 100644
--- a/test/loaders.js
+++ b/test/loaders.js
@@ -623,3 +623,288 @@ test('subscriptions properly execute loaders', t => {
     })
   })
 })
+
+test('Pass info to loader if cache is disabled', async (t) => {
+  const app = Fastify()
+
+  const dogs = [{
+    dogName: 'Max',
+    age: 10
+  }, {
+    dogName: 'Charlie',
+    age: 13
+  }, {
+    dogName: 'Buddy',
+    age: 15
+  }, {
+    dogName: 'Max',
+    age: 17
+  }]
+
+  const cats = [{
+    catName: 'Charlie',
+    age: 10
+  }, {
+    catName: 'Max',
+    age: 13
+  }, {
+    catName: 'Buddy',
+    age: 15
+  }]
+
+  const owners = {
+    Max: {
+      nickName: 'Jennifer',
+      age: 25
+    },
+    Charlie: {
+      nickName: 'Sarah',
+      age: 35
+    },
+    Buddy: {
+      nickName: 'Tracy',
+      age: 45
+    }
+  }
+
+  const schema = `
+    type Human {
+      nickName: String!
+      age: Int!
+    }
+
+    type Dog {
+      dogName: String!
+      age: Int!
+      owner: Human
+    }
+
+    type Cat {
+      catName: String!
+      age: Int!
+      owner: Human
+    }
+
+    type Query {
+      dogs: [Dog]
+      cats: [Cat]
+    }
+  `
+
+  const query = `{
+    dogs {
+      dogName
+      age
+      owner {
+        nickName
+        age
+      }
+    }
+    cats {
+      catName
+      owner {
+        age
+      }
+    }
+  }`
+  const resolvers = {
+    Query: {
+      dogs: (_, params, context) => {
+        return dogs
+      },
+      cats: (_, params, context) => {
+        return cats
+      }
+    }
+  }
+
+  const loaders = {
+    Dog: {
+      async owner (queries, context) {
+        t.equal(context.app, app)
+        return queries.map(({ obj, info }) => {
+          // verify info properties
+          t.equal(info.operation.operation, 'query')
+
+          const resolverOutputParams = info.operation.selectionSet.selections[0].selectionSet.selections
+          t.equal(resolverOutputParams.length, 3)
+          t.equal(resolverOutputParams[0].name.value, 'dogName')
+          t.equal(resolverOutputParams[1].name.value, 'age')
+          t.equal(resolverOutputParams[2].name.value, 'owner')
+
+          const loaderOutputParams = resolverOutputParams[2].selectionSet.selections
+
+          t.equal(loaderOutputParams.length, 2)
+          t.equal(loaderOutputParams[0].name.value, 'nickName')
+          t.equal(loaderOutputParams[1].name.value, 'age')
+
+          return owners[obj.dogName]
+        })
+      }
+    },
+    Cat: {
+      async owner (queries, context) {
+        t.equal(context.app, app)
+        return queries.map(({ obj, info }) => {
+          // verify info properties
+          t.equal(info.operation.operation, 'query')
+
+          const resolverOutputParams = info.operation.selectionSet.selections[1].selectionSet.selections
+          t.equal(resolverOutputParams.length, 2)
+          t.equal(resolverOutputParams[0].name.value, 'catName')
+          t.equal(resolverOutputParams[1].name.value, 'owner')
+
+          const loaderOutputParams = resolverOutputParams[1].selectionSet.selections
+
+          t.equal(loaderOutputParams.length, 1)
+          t.equal(loaderOutputParams[0].name.value, 'age')
+
+          return owners[obj.catName]
+        })
+      }
+    }
+  }
+
+  app.register(GQL, {
+    schema,
+    resolvers,
+    loaders,
+    cache: false
+  })
+
+  await app.ready()
+
+  const res = await app.inject({
+    method: 'POST',
+    url: '/graphql',
+    body: {
+      query
+    }
+  })
+
+  t.equal(res.statusCode, 200)
+  t.strictSame(JSON.parse(res.body), {
+    data: {
+      dogs: [
+        {
+          dogName: 'Max',
+          age: 10,
+          owner: {
+            nickName: 'Jennifer',
+            age: 25
+          }
+        },
+        {
+          dogName: 'Charlie',
+          age: 13,
+          owner: {
+            nickName: 'Sarah',
+            age: 35
+          }
+        },
+        {
+          dogName: 'Buddy',
+          age: 15,
+          owner: {
+            nickName: 'Tracy',
+            age: 45
+          }
+        },
+        {
+          dogName: 'Max',
+          age: 17,
+          owner: {
+            nickName: 'Jennifer',
+            age: 25
+          }
+        }
+      ],
+      cats: [
+        {
+          catName: 'Charlie',
+          owner: {
+            age: 35
+          }
+        },
+        {
+          catName: 'Max',
+          owner: {
+            age: 25
+          }
+        },
+        {
+          catName: 'Buddy',
+          owner: {
+            age: 45
+          }
+        }
+      ]
+    }
+  })
+})
+
+test('should not pass info to loader if cache is enabled', async (t) => {
+  const app = Fastify()
+
+  const resolvers = {
+    Query: {
+      dogs: (_, params, context) => {
+        return dogs
+      }
+    }
+  }
+
+  const loaders = {
+    Dog: {
+      async owner (queries) {
+        t.equal(queries[0].info, undefined)
+        return queries.map(({ obj }) => owners[obj.name])
+      }
+    }
+  }
+
+  app.register(GQL, {
+    schema,
+    resolvers,
+    loaders,
+    cache: true
+  })
+
+  // needed so that graphql is defined
+  await app.ready()
+
+  const query = 'query { dogs { name owner { name } } }'
+  const res = await app.inject({
+    method: 'POST',
+    url: '/graphql',
+    body: {
+      query
+    }
+  })
+
+  t.same(JSON.parse(res.body), {
+    data: {
+      dogs: [{
+        name: 'Max',
+        owner: {
+          name: 'Jennifer'
+        }
+      }, {
+        name: 'Charlie',
+        owner: {
+          name: 'Sarah'
+        }
+      }, {
+        name: 'Buddy',
+        owner: {
+          name: 'Tracy'
+        }
+      }, {
+        name: 'Max',
+        owner: {
+          name: 'Jennifer'
+        }
+      }]
+    }
+  })
+})
diff --git a/test/reply-decorator.js b/test/reply-decorator.js
index 21a3afb..f230775 100644
--- a/test/reply-decorator.js
+++ b/test/reply-decorator.js
@@ -123,3 +123,55 @@ test('reply decorator set status code to 400 with bad query', async (t) => {
   t.equal(res.statusCode, 400)
   t.matchSnapshot(JSON.stringify(JSON.parse(res.body)))
 })
+
+test('reply decorator supports encapsulation when loaders are defined in parent object', async (t) => {
+  const app = Fastify()
+  const schema = `
+    type Query {
+      add(x: Int, y: Int): Int
+    }
+  `
+  const resolvers = {
+    add: async ({ x, y }) => x + y
+  }
+
+  app.register(GQL, {
+    schema,
+    resolvers,
+    loaders: {}
+  })
+
+  app.register(async (app) => {
+    const schema = `
+      type Query {
+        multiply(x: Int, y: Int): Int
+      }
+    `
+    const resolvers = {
+      multiply: async ({ x, y }) => x * y
+    }
+
+    app.register(GQL, {
+      schema,
+      resolvers,
+      prefix: '/prefix'
+    })
+  })
+
+  const res = await app.inject({
+    method: 'POST',
+    url: '/prefix/graphql',
+    payload: {
+      query: '{ multiply(x: 5, y: 5) }'
+    }
+  })
+
+  t.equal(res.statusCode, 200)
+  t.same(JSON.parse(res.body), {
+    data: {
+      multiply: 25
+    }
+  })
+
+  t.matchSnapshot(JSON.stringify(JSON.parse(res.body)))
+})
diff --git a/test/subscription-connection.js b/test/subscription-connection.js
index 99b8bbe..efb8b42 100644
--- a/test/subscription-connection.js
+++ b/test/subscription-connection.js
@@ -6,7 +6,7 @@ const fastify = require('fastify')
 const mq = require('mqemitter')
 const SubscriptionConnection = require('../lib/subscription-connection')
 const { PubSub } = require('../lib/subscriber')
-const { GRAPHQL_WS } = require('../lib/subscription-protocol')
+const { GRAPHQL_WS, GRAPHQL_TRANSPORT_WS } = require('../lib/subscription-protocol')
 
 test('socket is closed on unhandled promise rejection in handleMessage', t => {
   t.plan(1)
@@ -95,7 +95,7 @@ test('subscription connection handles GQL_CONNECTION_TERMINATE message correctly
   }))
 })
 
-test('subscription connection closes context on GQL_STOP message correctly', async (t) => {
+test('subscription connection closes context on GQL_STOP message correctly (protocol: graphql-ws)', async (t) => {
   t.plan(2)
   const sc = new SubscriptionConnection({
     on () {},
@@ -119,7 +119,31 @@ test('subscription connection closes context on GQL_STOP message correctly', asy
   t.equal(sc.subscriptionContexts.size, 0)
 })
 
-test('subscription connection completes resolver iterator on GQL_STOP message correctly', async (t) => {
+test('subscription connection closes context on GQL_STOP message correctly (protocol: graphql-transport-ws)', async (t) => {
+  t.plan(2)
+  const sc = new SubscriptionConnection({
+    on () {},
+    close () {},
+    send (message) {},
+    protocol: GRAPHQL_TRANSPORT_WS
+  }, {})
+
+  sc.subscriptionContexts = new Map()
+  sc.subscriptionContexts.set(1, {
+    close () {
+      t.pass()
+    }
+  })
+
+  await sc.handleMessage(JSON.stringify({
+    id: 1,
+    type: 'complete'
+  }))
+
+  t.equal(sc.subscriptionContexts.size, 0)
+})
+
+test('subscription connection completes resolver iterator on GQL_STOP message correctly (protocol: graphql-ws)', async (t) => {
   t.plan(2)
   const sc = new SubscriptionConnection({
     on () {},
@@ -143,6 +167,30 @@ test('subscription connection completes resolver iterator on GQL_STOP message co
   t.equal(sc.subscriptionIters.size, 0)
 })
 
+test('subscription connection completes resolver iterator on GQL_STOP message correctly (protocol: graphql-transport-ws)', async (t) => {
+  t.plan(2)
+  const sc = new SubscriptionConnection({
+    on () {},
+    close () {},
+    send (message) {},
+    protocol: GRAPHQL_TRANSPORT_WS
+  }, {})
+
+  sc.subscriptionIters = new Map()
+  sc.subscriptionIters.set(1, {
+    return () {
+      t.pass()
+    }
+  })
+
+  await sc.handleMessage(JSON.stringify({
+    id: 1,
+    type: 'complete'
+  }))
+
+  t.equal(sc.subscriptionIters.size, 0)
+})
+
 test('handles error in send and closes connection', async t => {
   t.plan(1)
 
@@ -163,7 +211,7 @@ test('handles error in send and closes connection', async t => {
   await sc.sendMessage('foo')
 })
 
-test('subscription connection handles GQL_STOP message correctly, with no data', async (t) => {
+test('subscription connection handles GQL_STOP message correctly, with no data (protocol: graphql-ws)', async (t) => {
   const sc = new SubscriptionConnection({
     on () {},
     close () {},
@@ -179,6 +227,22 @@ test('subscription connection handles GQL_STOP message correctly, with no data',
   t.notOk(sc.subscriptionContexts.get(0))
 })
 
+test('subscription connection handles GQL_STOP message correctly, with no data (protocol: graphql-transport-ws)', async (t) => {
+  const sc = new SubscriptionConnection({
+    on () {},
+    close () {},
+    send (message) {},
+    protocol: GRAPHQL_TRANSPORT_WS
+  }, {})
+
+  await sc.handleMessage(JSON.stringify({
+    id: 1,
+    type: 'complete'
+  }))
+
+  t.notOk(sc.subscriptionContexts.get(0))
+})
+
 test('subscription connection send error message when GQL_START handler errs', async (t) => {
   t.plan(1)
 
diff --git a/test/subscription-protocol.js b/test/subscription-protocol.js
index b7c8472..4bc8c26 100644
--- a/test/subscription-protocol.js
+++ b/test/subscription-protocol.js
@@ -26,7 +26,7 @@ test('getProtocolByName returns correct protocol message types', t => {
     GQL_DATA: 'next',
     GQL_ERROR: 'error',
     GQL_COMPLETE: 'complete',
-    GQL_STOP: 'stop'
+    GQL_STOP: 'complete'
   })
   t.equal(getProtocolByName('unsupported-protocol'), null)
 })