feat(wip): support bypassing subscriptions

Author: kettanaito

src/core/handlers/GraphQLSubscriptionHandler.ts

@@ -1,5 +1,9 @@
+import { invariant } from 'outvariant'
 import { OperationTypeNode } from 'graphql'
-import type { WebSocketConnectionData } from '@mswjs/interceptors/WebSocket'
+import type {
+  WebSocketConnectionData,
+  WebSocketServerConnection,
+} from '@mswjs/interceptors/WebSocket'
 import {
   GraphQLHandler,
   type GraphQLHandlerInfo,
@@ -62,6 +66,10 @@ type GraphQLWebSocketNextMessage = {
   payload: unknown
 }
 
+type GraphQLSubscriptionIncomingMessage =
+  | GraphQLWebSocketAcknowledgeMessage
+  | GraphQLWebSocketNextMessage
+
 interface GraphQLWebSocketSubscribeMessagePayload<
   Variables extends GraphQLVariables = GraphQLVariables,
 > {
@@ -73,6 +81,8 @@ interface GraphQLWebSocketSubscribeMessagePayload<
 export class GraphQLInternalPubsub {
   public pubsub: GraphQLPubsub
   public webSocketLink: WebSocketLink
+  public server?: WebSocketServerConnection
+
   private subscriptions: Map<string, GraphQLWebSocketSubscribeMessage>
 
   constructor(public readonly url: Path) {
@@ -89,7 +99,9 @@ export class GraphQLInternalPubsub {
 
     const webSocketHandler = this.webSocketLink.addEventListener(
       'connection',
-      ({ client }) => {
+      ({ client, server }) => {
+        this.server = server
+
         client.addEventListener('message', (event) => {
           if (typeof event.data !== 'string') {
             return
@@ -138,13 +150,13 @@ export class GraphQLInternalPubsub {
     }
   }
 
-  private createAcknowledgeMessage() {
+  protected createAcknowledgeMessage() {
     return JSON.stringify({
       type: 'connection_ack',
     } satisfies GraphQLWebSocketAcknowledgeMessage)
   }
 
-  private createSubscribeMessage(args: { id: string; payload: unknown }) {
+  protected createSubscribeMessage(args: { id: string; payload: unknown }) {
     return JSON.stringify({
       id: args.id,
       type: 'next',
@@ -178,7 +190,7 @@ export class GraphQLSubscription<
   }
 
   /**
-   * Publishes data to this subscription.
+   * Publishes data to all subscribed GraphQL clients.
    */
   public publish(payload: { data?: Query }): void {
     this.args.internalPubsub.pubsub.publish(payload, ({ subscription }) => {
@@ -197,6 +209,31 @@ export class GraphQLSubscription<
       } satisfies GraphQLWebSocketCompleteMessage),
     )
   }
+
+  /**
+   * Establishes this GraphQL subscription to the actual server.
+   */
+  public subscribe() {
+    const { server } = this.args.internalPubsub
+
+    /**
+     * @note One can only call `subscription.subscribe()` inside the
+     * GraphQL subscription handler. By that point, the WebSocket connection
+     * has been established and intercepted so the `server` reference
+     * is guaranteed.
+     */
+    invariant(
+      server,
+      'Failed to establish GraphQL subscription ("%s") to the actual server ("%s"): `subscription.subscribe()` was called outside the subscription handler',
+      this.args.message.payload.query,
+      this.args.internalPubsub.url,
+    )
+
+    return new GraphQLServerSubscription({
+      server,
+      message: this.args.message,
+    })
+  }
 }
 
 export type GraphQLSubscriptionName<
@@ -334,3 +371,80 @@ export function createGraphQLSubscriptionHandler(
     )
   }
 }
+
+interface GraphQLServerSubscriptionEventMap {
+  connection_ack: Event
+  next: MessageEvent<GraphQLWebSocketNextMessage>
+}
+
+/**
+ * Representation of a GraphQL subscription to the actual server.
+ * You interface with this object from the client's perspective.
+ */
+class GraphQLServerSubscription {
+  protected readonly server: WebSocketServerConnection
+
+  private eventTarget: EventTarget
+
+  constructor(
+    readonly args: {
+      server: WebSocketServerConnection
+      message: GraphQLWebSocketSubscribeMessage
+    },
+  ) {
+    this.eventTarget = new EventTarget()
+    this.server = args.server
+
+    this.server.connect()
+    this.server.addEventListener('open', () => {
+      // Once the server connetion has been established, send the
+      // subscription intent message. This is the same message as
+      // was sent from the GraphQL client.
+      this.server.send(JSON.stringify(this.args.message))
+    })
+
+    this.server.addEventListener('message', (event) => {
+      if (typeof event.data !== 'string') {
+        return
+      }
+
+      const message = jsonParse<GraphQLSubscriptionIncomingMessage>(event.data)
+      if (!message) {
+        return
+      }
+
+      switch (message.type) {
+        case 'connection_ack': {
+          this.eventTarget.dispatchEvent(new Event('connection_ack'))
+          break
+        }
+
+        case 'next': {
+          /**
+           * @fixme This isn't the same `event` from the server
+           * so calling `event.preventDefault()` won't prevent anything.
+           */
+          this.eventTarget.dispatchEvent(
+            new MessageEvent('next', {
+              data: message,
+            }),
+          )
+          break
+        }
+      }
+    })
+  }
+
+  public addEventListener<
+    EventType extends keyof GraphQLServerSubscriptionEventMap,
+  >(
+    event: EventType,
+    listener: (event: GraphQLServerSubscriptionEventMap[EventType]) => void,
+  ) {
+    this.eventTarget.addEventListener(event, { handleEvent: listener })
+  }
+
+  public unsubscribe(): void {
+    this.server.close()
+  }
+}

test/node/graphql-api/graphql-subscription.test.ts

@@ -119,3 +119,51 @@ subscription OnCommentAdded {
     service: 'user-service',
   })
 })
+
+it('supports bypassing a GraphQL subscription', async () => {
+  const api = graphql.link('https://api.example.com/graphql')
+  server.use(
+    api.subscription('OnCommentAdded', async ({ subscription }) => {
+      // This creates a NEW subscription for the same event in the actual server.
+      // - There is no client events to prevent! Subscription intent is sent ONCE.
+      const commentSubscription = subscription.subscribe()
+
+      // Event listeners here are modeled after the subscription protocol:
+      // - acknowledge, server confirmed the subscription.
+      // - next, server is sending data to the client.
+      // - error, error happened (is this event a thing?). Server connection errors!
+      // - complete, server has completed this subscription.
+      commentSubscription.addEventListener('next', (event) => {
+        // You can still prevent messages from the original server.
+        // By default, they are forwarded to the client, just like with mocking WebSockets.
+        event.preventDefault()
+        // The event data is already parsed to drop the implementation details of the subscription.
+        subscription.publish({
+          data: event.data.payload,
+        })
+      })
+
+      // You can unsubscribe from the original server subscription at any time.
+      commentSubscription.unsubscribe()
+    }),
+  )
+
+  const client = createClient({
+    url: 'wss://api.example.com/graphql',
+  })
+  const subscription = client.iterate({
+    query: `
+subscription OnCommentAdded {
+  commentAdded {
+    text
+  }
+}
+    `,
+  })
+
+  await subscription.next()
+
+  /**
+   * @fixme Expect the subscription payload from the ACTUAL server.
+   */
+})