From 55cc6068c31386afdb82e679c9b8a9de7eb8d95c Mon Sep 17 00:00:00 2001
From: Szymon Chmal <szymonchmal@gmail.com>
Date: Thu, 28 Apr 2022 03:58:11 +0200
Subject: [PATCH] fix(use-masonry): should update when positioner changes
 (#113)

---
 src/index.test.tsx  | 54 +++++++++++++++++++++++++++++++++++++++++++++
 src/use-masonry.tsx |  2 +-
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/src/index.test.tsx b/src/index.test.tsx
index 8127620..1da1045 100644
--- a/src/index.test.tsx
+++ b/src/index.test.tsx
@@ -1,6 +1,7 @@
 import { dimension } from "@shopify/jest-dom-mocks";
 import { act, render, screen } from "@testing-library/react";
 import { act as hookAct, renderHook } from "@testing-library/react-hooks";
+import type { RenderHookResult } from "@testing-library/react-hooks";
 import * as React from "react";
 import {
   createResizeObserver,
@@ -14,6 +15,7 @@ import {
   useResizeObserver,
   useScroller,
 } from "./index";
+import * as useForceUpdateModule from "./use-force-update";
 
 jest.useFakeTimers();
 
@@ -50,6 +52,7 @@ afterEach(() => {
   resetScroll();
   dimension.restore();
   perf.uninstall();
+  jest.restoreAllMocks();
 });
 
 describe("useMasonry()", () => {
@@ -312,6 +315,57 @@ describe("useMasonry()", () => {
     const { result } = renderBasicMasonry({ className: "foo" });
     expect(result.current.props.className).toEqual("foo");
   });
+
+  it('should render multiple batches if "itemHeightEstimate" isn\'t accurate', () => {
+    // eslint-disable-next-line prefer-const
+    let hook: RenderHookResult<
+      { items: { id: number; height: number }[] },
+      JSX.Element
+    >;
+
+    // Render hook again on useForceUpdate
+    jest.spyOn(useForceUpdateModule, "useForceUpdate").mockReturnValue(() => {
+      if (hook) {
+        hook.rerender();
+        render(hook.result.current);
+      }
+    });
+
+    // Render hook with items-dependent positioner
+    hook = renderHook(
+      (props) => {
+        const positioner = usePositioner({ width: 1280, columnWidth: 1280 }, [
+          props.items[0],
+        ]);
+        return useMasonry({
+          height: 1280,
+          positioner,
+          itemHeightEstimate: 640,
+          overscanBy: 1,
+          render: FakeCard,
+          scrollTop: 0,
+          ...props,
+        });
+      },
+      {
+        initialProps: {
+          items: getFakeItems(100, 80),
+        },
+      }
+    );
+
+    // Switch items, positioner will update itself
+    hook.rerender({
+      items: getFakeItems(100, 80),
+    });
+
+    // All items should have measured styles
+    for (let i = 0; i < 2; i++) {
+      expect(hook.result.current.props.children[i].props.style).not.toEqual(
+        prerenderStyles(1280)
+      );
+    }
+  });
 });
 
 describe("usePositioner()", () => {
diff --git a/src/use-masonry.tsx b/src/use-masonry.tsx
index a0f9ade..8e02cc3 100644
--- a/src/use-masonry.tsx
+++ b/src/use-masonry.tsx
@@ -187,7 +187,7 @@ export function useMasonry<Item>({
   React.useEffect(() => {
     if (needsFreshBatch) forceUpdate();
     // eslint-disable-next-line
-  }, [needsFreshBatch]);
+  }, [needsFreshBatch, positioner]);
 
   // gets the container style object based upon the estimated height and whether or not
   // the page is being scrolled