From 5d8946728d246126bb2af6854cc7669c831deed5 Mon Sep 17 00:00:00 2001
From: Lenz Weber-Tronic <lorenz.weber-tronic@apollographql.com>
Date: Tue, 21 Nov 2023 11:23:02 +0100
Subject: [PATCH] suggestion: use `immer` in `updateSnapshot`

---
 package-lock.json                             | 11 ++++++
 package.json                                  |  1 +
 .../__tests__/useBackgroundQuery.test.tsx     | 20 +++++-----
 src/testing/internal/profile/profile.tsx      | 37 ++++---------------
 4 files changed, 30 insertions(+), 39 deletions(-)

diff --git a/package-lock.json b/package-lock.json
index dbd669f34a2..f178204f214 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -69,6 +69,7 @@
         "glob": "8.1.0",
         "graphql": "16.8.1",
         "graphql-ws": "5.14.2",
+        "immer": "^10.0.3",
         "jest": "29.7.0",
         "jest-environment-jsdom": "29.7.0",
         "jest-junit": "16.0.0",
@@ -6356,6 +6357,16 @@
         "node": ">= 4"
       }
     },
+    "node_modules/immer": {
+      "version": "10.0.3",
+      "resolved": "https://registry.npmjs.org/immer/-/immer-10.0.3.tgz",
+      "integrity": "sha512-pwupu3eWfouuaowscykeckFmVTpqbzW+rXFCX8rQLkZzM9ftBmU/++Ra+o+L27mz03zJTlyV4UUr+fdKNffo4A==",
+      "dev": true,
+      "funding": {
+        "type": "opencollective",
+        "url": "https://opencollective.com/immer"
+      }
+    },
     "node_modules/import-fresh": {
       "version": "3.3.0",
       "resolved": "https://registry.npmjs.org/import-fresh/-/import-fresh-3.3.0.tgz",
diff --git a/package.json b/package.json
index ec11352a8b6..efcab6b19fa 100644
--- a/package.json
+++ b/package.json
@@ -148,6 +148,7 @@
     "glob": "8.1.0",
     "graphql": "16.8.1",
     "graphql-ws": "5.14.2",
+    "immer": "^10.0.3",
     "jest": "29.7.0",
     "jest-environment-jsdom": "29.7.0",
     "jest-junit": "16.0.0",
diff --git a/src/react/hooks/__tests__/useBackgroundQuery.test.tsx b/src/react/hooks/__tests__/useBackgroundQuery.test.tsx
index d2ee0109c86..b184a3c385a 100644
--- a/src/react/hooks/__tests__/useBackgroundQuery.test.tsx
+++ b/src/react/hooks/__tests__/useBackgroundQuery.test.tsx
@@ -434,9 +434,9 @@ function renderPaginatedIntegrationTest({
   }
 
   function SuspenseFallback() {
-    ProfiledApp.setSnapshot(({ suspenseCount }) => ({
-      suspenseCount: suspenseCount + 1,
-    }));
+    ProfiledApp.updateSnapshot((snapshot) => {
+      snapshot.suspenseCount++;
+    });
     return <div>loading</div>;
   }
 
@@ -449,9 +449,9 @@ function renderPaginatedIntegrationTest({
   }) {
     const { data, error } = useReadQuery(queryRef);
     // count renders in the child component
-    ProfiledApp.setSnapshot(({ count }) => ({
-      count: count + 1,
-    }));
+    ProfiledApp.updateSnapshot((snapshot) => {
+      snapshot.count++;
+    });
     return (
       <div>
         {error ? <div>{error.message}</div> : null}
@@ -502,10 +502,10 @@ function renderPaginatedIntegrationTest({
         <ErrorBoundary
           fallback={<div>Error</div>}
           onError={(error) => {
-            ProfiledApp.setSnapshot(({ errorCount, errors }) => ({
-              errorCount: errorCount + 1,
-              errors: errors.concat(error),
-            }));
+            ProfiledApp.updateSnapshot((snapshot) => {
+              snapshot.errorCount++;
+              snapshot.errors.push(error);
+            });
           }}
         >
           <Suspense fallback={<SuspenseFallback />}>
diff --git a/src/testing/internal/profile/profile.tsx b/src/testing/internal/profile/profile.tsx
index 8645f0b281c..cdcc3d6344c 100644
--- a/src/testing/internal/profile/profile.tsx
+++ b/src/testing/internal/profile/profile.tsx
@@ -1,5 +1,6 @@
 import * as React from "react";
-
+import type { Draft } from "immer";
+import { produce } from "immer";
 import { TextEncoder, TextDecoder } from "util";
 
 global.TextEncoder ??= TextEncoder;
@@ -27,21 +28,10 @@ export interface ProfiledComponent<Props, Snapshot>
 
 interface UpdateSnapshot<Snapshot> {
   (newSnapshot: Snapshot): void;
-  (updateSnapshot: (lastSnapshot: Readonly<Snapshot>) => Snapshot): void;
-}
-
-interface SetSnapshot<Snapshot> {
-  (partialSnapshot: Partial<Snapshot>): void;
-  (
-    updatePartialSnapshot: (
-      lastSnapshot: Readonly<Snapshot>
-    ) => Partial<Snapshot>
-  ): void;
+  (updateSnapshot: (draft: Draft<Snapshot>) => Snapshot | void): void;
 }
 
 interface ProfiledComponentOnlyFields<Props, Snapshot> {
-  // Allows for partial updating of the snapshot by shallow merging the results
-  setSnapshot: SetSnapshot<Snapshot>;
   // Performs a full replacement of the snapshot
   updateSnapshot: UpdateSnapshot<Snapshot>;
 }
@@ -112,26 +102,16 @@ export function profile<
           "Cannot use a function to update the snapshot if no initial snapshot was provided."
         );
       }
-      snapshotRef.current = snap(
-        typeof snapshotRef.current === "object"
-          ? // "cheap best effort" to prevent accidental mutation of the last snapshot
-            { ...snapshotRef.current! }
-          : snapshotRef.current!
+      snapshotRef.current = produce<Snapshot>(
+        snapshotRef.current!,
+        snap as any
       );
     } else {
-      snapshotRef.current = snap;
+      // move the snapshot through `produce` to autofreeze it
+      snapshotRef.current = produce<Snapshot>(snap, (x) => x);
     }
   };
 
-  const setSnapshot: SetSnapshot<Snapshot> = (partialSnapshot) => {
-    updateSnapshot((snapshot) => ({
-      ...snapshot,
-      ...(typeof partialSnapshot === "function"
-        ? partialSnapshot(snapshot)
-        : partialSnapshot),
-    }));
-  };
-
   const profilerOnRender: React.ProfilerOnRenderCallback = (
     id,
     phase,
@@ -191,7 +171,6 @@ export function profile<
     ),
     {
       updateSnapshot,
-      setSnapshot,
     } satisfies ProfiledComponentOnlyFields<Props, Snapshot>,
     {
       renders: new Array<