From bb1750986e17250b3454b630812cd5568d4a1f77 Mon Sep 17 00:00:00 2001
From: Jerel Miller <jerelmiller@gmail.com>
Date: Tue, 21 Nov 2023 12:17:13 -0700
Subject: [PATCH] Improvements to profiler testing utility (#11376)

---
 .../__tests__/client/Query.test.tsx           |  2 +-
 .../hoc/__tests__/queries/lifecycle.test.tsx  |  2 +-
 .../hoc/__tests__/queries/loading.test.tsx    |  2 +-
 .../__tests__/useBackgroundQuery.test.tsx     | 19 ++--
 .../hooks/__tests__/useFragment.test.tsx      |  2 +-
 .../hooks/__tests__/useLazyQuery.test.tsx     |  2 +-
 .../hooks/__tests__/useSuspenseQuery.test.tsx |  2 +-
 src/testing/internal/profile/profile.tsx      | 96 ++++++++++++-------
 src/testing/matchers/ProfiledComponent.ts     |  6 +-
 9 files changed, 76 insertions(+), 57 deletions(-)

diff --git a/src/react/components/__tests__/client/Query.test.tsx b/src/react/components/__tests__/client/Query.test.tsx
index 3aee25d2685..acdd2015301 100644
--- a/src/react/components/__tests__/client/Query.test.tsx
+++ b/src/react/components/__tests__/client/Query.test.tsx
@@ -1491,7 +1491,7 @@ describe("Query component", () => {
       return (
         <AllPeopleQuery2 query={query} notifyOnNetworkStatusChange={true}>
           {(r: any) => {
-            ProfiledContainer.updateSnapshot(r);
+            ProfiledContainer.replaceSnapshot(r);
             return null;
           }}
         </AllPeopleQuery2>
diff --git a/src/react/hoc/__tests__/queries/lifecycle.test.tsx b/src/react/hoc/__tests__/queries/lifecycle.test.tsx
index 99e7fbc6c5e..cf460af964a 100644
--- a/src/react/hoc/__tests__/queries/lifecycle.test.tsx
+++ b/src/react/hoc/__tests__/queries/lifecycle.test.tsx
@@ -52,7 +52,7 @@ describe("[queries] lifecycle", () => {
     })(
       class extends React.Component<ChildProps<Vars, Data, Vars>> {
         render() {
-          ProfiledApp.updateSnapshot(this.props.data!);
+          ProfiledApp.replaceSnapshot(this.props.data!);
           return null;
         }
       }
diff --git a/src/react/hoc/__tests__/queries/loading.test.tsx b/src/react/hoc/__tests__/queries/loading.test.tsx
index c2a40c0b9ec..387a6803fb5 100644
--- a/src/react/hoc/__tests__/queries/loading.test.tsx
+++ b/src/react/hoc/__tests__/queries/loading.test.tsx
@@ -407,7 +407,7 @@ describe("[queries] loading", () => {
     })(
       class extends React.Component<ChildProps<{}, Data>> {
         render() {
-          ProfiledContainer.updateSnapshot(this.props.data!);
+          ProfiledContainer.replaceSnapshot(this.props.data!);
           return null;
         }
       }
diff --git a/src/react/hooks/__tests__/useBackgroundQuery.test.tsx b/src/react/hooks/__tests__/useBackgroundQuery.test.tsx
index 08121af7f27..131364939cd 100644
--- a/src/react/hooks/__tests__/useBackgroundQuery.test.tsx
+++ b/src/react/hooks/__tests__/useBackgroundQuery.test.tsx
@@ -335,7 +335,7 @@ function renderVariablesIntegrationTest({
   const ProfiledApp = profile<Renders, ComponentProps<typeof App>>({
     Component: App,
     snapshotDOM: true,
-    onRender: ({ updateSnapshot }) => updateSnapshot(cloneDeep(renders)),
+    onRender: ({ replaceSnapshot }) => replaceSnapshot(cloneDeep(renders)),
   });
 
   const { ...rest } = render(
@@ -434,9 +434,8 @@ function renderPaginatedIntegrationTest({
   }
 
   function SuspenseFallback() {
-    ProfiledApp.updateSnapshot((snapshot) => ({
-      ...snapshot,
-      suspenseCount: snapshot.suspenseCount + 1,
+    ProfiledApp.mergeSnapshot(({ suspenseCount }) => ({
+      suspenseCount: suspenseCount + 1,
     }));
     return <div>loading</div>;
   }
@@ -450,9 +449,8 @@ function renderPaginatedIntegrationTest({
   }) {
     const { data, error } = useReadQuery(queryRef);
     // count renders in the child component
-    ProfiledApp.updateSnapshot((snapshot) => ({
-      ...snapshot,
-      count: snapshot.count + 1,
+    ProfiledApp.mergeSnapshot(({ count }) => ({
+      count: count + 1,
     }));
     return (
       <div>
@@ -504,10 +502,9 @@ function renderPaginatedIntegrationTest({
         <ErrorBoundary
           fallback={<div>Error</div>}
           onError={(error) => {
-            ProfiledApp.updateSnapshot((snapshot) => ({
-              ...snapshot,
-              errorCount: snapshot.errorCount + 1,
-              errors: snapshot.errors.concat(error),
+            ProfiledApp.mergeSnapshot(({ errorCount, errors }) => ({
+              errorCount: errorCount + 1,
+              errors: errors.concat(error),
             }));
           }}
         >
diff --git a/src/react/hooks/__tests__/useFragment.test.tsx b/src/react/hooks/__tests__/useFragment.test.tsx
index 6ef04ad4d01..21b9e083a03 100644
--- a/src/react/hooks/__tests__/useFragment.test.tsx
+++ b/src/react/hooks/__tests__/useFragment.test.tsx
@@ -1476,7 +1476,7 @@ describe("has the same timing as `useQuery`", () => {
         from: initialItem,
       });
 
-      ProfiledComponent.updateSnapshot({ queryData, fragmentData });
+      ProfiledComponent.replaceSnapshot({ queryData, fragmentData });
 
       return complete ? JSON.stringify(fragmentData) : "loading";
     }
diff --git a/src/react/hooks/__tests__/useLazyQuery.test.tsx b/src/react/hooks/__tests__/useLazyQuery.test.tsx
index a36c0725bea..150ce125fca 100644
--- a/src/react/hooks/__tests__/useLazyQuery.test.tsx
+++ b/src/react/hooks/__tests__/useLazyQuery.test.tsx
@@ -1115,7 +1115,7 @@ describe("useLazyQuery Hook", () => {
       ),
     });
 
-    const [execute] = ProfiledHook.getCurrentSnapshot();
+    const [execute] = await ProfiledHook.peekSnapshot();
 
     {
       const [, result] = await ProfiledHook.takeSnapshot();
diff --git a/src/react/hooks/__tests__/useSuspenseQuery.test.tsx b/src/react/hooks/__tests__/useSuspenseQuery.test.tsx
index 468fe8f4fc5..642be7d023a 100644
--- a/src/react/hooks/__tests__/useSuspenseQuery.test.tsx
+++ b/src/react/hooks/__tests__/useSuspenseQuery.test.tsx
@@ -357,7 +357,7 @@ describe("useSuspenseQuery", () => {
 
     const Component = () => {
       const result = useSuspenseQuery(query);
-      ProfiledApp.updateSnapshot(result);
+      ProfiledApp.replaceSnapshot(result);
       return <div>{result.data.greeting}</div>;
     };
 
diff --git a/src/testing/internal/profile/profile.tsx b/src/testing/internal/profile/profile.tsx
index f4abb64af0e..8ae43b64c01 100644
--- a/src/testing/internal/profile/profile.tsx
+++ b/src/testing/internal/profile/profile.tsx
@@ -23,15 +23,27 @@ export interface NextRenderOptions {
 export interface ProfiledComponent<Props, Snapshot>
   extends React.FC<Props>,
     ProfiledComponentFields<Props, Snapshot>,
-    ProfiledComponenOnlyFields<Props, Snapshot> {}
+    ProfiledComponentOnlyFields<Props, Snapshot> {}
 
-interface UpdateSnapshot<Snapshot> {
+interface ReplaceSnapshot<Snapshot> {
   (newSnapshot: Snapshot): void;
   (updateSnapshot: (lastSnapshot: Readonly<Snapshot>) => Snapshot): void;
 }
 
-interface ProfiledComponenOnlyFields<Props, Snapshot> {
-  updateSnapshot: UpdateSnapshot<Snapshot>;
+interface MergeSnapshot<Snapshot> {
+  (partialSnapshot: Partial<Snapshot>): void;
+  (
+    updatePartialSnapshot: (
+      lastSnapshot: Readonly<Snapshot>
+    ) => Partial<Snapshot>
+  ): void;
+}
+
+interface ProfiledComponentOnlyFields<Props, Snapshot> {
+  // Allows for partial updating of the snapshot by shallow merging the results
+  mergeSnapshot: MergeSnapshot<Snapshot>;
+  // Performs a full replacement of the snapshot
+  replaceSnapshot: ReplaceSnapshot<Snapshot>;
 }
 interface ProfiledComponentFields<Props, Snapshot> {
   /**
@@ -54,21 +66,14 @@ interface ProfiledComponentFields<Props, Snapshot> {
    */
   takeRender(options?: NextRenderOptions): Promise<Render<Snapshot>>;
   /**
-   * Returns the current render count.
+   * Returns the total number of renders.
    */
-  currentRenderCount(): number;
+  totalRenderCount(): number;
   /**
    * Returns the current render.
    * @throws {Error} if no render has happened yet
    */
   getCurrentRender(): Render<Snapshot>;
-  /**
-   * Iterates the renders until the render count is reached.
-   */
-  takeUntilRenderCount(
-    count: number,
-    optionsPerRender?: NextRenderOptions
-  ): Promise<void>;
   /**
    * Waits for the next render to happen.
    * Does not advance the render iterator.
@@ -90,18 +95,18 @@ export function profile<
   onRender?: (
     info: BaseRender & {
       snapshot: Snapshot;
-      updateSnapshot: UpdateSnapshot<Snapshot>;
+      replaceSnapshot: ReplaceSnapshot<Snapshot>;
+      mergeSnapshot: MergeSnapshot<Snapshot>;
     }
   ) => void;
   snapshotDOM?: boolean;
   initialSnapshot?: Snapshot;
 }) {
-  let currentRender: Render<Snapshot> | undefined;
   let nextRender: Promise<Render<Snapshot>> | undefined;
   let resolveNextRender: ((render: Render<Snapshot>) => void) | undefined;
   let rejectNextRender: ((error: unknown) => void) | undefined;
   const snapshotRef = { current: initialSnapshot };
-  const updateSnapshot: UpdateSnapshot<Snapshot> = (snap) => {
+  const replaceSnapshot: ReplaceSnapshot<Snapshot> = (snap) => {
     if (typeof snap === "function") {
       if (!initialSnapshot) {
         throw new Error(
@@ -118,6 +123,16 @@ export function profile<
       snapshotRef.current = snap;
     }
   };
+
+  const mergeSnapshot: MergeSnapshot<Snapshot> = (partialSnapshot) => {
+    replaceSnapshot((snapshot) => ({
+      ...snapshot,
+      ...(typeof partialSnapshot === "function"
+        ? partialSnapshot(snapshot)
+        : partialSnapshot),
+    }));
+  };
+
   const profilerOnRender: React.ProfilerOnRenderCallback = (
     id,
     phase,
@@ -145,7 +160,8 @@ export function profile<
        */
       onRender?.({
         ...baseRender,
-        updateSnapshot,
+        replaceSnapshot,
+        mergeSnapshot,
         snapshot: snapshotRef.current!,
       });
 
@@ -154,8 +170,6 @@ export function profile<
         ? window.document.body.innerHTML
         : undefined;
       const render = new RenderInstance(baseRender, snapshot, domSnapshot);
-      // eslint-disable-next-line testing-library/render-result-naming-convention
-      currentRender = render;
       Profiled.renders.push(render);
       resolveNextRender?.(render);
     } catch (error) {
@@ -178,29 +192,31 @@ export function profile<
       </React.Profiler>
     ),
     {
-      updateSnapshot,
-    } satisfies ProfiledComponenOnlyFields<Props, Snapshot>,
+      replaceSnapshot,
+      mergeSnapshot,
+    } satisfies ProfiledComponentOnlyFields<Props, Snapshot>,
     {
       renders: new Array<
         | Render<Snapshot>
         | { phase: "snapshotError"; count: number; error: unknown }
       >(),
-      currentRenderCount() {
+      totalRenderCount() {
         return Profiled.renders.length;
       },
       async peekRender(options: NextRenderOptions = {}) {
         if (iteratorPosition < Profiled.renders.length) {
           const render = Profiled.renders[iteratorPosition];
+
           if (render.phase === "snapshotError") {
             throw render.error;
           }
+
           return render;
         }
-        const render = Profiled.waitForNextRender({
+        return Profiled.waitForNextRender({
           [_stackTrace]: captureStackTrace(Profiled.peekRender),
           ...options,
         });
-        return render;
       },
       async takeRender(options: NextRenderOptions = {}) {
         let error: unknown = undefined;
@@ -219,18 +235,25 @@ export function profile<
         }
       },
       getCurrentRender() {
-        if (!currentRender) {
-          throw new Error("Has not been rendered yet!");
+        // The "current" render should point at the same render that the most
+        // recent `takeRender` call returned, so we need to get the "previous"
+        // iterator position, otherwise `takeRender` advances the iterator
+        // to the next render. This means we need to call `takeRender` at least
+        // once before we can get a current render.
+        const currentPosition = iteratorPosition - 1;
+
+        if (currentPosition < 0) {
+          throw new Error(
+            "No current render available. You need to call `takeRender` before you can get the current render."
+          );
         }
-        return currentRender;
-      },
-      async takeUntilRenderCount(
-        count: number,
-        optionsPerRender?: NextRenderOptions
-      ) {
-        while (Profiled.renders.length < count) {
-          await Profiled.takeRender(optionsPerRender);
+
+        const render = Profiled.renders[currentPosition];
+
+        if (render.phase === "snapshotError") {
+          throw render.error;
         }
+        return render;
       },
       waitForNextRender({
         timeout = 1000,
@@ -306,7 +329,7 @@ export function profileHook<ReturnValue extends ValidSnapshot, Props>(
 ): ProfiledHook<Props, ReturnValue> {
   let returnValue: ReturnValue;
   const Component = (props: Props) => {
-    ProfiledComponent.updateSnapshot(renderCallback(props));
+    ProfiledComponent.replaceSnapshot(renderCallback(props));
     return null;
   };
   const ProfiledComponent = profile<ReturnValue, Props>({
@@ -322,7 +345,7 @@ export function profileHook<ReturnValue extends ValidSnapshot, Props>(
     },
     {
       renders: ProfiledComponent.renders,
-      currentSnapshotCount: ProfiledComponent.currentRenderCount,
+      totalSnapshotCount: ProfiledComponent.totalRenderCount,
       async peekSnapshot(options) {
         return (await ProfiledComponent.peekRender(options)).snapshot;
       },
@@ -332,7 +355,6 @@ export function profileHook<ReturnValue extends ValidSnapshot, Props>(
       getCurrentSnapshot() {
         return ProfiledComponent.getCurrentRender().snapshot;
       },
-      takeUntilSnapshotCount: ProfiledComponent.takeUntilRenderCount,
       async waitForNextSnapshot(options) {
         return (await ProfiledComponent.waitForNextRender(options)).snapshot;
       },
diff --git a/src/testing/matchers/ProfiledComponent.ts b/src/testing/matchers/ProfiledComponent.ts
index 469cfe00995..8a4e72025a9 100644
--- a/src/testing/matchers/ProfiledComponent.ts
+++ b/src/testing/matchers/ProfiledComponent.ts
@@ -52,11 +52,11 @@ export const toRenderExactlyTimes: MatcherFunction<
   const hint = this.utils.matcherHint("toRenderExactlyTimes");
   let pass = true;
   try {
-    if (profiled.currentRenderCount() > times) {
+    if (profiled.totalRenderCount() > times) {
       throw failed;
     }
     try {
-      while (profiled.currentRenderCount() < times) {
+      while (profiled.totalRenderCount() < times) {
         await profiled.waitForNextRender(options);
       }
     } catch (e) {
@@ -84,7 +84,7 @@ export const toRenderExactlyTimes: MatcherFunction<
       return (
         hint +
         ` Expected component to${pass ? " not" : ""} render exactly ${times}.` +
-        ` It rendered ${profiled.currentRenderCount()} times.`
+        ` It rendered ${profiled.totalRenderCount()} times.`
       );
     },
   };