Repository: apollographql/apollo-client

Problem statement:
Add `subscribeToMore` functionality to `useQueryRefHandlers`, `useBackgroundQuery`
Context: https://community.apollographql.com/t/how-to-use-subscriptions-together-with-preloadquery/7434/3

Some of our hooks include `subscribeToMore` in the results object.  We want to ensure that all relevant hooks include `subscribeToMore`, including `useQueryRefHandlers` and `useBackgroundQuery`.


Hints:

Patch:
diff --git a/src/core/ObservableQuery.ts b/src/core/ObservableQuery.ts
--- a/src/core/ObservableQuery.ts
+++ b/src/core/ObservableQuery.ts
@@ -163,6 +163,8 @@ export class ObservableQuery<
     this.waitForOwnResult = skipCacheDataFor(options.fetchPolicy);
     this.isTornDown = false;
 
+    this.subscribeToMore = this.subscribeToMore.bind(this);
+
     const {
       watchQuery: { fetchPolicy: defaultFetchPolicy = "cache-first" } = {},
     } = queryManager.defaultOptions;
diff --git a/src/react/hooks/useBackgroundQuery.ts b/src/react/hooks/useBackgroundQuery.ts
--- a/src/react/hooks/useBackgroundQuery.ts
+++ b/src/react/hooks/useBackgroundQuery.ts
@@ -17,7 +17,11 @@ import type { CacheKey, QueryRef } from "../internal/index.js";
 import type { BackgroundQueryHookOptions, NoInfer } from "../types/types.js";
 import { wrapHook } from "./internal/index.js";
 import { useWatchQueryOptions } from "./useSuspenseQuery.js";
-import type { FetchMoreFunction, RefetchFunction } from "./useSuspenseQuery.js";
+import type {
+  FetchMoreFunction,
+  RefetchFunction,
+  SubscribeToMoreFunction,
+} from "./useSuspenseQuery.js";
 import { canonicalStringify } from "../../cache/index.js";
 import type { DeepPartial } from "../../utilities/index.js";
 import type { SkipToken } from "./constants.js";
@@ -26,7 +30,11 @@ export type UseBackgroundQueryResult<
   TData = unknown,
   TVariables extends OperationVariables = OperationVariables,
 > = {
+  /** {@inheritDoc @apollo/client!ObservableQuery#subscribeToMore:member(1)} */
+  subscribeToMore: SubscribeToMoreFunction<TData, TVariables>;
+  /** {@inheritDoc @apollo/client!ObservableQuery#fetchMore:member(1)} */
   fetchMore: FetchMoreFunction<TData, TVariables>;
+  /** {@inheritDoc @apollo/client!ObservableQuery#refetch:member(1)} */
   refetch: RefetchFunction<TData, TVariables>;
 };
 
@@ -281,6 +289,10 @@ function _useBackgroundQuery<
 
   return [
     didFetchResult.current ? wrappedQueryRef : void 0,
-    { fetchMore, refetch },
+    {
+      fetchMore,
+      refetch,
+      subscribeToMore: queryRef.observable.subscribeToMore,
+    },
   ];
 }
diff --git a/src/react/hooks/useLoadableQuery.ts b/src/react/hooks/useLoadableQuery.ts
--- a/src/react/hooks/useLoadableQuery.ts
+++ b/src/react/hooks/useLoadableQuery.ts
@@ -18,7 +18,11 @@ import type { CacheKey, QueryRef } from "../internal/index.js";
 import type { LoadableQueryHookOptions } from "../types/types.js";
 import { __use, useRenderGuard } from "./internal/index.js";
 import { useWatchQueryOptions } from "./useSuspenseQuery.js";
-import type { FetchMoreFunction, RefetchFunction } from "./useSuspenseQuery.js";
+import type {
+  FetchMoreFunction,
+  RefetchFunction,
+  SubscribeToMoreFunction,
+} from "./useSuspenseQuery.js";
 import { canonicalStringify } from "../../cache/index.js";
 import type {
   DeepPartial,
@@ -49,6 +53,8 @@ export type UseLoadableQueryResult<
     fetchMore: FetchMoreFunction<TData, TVariables>;
     /** {@inheritDoc @apollo/client!QueryResultDocumentation#refetch:member} */
     refetch: RefetchFunction<TData, TVariables>;
+    /** {@inheritDoc @apollo/client!ObservableQuery#subscribeToMore:member(1)} */
+    subscribeToMore: SubscribeToMoreFunction<TData, TVariables>;
     /**
      * A function that resets the `queryRef` back to `null`.
      */
@@ -255,9 +261,22 @@ export function useLoadableQuery<
     ]
   );
 
+  const subscribeToMore: SubscribeToMoreFunction<TData, TVariables> =
+    React.useCallback(
+      (options) => {
+        invariant(
+          internalQueryRef,
+          "The query has not been loaded. Please load the query."
+        );
+
+        return internalQueryRef.observable.subscribeToMore(options);
+      },
+      [internalQueryRef]
+    );
+
   const reset: ResetFunction = React.useCallback(() => {
     setQueryRef(null);
   }, []);
 
-  return [loadQuery, queryRef, { fetchMore, refetch, reset }];
+  return [loadQuery, queryRef, { fetchMore, refetch, reset, subscribeToMore }];
 }
diff --git a/src/react/hooks/useQueryRefHandlers.ts b/src/react/hooks/useQueryRefHandlers.ts
--- a/src/react/hooks/useQueryRefHandlers.ts
+++ b/src/react/hooks/useQueryRefHandlers.ts
@@ -8,7 +8,11 @@ import {
 } from "../internal/index.js";
 import type { QueryRef } from "../internal/index.js";
 import type { OperationVariables } from "../../core/types.js";
-import type { RefetchFunction, FetchMoreFunction } from "./useSuspenseQuery.js";
+import type {
+  RefetchFunction,
+  FetchMoreFunction,
+  SubscribeToMoreFunction,
+} from "./useSuspenseQuery.js";
 import type { FetchMoreQueryOptions } from "../../core/watchQueryOptions.js";
 import { useApolloClient } from "./useApolloClient.js";
 import { wrapHook } from "./internal/index.js";
@@ -21,6 +25,8 @@ export interface UseQueryRefHandlersResult<
   refetch: RefetchFunction<TData, TVariables>;
   /** {@inheritDoc @apollo/client!ObservableQuery#fetchMore:member(1)} */
   fetchMore: FetchMoreFunction<TData, TVariables>;
+  /** {@inheritDoc @apollo/client!ObservableQuery#subscribeToMore:member(1)} */
+  subscribeToMore: SubscribeToMoreFunction<TData, TVariables>;
 }
 
 /**
@@ -112,5 +118,9 @@ function _useQueryRefHandlers<
     [internalQueryRef]
   );
 
-  return { refetch, fetchMore };
+  return {
+    refetch,
+    fetchMore,
+    subscribeToMore: internalQueryRef.observable.subscribeToMore,
+  };
 }
diff --git a/src/react/hooks/useSuspenseQuery.ts b/src/react/hooks/useSuspenseQuery.ts
--- a/src/react/hooks/useSuspenseQuery.ts
+++ b/src/react/hooks/useSuspenseQuery.ts
@@ -274,13 +274,7 @@ function _useSuspenseQuery<
     [queryRef]
   );
 
-  const subscribeToMore: SubscribeToMoreFunction<
-    TData | undefined,
-    TVariables
-  > = React.useCallback(
-    (options) => queryRef.observable.subscribeToMore(options),
-    [queryRef]
-  );
+  const subscribeToMore = queryRef.observable.subscribeToMore;
 
   return React.useMemo<
     UseSuspenseQueryResult<TData | undefined, TVariables>


Test patch:
diff --git a/src/react/hooks/__tests__/useBackgroundQuery.test.tsx b/src/react/hooks/__tests__/useBackgroundQuery.test.tsx
--- a/src/react/hooks/__tests__/useBackgroundQuery.test.tsx
+++ b/src/react/hooks/__tests__/useBackgroundQuery.test.tsx
@@ -16,6 +16,7 @@ import {
   TypedDocumentNode,
   ApolloLink,
   Observable,
+  split,
 } from "../../../core";
 import {
   MockedResponse,
@@ -29,6 +30,7 @@ import {
   concatPagination,
   offsetLimitPagination,
   DeepPartial,
+  getMainDefinition,
 } from "../../../utilities";
 import { useBackgroundQuery } from "../useBackgroundQuery";
 import { UseReadQueryResult, useReadQuery } from "../useReadQuery";
@@ -37,7 +39,10 @@ import { QueryRef, QueryReference } from "../../internal";
 import { InMemoryCache } from "../../../cache";
 import { SuspenseQueryHookFetchPolicy } from "../../types/types";
 import equal from "@wry/equality";
-import { RefetchWritePolicy } from "../../../core/watchQueryOptions";
+import {
+  RefetchWritePolicy,
+  SubscribeToMoreOptions,
+} from "../../../core/watchQueryOptions";
 import { skipToken } from "../constants";
 import {
   PaginatedCaseData,
@@ -54,6 +59,7 @@ import {
   spyOnConsole,
   useTrackRenders,
 } from "../../../testing/internal";
+import { SubscribeToMoreFunction } from "../useSuspenseQuery";
 
 afterEach(() => {
   jest.useRealTimers();
@@ -6052,6 +6058,135 @@ describe("fetchMore", () => {
 
     await expect(Profiler).not.toRerender();
   });
+
+  it("can subscribe to subscriptions and react to cache updates via `subscribeToMore`", async () => {
+    interface SubscriptionData {
+      greetingUpdated: string;
+    }
+
+    type UpdateQueryFn = NonNullable<
+      SubscribeToMoreOptions<
+        SimpleCaseData,
+        Record<string, never>,
+        SubscriptionData
+      >["updateQuery"]
+    >;
+
+    const subscription: TypedDocumentNode<
+      SubscriptionData,
+      Record<string, never>
+    > = gql`
+      subscription {
+        greetingUpdated
+      }
+    `;
+
+    const { mocks, query } = setupSimpleCase();
+
+    const wsLink = new MockSubscriptionLink();
+    const mockLink = new MockLink(mocks);
+
+    const link = split(
+      ({ query }) => {
+        const definition = getMainDefinition(query);
+
+        return (
+          definition.kind === "OperationDefinition" &&
+          definition.operation === "subscription"
+        );
+      },
+      wsLink,
+      mockLink
+    );
+
+    const client = new ApolloClient({ link, cache: new InMemoryCache() });
+
+    const Profiler = createProfiler({
+      initialSnapshot: {
+        subscribeToMore: null as SubscribeToMoreFunction<
+          SimpleCaseData,
+          Record<string, never>
+        > | null,
+        result: null as UseReadQueryResult<SimpleCaseData> | null,
+      },
+    });
+
+    const { SuspenseFallback, ReadQueryHook } =
+      createDefaultTrackedComponents(Profiler);
+
+    function App() {
+      useTrackRenders();
+      const [queryRef, { subscribeToMore }] = useBackgroundQuery(query);
+
+      Profiler.mergeSnapshot({ subscribeToMore });
+
+      return (
+        <Suspense fallback={<SuspenseFallback />}>
+          <ReadQueryHook queryRef={queryRef} />
+        </Suspense>
+      );
+    }
+
+    renderWithClient(<App />, { client, wrapper: Profiler });
+
+    {
+      const { renderedComponents } = await Profiler.takeRender();
+
+      expect(renderedComponents).toStrictEqual([App, SuspenseFallback]);
+    }
+
+    {
+      const { renderedComponents, snapshot } = await Profiler.takeRender();
+
+      expect(renderedComponents).toStrictEqual([ReadQueryHook]);
+      expect(snapshot.result).toEqual({
+        data: { greeting: "Hello" },
+        error: undefined,
+        networkStatus: NetworkStatus.ready,
+      });
+    }
+
+    const updateQuery = jest.fn<
+      ReturnType<UpdateQueryFn>,
+      Parameters<UpdateQueryFn>
+    >((_, { subscriptionData: { data } }) => {
+      return { greeting: data.greetingUpdated };
+    });
+
+    const { snapshot } = Profiler.getCurrentRender();
+
+    snapshot.subscribeToMore!({ document: subscription, updateQuery });
+
+    wsLink.simulateResult({
+      result: {
+        data: {
+          greetingUpdated: "Subscription hello",
+        },
+      },
+    });
+
+    {
+      const { snapshot, renderedComponents } = await Profiler.takeRender();
+
+      expect(renderedComponents).toStrictEqual([ReadQueryHook]);
+      expect(snapshot.result).toEqual({
+        data: { greeting: "Subscription hello" },
+        error: undefined,
+        networkStatus: NetworkStatus.ready,
+      });
+    }
+
+    expect(updateQuery).toHaveBeenCalledTimes(1);
+    expect(updateQuery).toHaveBeenCalledWith(
+      { greeting: "Hello" },
+      {
+        subscriptionData: {
+          data: { greetingUpdated: "Subscription hello" },
+        },
+        variables: {},
+      }
+    );
+  });
 });
 
 describe.skip("type tests", () => {
diff --git a/src/react/hooks/__tests__/useLoadableQuery.test.tsx b/src/react/hooks/__tests__/useLoadableQuery.test.tsx
--- a/src/react/hooks/__tests__/useLoadableQuery.test.tsx
+++ b/src/react/hooks/__tests__/useLoadableQuery.test.tsx
@@ -22,6 +22,8 @@ import {
   Observable,
   OperationVariables,
   RefetchWritePolicy,
+  SubscribeToMoreOptions,
+  split,
 } from "../../../core";
 import {
   MockedProvider,
@@ -35,6 +37,7 @@ import {
   concatPagination,
   offsetLimitPagination,
   DeepPartial,
+  getMainDefinition,
 } from "../../../utilities";
 import { useLoadableQuery } from "../useLoadableQuery";
 import type { UseReadQueryResult } from "../useReadQuery";
@@ -43,7 +46,11 @@ import { ApolloProvider } from "../../context";
 import { InMemoryCache } from "../../../cache";
 import { LoadableQueryHookFetchPolicy } from "../../types/types";
 import { QueryRef } from "../../../react";
-import { FetchMoreFunction, RefetchFunction } from "../useSuspenseQuery";
+import {
+  FetchMoreFunction,
+  RefetchFunction,
+  SubscribeToMoreFunction,
+} from "../useSuspenseQuery";
 import invariant, { InvariantError } from "ts-invariant";
 import {
   Profiler,
@@ -4667,6 +4674,218 @@ it("allows loadQuery to be called in useEffect on first render", async () => {
   expect(() => renderWithMocks(<App />, { mocks })).not.toThrow();
 });
 
+it("can subscribe to subscriptions and react to cache updates via `subscribeToMore`", async () => {
+  interface SubscriptionData {
+    greetingUpdated: string;
+  }
+
+  type UpdateQueryFn = NonNullable<
+    SubscribeToMoreOptions<
+      SimpleCaseData,
+      Record<string, never>,
+      SubscriptionData
+    >["updateQuery"]
+  >;
+
+  const subscription: TypedDocumentNode<
+    SubscriptionData,
+    Record<string, never>
+  > = gql`
+    subscription {
+      greetingUpdated
+    }
+  `;
+
+  const { mocks, query } = setupSimpleCase();
+
+  const wsLink = new MockSubscriptionLink();
+  const mockLink = new MockLink(mocks);
+
+  const link = split(
+    ({ query }) => {
+      const definition = getMainDefinition(query);
+
+      return (
+        definition.kind === "OperationDefinition" &&
+        definition.operation === "subscription"
+      );
+    },
+    wsLink,
+    mockLink
+  );
+
+  const client = new ApolloClient({ link, cache: new InMemoryCache() });
+
+  const Profiler = createProfiler({
+    initialSnapshot: {
+      subscribeToMore: null as SubscribeToMoreFunction<
+        SimpleCaseData,
+        Record<string, never>
+      > | null,
+      result: null as UseReadQueryResult<SimpleCaseData> | null,
+    },
+  });
+
+  const { SuspenseFallback, ReadQueryHook } =
+    createDefaultProfiledComponents(Profiler);
+
+  function App() {
+    useTrackRenders();
+    const [loadQuery, queryRef, { subscribeToMore }] = useLoadableQuery(query);
+
+    Profiler.mergeSnapshot({ subscribeToMore });
+
+    return (
+      <div>
+        <button onClick={() => loadQuery()}>Load query</button>
+        <Suspense fallback={<SuspenseFallback />}>
+          {queryRef && <ReadQueryHook queryRef={queryRef} />}
+        </Suspense>
+      </div>
+    );
+  }
+
+  const { user } = renderWithClient(<App />, { client, wrapper: Profiler });
+  // initial render
+  await Profiler.takeRender();
+
+  await act(() => user.click(screen.getByText("Load query")));
+
+  {
+    const { renderedComponents } = await Profiler.takeRender();
+
+    expect(renderedComponents).toStrictEqual([App, SuspenseFallback]);
+  }
+
+  {
+    const { renderedComponents, snapshot } = await Profiler.takeRender();
+
+    expect(renderedComponents).toStrictEqual([ReadQueryHook]);
+    expect(snapshot.result).toEqual({
+      data: { greeting: "Hello" },
+      error: undefined,
+      networkStatus: NetworkStatus.ready,
+    });
+  }
+
+  const updateQuery = jest.fn<
+    ReturnType<UpdateQueryFn>,
+    Parameters<UpdateQueryFn>
+  >((_, { subscriptionData: { data } }) => {
+    return { greeting: data.greetingUpdated };
+  });
+
+  const { snapshot } = Profiler.getCurrentRender();
+
+  snapshot.subscribeToMore!({ document: subscription, updateQuery });
+
+  wsLink.simulateResult({
+    result: {
+      data: {
+        greetingUpdated: "Subscription hello",
+      },
+    },
+  });
+
+  {
+    const { snapshot, renderedComponents } = await Profiler.takeRender();
+
+    expect(renderedComponents).toStrictEqual([ReadQueryHook]);
+    expect(snapshot.result).toEqual({
+      data: { greeting: "Subscription hello" },
+      error: undefined,
+      networkStatus: NetworkStatus.ready,
+    });
+  }
+
+  expect(updateQuery).toHaveBeenCalledTimes(1);
+  expect(updateQuery).toHaveBeenCalledWith(
+    { greeting: "Hello" },
+    {
+      subscriptionData: {
+        data: { greetingUpdated: "Subscription hello" },
+      },
+      variables: {},
+    }
+  );
+});
+
+it("throws when calling `subscribeToMore` before loading the query", async () => {
+  interface SubscriptionData {
+    greetingUpdated: string;
+  }
+
+  const subscription: TypedDocumentNode<
+    SubscriptionData,
+    Record<string, never>
+  > = gql`
+    subscription {
+      greetingUpdated
+    }
+  `;
+
+  const { mocks, query } = setupSimpleCase();
+
+  const wsLink = new MockSubscriptionLink();
+  const mockLink = new MockLink(mocks);
+
+  const link = split(
+    ({ query }) => {
+      const definition = getMainDefinition(query);
+
+      return (
+        definition.kind === "OperationDefinition" &&
+        definition.operation === "subscription"
+      );
+    },
+    wsLink,
+    mockLink
+  );
+
+  const client = new ApolloClient({ link, cache: new InMemoryCache() });
+
+  const Profiler = createProfiler({
+    initialSnapshot: {
+      subscribeToMore: null as SubscribeToMoreFunction<
+        SimpleCaseData,
+        Record<string, never>
+      > | null,
+      result: null as UseReadQueryResult<SimpleCaseData> | null,
+    },
+  });
+
+  const { SuspenseFallback, ReadQueryHook } =
+    createDefaultProfiledComponents(Profiler);
+
+  function App() {
+    useTrackRenders();
+    const [loadQuery, queryRef, { subscribeToMore }] = useLoadableQuery(query);
+
+    Profiler.mergeSnapshot({ subscribeToMore });
+
+    return (
+      <div>
+        <button onClick={() => loadQuery()}>Load query</button>
+        <Suspense fallback={<SuspenseFallback />}>
+          {queryRef && <ReadQueryHook queryRef={queryRef} />}
+        </Suspense>
+      </div>
+    );
+  }
+
+  renderWithClient(<App />, { client, wrapper: Profiler });
+  // initial render
+  await Profiler.takeRender();
+
+  const { snapshot } = Profiler.getCurrentRender();
+
+  expect(() => {
+    snapshot.subscribeToMore!({ document: subscription });
+  }).toThrow(
+    new InvariantError("The query has not been loaded. Please load the query.")
+  );
+});
+
 describe.skip("type tests", () => {
   it("returns unknown when TData cannot be inferred", () => {
     const query = gql``;
diff --git a/src/react/hooks/__tests__/useQueryRefHandlers.test.tsx b/src/react/hooks/__tests__/useQueryRefHandlers.test.tsx
--- a/src/react/hooks/__tests__/useQueryRefHandlers.test.tsx
+++ b/src/react/hooks/__tests__/useQueryRefHandlers.test.tsx
@@ -4,10 +4,16 @@ import {
   ApolloClient,
   InMemoryCache,
   NetworkStatus,
+  SubscribeToMoreOptions,
   TypedDocumentNode,
   gql,
+  split,
 } from "../../../core";
-import { MockLink, MockedResponse } from "../../../testing";
+import {
+  MockLink,
+  MockSubscriptionLink,
+  MockedResponse,
+} from "../../../testing";
 import {
   PaginatedCaseData,
   SimpleCaseData,
@@ -19,13 +25,14 @@ import {
 } from "../../../testing/internal";
 import { useQueryRefHandlers } from "../useQueryRefHandlers";
 import { UseReadQueryResult, useReadQuery } from "../useReadQuery";
+import type { SubscribeToMoreFunction } from "../useSuspenseQuery";
 import { Suspense } from "react";
 import { createQueryPreloader } from "../../query-preloader/createQueryPreloader";
 import userEvent from "@testing-library/user-event";
 import { QueryRef } from "../../internal";
 import { useBackgroundQuery } from "../useBackgroundQuery";
 import { useLoadableQuery } from "../useLoadableQuery";
-import { concatPagination } from "../../../utilities";
+import { concatPagination, getMainDefinition } from "../../../utilities";
 
 test("does not interfere with updates from useReadQuery", async () => {
   const { query, mocks } = setupSimpleCase();
@@ -1927,3 +1934,147 @@ test("`fetchMore` works with startTransition from useBackgroundQuery and useQuer
 
   await expect(Profiler).not.toRerender();
 });
+
+test("can subscribe to subscriptions and react to cache updates via `subscribeToMore`", async () => {
+  interface SubscriptionData {
+    greetingUpdated: string;
+  }
+
+  type UpdateQueryFn = NonNullable<
+    SubscribeToMoreOptions<
+      SimpleCaseData,
+      Record<string, never>,
+      SubscriptionData
+    >["updateQuery"]
+  >;
+
+  const subscription: TypedDocumentNode<
+    SubscriptionData,
+    Record<string, never>
+  > = gql`
+    subscription {
+      greetingUpdated
+    }
+  `;
+
+  const { mocks, query } = setupSimpleCase();
+
+  const wsLink = new MockSubscriptionLink();
+  const mockLink = new MockLink(mocks);
+
+  const link = split(
+    ({ query }) => {
+      const definition = getMainDefinition(query);
+
+      return (
+        definition.kind === "OperationDefinition" &&
+        definition.operation === "subscription"
+      );
+    },
+    wsLink,
+    mockLink
+  );
+
+  const client = new ApolloClient({ link, cache: new InMemoryCache() });
+
+  const preloadQuery = createQueryPreloader(client);
+  const queryRef = preloadQuery(query);
+
+  const Profiler = createProfiler({
+    initialSnapshot: {
+      subscribeToMore: null as SubscribeToMoreFunction<
+        SimpleCaseData,
+        Record<string, never>
+      > | null,
+      result: null as UseReadQueryResult<SimpleCaseData> | null,
+    },
+  });
+
+  function SuspenseFallback() {
+    useTrackRenders();
+    return <p>Loading</p>;
+  }
+
+  function ReadQueryHook() {
+    useTrackRenders();
+    Profiler.mergeSnapshot({ result: useReadQuery(queryRef) });
+
+    return null;
+  }
+
+  function App() {
+    useTrackRenders();
+    // We can ignore the return result here since we are testing the mechanics
+    // of this hook to ensure it doesn't interfere with the updates from
+    // useReadQuery
+    const { subscribeToMore } = useQueryRefHandlers(queryRef);
+
+    Profiler.mergeSnapshot({ subscribeToMore });
+
+    return (
+      <Suspense fallback={<SuspenseFallback />}>
+        <ReadQueryHook />
+      </Suspense>
+    );
+  }
+
+  renderWithClient(<App />, { client, wrapper: Profiler });
+
+  {
+    const { renderedComponents } = await Profiler.takeRender();
+
+    expect(renderedComponents).toStrictEqual([App, SuspenseFallback]);
+  }
+
+  {
+    const { snapshot, renderedComponents } = await Profiler.takeRender();
+
+    expect(renderedComponents).toStrictEqual([ReadQueryHook]);
+    expect(snapshot.result).toEqual({
+      data: { greeting: "Hello" },
+      error: undefined,
+      networkStatus: NetworkStatus.ready,
+    });
+  }
+
+  const updateQuery = jest.fn<
+    ReturnType<UpdateQueryFn>,
+    Parameters<UpdateQueryFn>
+  >((_, { subscriptionData: { data } }) => {
+    return { greeting: data.greetingUpdated };
+  });
+
+  const { snapshot } = Profiler.getCurrentRender();
+
+  snapshot.subscribeToMore!({ document: subscription, updateQuery });
+
+  wsLink.simulateResult({
+    result: {
+      data: {
+        greetingUpdated: "Subscription hello",
+      },
+    },
+  });
+
+  {
+    const { snapshot, renderedComponents } = await Profiler.takeRender();
+
+    expect(renderedComponents).toStrictEqual([ReadQueryHook]);
+    expect(snapshot.result).toEqual({
+      data: { greeting: "Subscription hello" },
+      error: undefined,
+      networkStatus: NetworkStatus.ready,
+    });
+  }
+
+  expect(updateQuery).toHaveBeenCalledTimes(1);
+  expect(updateQuery).toHaveBeenCalledWith(
+    { greeting: "Hello" },
+    {
+      subscriptionData: {
+        data: { greetingUpdated: "Subscription hello" },
+      },
+      variables: {},
+    }
+  );
+});