From ef27883bb6343e0286e5e4f547e6dee7827fcf9c Mon Sep 17 00:00:00 2001
From: Jerel Miller <jerelmiller@gmail.com>
Date: Thu, 19 Dec 2024 12:39:35 -0700
Subject: [PATCH] Add tests to ensure suspense fragment is torn down

---
 .../__tests__/useSuspenseFragment.test.tsx    | 112 +++++++++++++++++-
 1 file changed, 111 insertions(+), 1 deletion(-)

diff --git a/src/react/hooks/__tests__/useSuspenseFragment.test.tsx b/src/react/hooks/__tests__/useSuspenseFragment.test.tsx
index d4ef67097a..7f9a54545b 100644
--- a/src/react/hooks/__tests__/useSuspenseFragment.test.tsx
+++ b/src/react/hooks/__tests__/useSuspenseFragment.test.tsx
@@ -23,7 +23,7 @@ import {
 import { spyOnConsole } from "../../../testing/internal";
 import { renderHook } from "@testing-library/react";
 import { InvariantError } from "ts-invariant";
-import { MockedProvider } from "../../../testing";
+import { MockedProvider, wait } from "../../../testing";
 import { expectTypeOf } from "expect-type";
 
 function createDefaultRenderStream<TData = unknown>() {
@@ -1456,6 +1456,116 @@ test("updates child fragments for cache updates to masked fields", async () => {
   await expect(takeRender).not.toRerender();
 });
 
+test("tears down the subscription on unmount", async () => {
+  interface ItemFragment {
+    __typename: "Item";
+    id: number;
+    text: string;
+  }
+
+  const fragment: TypedDocumentNode<ItemFragment> = gql`
+    fragment ItemFragment on Item {
+      id
+      text
+    }
+  `;
+
+  const cache = new InMemoryCache();
+  const client = new ApolloClient({ cache });
+
+  client.writeFragment({
+    fragment,
+    data: { __typename: "Item", id: 1, text: "Item #1" },
+  });
+
+  using _disabledAct = disableActEnvironment();
+  const { unmount, takeSnapshot } = await renderHookToSnapshotStream(
+    () =>
+      useSuspenseFragment({ fragment, from: { __typename: "Item", id: 1 } }),
+    {
+      wrapper: ({ children }) => (
+        <ApolloProvider client={client}>{children}</ApolloProvider>
+      ),
+    }
+  );
+
+  {
+    const { data } = await takeSnapshot();
+
+    expect(data).toEqual({ __typename: "Item", id: 1, text: "Item #1" });
+  }
+
+  expect(cache["watches"].size).toBe(1);
+
+  unmount();
+  // We need to wait a tick since the cleanup is run in a setTimeout to
+  // prevent strict mode bugs.
+  await wait(0);
+
+  expect(cache["watches"].size).toBe(0);
+});
+
+test("tears down all watches when rendering multiple records", async () => {
+  interface ItemFragment {
+    __typename: "Item";
+    id: number;
+    text: string;
+  }
+
+  const fragment: TypedDocumentNode<ItemFragment> = gql`
+    fragment ItemFragment on Item {
+      id
+      text
+    }
+  `;
+
+  const cache = new InMemoryCache();
+  const client = new ApolloClient({ cache });
+
+  client.writeFragment({
+    fragment,
+    data: { __typename: "Item", id: 1, text: "Item #1" },
+  });
+
+  client.writeFragment({
+    fragment,
+    data: { __typename: "Item", id: 2, text: "Item #2" },
+  });
+
+  using _disabledAct = disableActEnvironment();
+  const { unmount, rerender, takeSnapshot } = await renderHookToSnapshotStream(
+    ({ id }) =>
+      useSuspenseFragment({ fragment, from: { __typename: "Item", id } }),
+    {
+      initialProps: { id: 1 },
+      wrapper: ({ children }) => (
+        <ApolloProvider client={client}>{children}</ApolloProvider>
+      ),
+    }
+  );
+
+  {
+    const { data } = await takeSnapshot();
+
+    expect(data).toEqual({ __typename: "Item", id: 1, text: "Item #1" });
+  }
+
+  await rerender({ id: 2 });
+
+  {
+    const { data } = await takeSnapshot();
+
+    expect(data).toEqual({ __typename: "Item", id: 2, text: "Item #2" });
+  }
+
+  unmount();
+  // We need to wait a tick since the cleanup is run in a setTimeout to
+  // prevent strict mode bugs.
+  await wait(0);
+
+  expect(cache["watches"].size).toBe(0);
+});
+
 describe.skip("type tests", () => {
   test("returns TData when from is a non-null value", () => {
     const fragment: TypedDocumentNode<{ foo: string }> = gql``;