test(nestjs-gateway): restore integration coverage

packages/nestjs-gateway/integration/test/gateway.test.ts → packages/nestjs-gateway/integration/gateway.test.ts

@@ -1,58 +1,127 @@
-import type { MeshPubSub }                                               from '@graphql-mesh/types'
-import type { INestApplication }                                         from '@nestjs/common'
-import type { INestMicroservice }                                        from '@nestjs/common'
-
-import type { GatewaySourceType as GatewaySourceTypeEnum }               from '../../src/index.js'
-import type { GATEWAY_MESH_PUBSUB as GatewayMeshPubSubToken }            from '../../src/index.js'
-import type { GATEWAY_MODULE_OPTIONS as GatewayModuleOptionsToken }      from '../../src/index.js'
-import type { GatewayIntegrationModule as GatewayIntegrationModuleType } from '../src/index.js'
-
-import assert                                                            from 'node:assert/strict'
-import path                                                              from 'node:path'
-import { before }                                                        from 'node:test'
-import { after }                                                         from 'node:test'
-import { describe }                                                      from 'node:test'
-import { it }                                                            from 'node:test'
-import { fileURLToPath }                                                 from 'node:url'
-
-import { Transport } from '@nestjs/microservices'
-import { Test }                                                          from '@nestjs/testing'
-import { WebSocket }                                                     from 'ws'
-import { buildClientSchema }                                             from 'graphql'
-import { printSchema }                                                   from 'graphql'
-import { getIntrospectionQuery }                                         from 'graphql'
-import { createClient }                                                  from 'graphql-ws'
-import getPort                                                           from 'get-port'
-import request                                                           from 'supertest'
+import type { MeshPubSub }                                          from '@graphql-mesh/types'
+import type { INestApplication }                                    from '@nestjs/common'
+import type { INestMicroservice }                                   from '@nestjs/common'
+
+import type { GatewaySourceType as GatewaySourceTypeEnum }          from '../src/index.js'
+import type { GATEWAY_MESH_PUBSUB as GatewayMeshPubSubToken }       from '../src/index.js'
+import type { GATEWAY_MODULE_OPTIONS as GatewayModuleOptionsToken } from '../src/index.js'
+import type { SubscriptionResult }                                  from './interfaces.js'
+import type { ApplicationModule as ApplicationModuleType }          from './src/index.js'
+import type { ServiceModule as ServiceModuleType }                  from './src/index.js'
+
+import assert                                                       from 'node:assert/strict'
+import path                                                         from 'node:path'
+import { before }                                                   from 'node:test'
+import { after }                                                    from 'node:test'
+import { describe }                                                 from 'node:test'
+import { it }                                                       from 'node:test'
+import { fileURLToPath }                                            from 'node:url'
+
+import { Transport }                                                from '@nestjs/microservices'
+import { Test }                                                     from '@nestjs/testing'
+import { WebSocket }                                                from 'ws'
+import { buildClientSchema }                                        from 'graphql'
+import { printSchema }                                              from 'graphql'
+import { getIntrospectionQuery }                                    from 'graphql'
+import { createClient }                                             from 'graphql-ws'
+import getPort                                                      from 'get-port'
+import request                                                      from 'supertest'
 
 const moduleDir = path.dirname(fileURLToPath(import.meta.url))
+const SUBSCRIPTION_TIMEOUT = 10000
+const SUBSCRIPTION_TOPIC = 'eventTriggered'
+
+const withTimeout = async <T>(promise: Promise<T>, message: string): Promise<T> => {
+  let timeout: ReturnType<typeof setTimeout> | undefined
+
+  try {
+    return await Promise.race([
+      promise,
+      new Promise<T>((_, reject) => {
+        timeout = setTimeout(() => {
+          reject(new Error(message))
+        }, SUBSCRIPTION_TIMEOUT)
+      }),
+    ])
+  } finally {
+    if (timeout) {
+      clearTimeout(timeout)
+    }
+  }
+}
+
+const createSubscriptionReadyTracker = (meshPubSub: MeshPubSub, triggerName: string) => {
+  const originalSubscribe = meshPubSub.subscribe
+  const originalAsyncIterator = meshPubSub.asyncIterator
+  const subscribe = originalSubscribe.bind(meshPubSub)
+  const asyncIterator = originalAsyncIterator.bind(meshPubSub)
+
+  let markReady!: () => void
+  let isReady = false
+
+  const ready = new Promise<void>((resolve) => {
+    markReady = () => {
+      if (!isReady) {
+        isReady = true
+        resolve()
+      }
+    }
+  })
+
+  const trackReady = (trigger: string) => {
+    if (trigger === triggerName) {
+      markReady()
+    }
+  }
+
+  meshPubSub.subscribe = ((trigger, onMessage, options) => {
+    trackReady(trigger)
+
+    return subscribe(trigger, onMessage, options)
+  }) as MeshPubSub['subscribe']
 
-// TODO: fix gateway integration test stability and re-enable suite.
-describe.skip('gateway', () => {
+  meshPubSub.asyncIterator = ((trigger) => {
+    trackReady(trigger)
+
+    return asyncIterator(trigger)
+  }) as MeshPubSub['asyncIterator']
+
+  return {
+    ready,
+    restore: () => {
+      meshPubSub.subscribe = originalSubscribe
+      meshPubSub.asyncIterator = originalAsyncIterator
+    },
+  }
+}
+
+describe('gateway', () => {
   let GatewaySourceType: typeof GatewaySourceTypeEnum
   let GATEWAY_MESH_PUBSUB: typeof GatewayMeshPubSubToken
   let GATEWAY_MODULE_OPTIONS: typeof GatewayModuleOptionsToken
-  let GatewayIntegrationModule: typeof GatewayIntegrationModuleType
+  let ApplicationModule: typeof ApplicationModuleType
+  let ServiceModule: typeof ServiceModuleType
 
   let service: INestMicroservice
   let app: INestApplication
   let pubsub: MeshPubSub
   let url: string
 
   before(async () => {
-    const gatewayCore = await import('../../src/index.js')
-    const gatewayIntegration = await import('../src/index.js')
+    const gatewayCore = await import('../src/index.js')
+    const gatewayIntegration = await import('./src/index.js')
 
     GatewaySourceType = gatewayCore.GatewaySourceType
     GATEWAY_MESH_PUBSUB = gatewayCore.GATEWAY_MESH_PUBSUB
     GATEWAY_MODULE_OPTIONS = gatewayCore.GATEWAY_MODULE_OPTIONS
-    GatewayIntegrationModule = gatewayIntegration.GatewayIntegrationModule
+    ApplicationModule = gatewayIntegration.ApplicationModule
+    ServiceModule = gatewayIntegration.ServiceModule
 
     const servicePort = await getPort()
     const appPort = await getPort()
 
-    const testingModule = await Test.createTestingModule({
-      imports: [GatewayIntegrationModule],
+    const applicationTestingModule = await Test.createTestingModule({
+      imports: [ApplicationModule],
     })
       .overrideProvider(GATEWAY_MODULE_OPTIONS)
       .useValue({
@@ -64,7 +133,7 @@ describe.skip('gateway', () => {
             handler: {
               endpoint: `localhost:${servicePort}`,
               protoFilePath: {
-                file: path.join(moduleDir, '../src/service.proto'),
+                file: path.join(moduleDir, 'src/service.proto'),
                 load: { includeDirs: [] },
               },
               serviceName: 'ExampleService',
@@ -138,13 +207,17 @@ describe.skip('gateway', () => {
       })
       .compile()
 
-    app = testingModule.createNestApplication()
+    const serviceTestingModule = await Test.createTestingModule({
+      imports: [ServiceModule],
+    }).compile()
+
+    app = applicationTestingModule.createNestApplication()
 
-    service = testingModule.createNestMicroservice({
+    service = serviceTestingModule.createNestMicroservice({
       transport: Transport.GRPC,
       options: {
         package: ['tech.atls'],
-        protoPath: [path.join(moduleDir, '../src/service.proto')],
+        protoPath: [path.join(moduleDir, 'src/service.proto')],
         url: `0.0.0.0:${servicePort}`,
         loader: {
           arrays: true,
@@ -156,7 +229,6 @@ describe.skip('gateway', () => {
       },
     })
 
-    await app.init()
     await service.init()
 
     await app.listen(appPort, '0.0.0.0')
@@ -250,17 +322,39 @@ describe.skip('gateway', () => {
     assert.strictEqual(exception.message, 'Test')
   })
 
-  // TODO: check the test and implemenation. Event doesn't resolve
-  it.skip('check subscriptions', async () => {
+  it('check subscriptions', async () => {
+    let openConnection!: () => void
+    const connected = new Promise<void>((resolve) => {
+      openConnection = resolve
+    })
+    const subscriptionReady = createSubscriptionReadyTracker(pubsub, SUBSCRIPTION_TOPIC)
+
     const client = createClient({
       url: url.replace('http:', 'ws:'),
       webSocketImpl: WebSocket,
+      on: {
+        connected: openConnection,
+      },
     })
 
-    const event = new Promise((resolve, reject) => {
-      let result: { id: string } | undefined
+    let dispose: (() => void) | undefined
+    let timeout: ReturnType<typeof setTimeout> | undefined
+
+    const disposeClient = () => {
+      if (timeout) {
+        clearTimeout(timeout)
+      }
+      dispose?.()
+      client.dispose()
+    }
 
-      client.subscribe(
+    const event = new Promise<SubscriptionResult>((resolve, reject) => {
+      timeout = setTimeout(() => {
+        disposeClient()
+        reject(new Error('Subscription result missing'))
+      }, SUBSCRIPTION_TIMEOUT)
+
+      dispose = client.subscribe(
         {
           query: `subscription onEventTriggered {
               eventTriggered {
@@ -270,23 +364,48 @@ describe.skip('gateway', () => {
         },
         {
           next: (data) => {
-            result = data as { id: string }
+            clearTimeout(timeout)
+            disposeClient()
+            resolve(data as SubscriptionResult)
+          },
+          error: (error) => {
+            clearTimeout(timeout)
+            disposeClient()
+            reject(error)
           },
-          error: reject,
           complete: () => {
-            if (!result) {
-              reject(new Error('Subscription result missing'))
-              return
-            }
-            resolve(result)
+            clearTimeout(timeout)
+            disposeClient()
+            reject(new Error('Subscription completed before result'))
           },
         }
       )
-
-      pubsub.publish('eventTriggered', { id: 'test' })
     })
 
-    const result = await event
-    assert.deepStrictEqual(result, { id: 'test' })
+    let eventAwaited = false
+
+    try {
+      await withTimeout(connected, 'WebSocket connection missing')
+      await withTimeout(subscriptionReady.ready, 'Subscription was not registered')
+      pubsub.publish(SUBSCRIPTION_TOPIC, { id: 'test' })
+
+      eventAwaited = true
+      const result = await event
+      assert.deepStrictEqual(result, {
+        data: {
+          eventTriggered: {
+            id: 'test',
+          },
+        },
+      })
+    } finally {
+      subscriptionReady.restore()
+
+      if (!eventAwaited) {
+        event.catch(() => undefined)
+      }
+
+      disposeClient()
+    }
   })
 })

packages/nestjs-gateway/integration/interfaces.ts

@@ -0,0 +1,7 @@
+export type SubscriptionResult = {
+  data?: {
+    eventTriggered?: {
+      id?: string
+    }
+  }
+}

packages/nestjs-gateway/integration/src/application.module.ts

@@ -0,0 +1,8 @@
+import { Module }        from '@nestjs/common'
+
+import { GatewayModule } from '../../src/index.js'
+
+@Module({
+  imports: [GatewayModule.register()],
+})
+export class ApplicationModule {}

packages/nestjs-gateway/integration/src/index.ts

@@ -1 +1,2 @@
-export * from './gateway-integration.module.js'
+export * from './application.module.js'
+export * from './service.module.js'

packages/nestjs-gateway/integration/src/gateway-integration.module.ts → packages/nestjs-gateway/integration/src/service.module.ts

@@ -1,10 +1,8 @@
 import { Module }           from '@nestjs/common'
 
-import { GatewayModule }    from '../../src/index.js'
 import { MoviesController } from './movies.controller.js'
 
 @Module({
-  imports: [GatewayModule.register()],
   controllers: [MoviesController],
 })
-export class GatewayIntegrationModule {}
+export class ServiceModule {}

packages/nestjs-gateway/src/mesh/graphql-mesh.handler.ts

@@ -1,11 +1,11 @@
 import type { OnModuleDestroy }               from '@nestjs/common'
 import type { OnModuleInit }                  from '@nestjs/common'
-import type { WebSocketServer }               from 'ws'
 
 import type { GatewayModuleOptions }          from '../module/interfaces.js'
 import type { GatewayHttpBoundary }           from './http/interfaces.js'
 import type { GatewayHttpServer }             from './http/interfaces.js'
 import type { GatewayGraphQLRuntime }         from './interfaces.js'
+import type { GatewaySubscriptionServer }     from './interfaces.js'
 
 import { Inject }                             from '@nestjs/common'
 import { Injectable }                         from '@nestjs/common'
@@ -22,7 +22,7 @@ import { GraphQLMeshRuntime }                 from './runtime.js'
 export class GraphQLMeshHandler implements OnModuleInit, OnModuleDestroy {
   private runtime?: GatewayGraphQLRuntime
 
-  private webSocketServer?: WebSocketServer
+  private subscriptionServer?: GatewaySubscriptionServer
 
   constructor(
     private readonly adapterHost: HttpAdapterHost,
@@ -46,21 +46,16 @@ export class GraphQLMeshHandler implements OnModuleInit, OnModuleDestroy {
     }
 
     this.runtime = runtime
-    this.webSocketServer = this.meshRuntime.registerSubscriptions(
+    this.subscriptionServer = this.meshRuntime.registerSubscriptions(
       runtime,
       this.adapterHost.httpAdapter.getHttpServer() as GatewayHttpServer,
       this.options.path || '/'
     )
   }
 
   async onModuleDestroy(): Promise<void> {
+    await this.subscriptionServer?.dispose()
     await this.runtime?.apolloServer.stop()
-
-    if (this.webSocketServer) {
-      for (const client of this.webSocketServer.clients) {
-        client.close(1001, 'Going away')
-      }
-    }
   }
 
   private getHttpGateway(): GatewayHttpBoundary {

packages/nestjs-gateway/src/mesh/http/interfaces.ts

@@ -9,6 +9,10 @@ export interface GatewayHttpServer {
     event: 'upgrade',
     handler: (req: IncomingMessage, socket: Socket, head: Buffer) => void
   ) => void
+  off: (
+    event: 'upgrade',
+    handler: (req: IncomingMessage, socket: Socket, head: Buffer) => void
+  ) => void
 }
 
 export interface GatewayHttpBoundary {