UI: Fix useTreeModel returning stale model for the given model sour…

…ce (#3761) (#3762)

(cherry picked from commit d118d8c)

Co-authored-by: Grigas <35135765+grigasp@users.noreply.github.com>

common/changes/@itwin/components-react/ui-tree-fix-out-of-sync-tree-model_2022-06-08-13-48.json

@@ -0,0 +1,10 @@
+{
+  "changes": [
+    {
+      "packageName": "@itwin/components-react",
+      "comment": "Fix `useTreeModel` returning stale model for the given model source",
+      "type": "none"
+    }
+  ],
+  "packageName": "@itwin/components-react"
+}

ui/components-react/src/components-react/tree/controlled/TreeHooks.ts

@@ -6,8 +6,9 @@
  * @module Tree
  */
 
-import { useCallback, useEffect, useMemo, useState } from "react";
+import { useCallback, useEffect, useMemo } from "react";
 import { useDisposable } from "@itwin/core-react";
+import { useRerender } from "../../common/UseRerender";
 import { TreeDataProvider } from "../TreeDataProvider";
 import { TreeEventHandler, TreeEventHandlerParams } from "./TreeEventHandler";
 import { TreeModel } from "./TreeModel";
@@ -21,17 +22,11 @@ import { PagedTreeNodeLoader, TreeNodeLoader } from "./TreeNodeLoader";
  * @public
  */
 export function useTreeModel(modelSource: TreeModelSource): TreeModel {
-  const [model, setModel] = useState(modelSource.getModel());
-
+  const { rerender } = useRerender();
   useEffect(() => {
-    const onModelChanged = () => {
-      setModel(modelSource.getModel());
-    };
-    onModelChanged();
-    return modelSource.onModelChanged.addListener(onModelChanged);
-  }, [modelSource]);
-
-  return model;
+    return modelSource.onModelChanged.addListener(rerender);
+  }, [modelSource, rerender]);
+  return modelSource.getModel();
 }
 
 /**

ui/components-react/src/test/tree/controlled/TreeHooks.test.tsx

@@ -35,27 +35,32 @@ describe("useTreeModel", () => {
       (props: { modelSource: TreeModelSource }) => useTreeModel(props.modelSource),
       { initialProps: { modelSource: modelSourceMock.object } },
     );
-
-    expect(result.current).to.not.be.undefined;
+    expect(result.all.length).to.eq(1);
+    expect(result.all[0]).to.eq(testModel);
     expect(spy).to.have.been.calledOnce;
   });
 
   it("resubscribes to onModelChangeEvent when model source changes", () => {
     const firstModelEventAddSpy = sinon.spy(onModelChangeEvent, "addListener");
     const firstModelEventRemoveSpy = sinon.spy(onModelChangeEvent, "removeListener");
-    const { rerender } = renderHook(
+    const { result, rerender } = renderHook(
       (props: { modelSource: TreeModelSource }) => useTreeModel(props.modelSource),
       { initialProps: { modelSource: modelSourceMock.object } },
     );
+    expect(result.all.length).to.eq(1);
+    expect(result.all[0]).to.eq(testModel);
     expect(firstModelEventAddSpy).to.have.been.calledOnce;
 
     const newOnModelChangeEvent = new BeUiEvent<[TreeModel, TreeModelChanges]>();
     const newModelEventAddSpy = sinon.spy(newOnModelChangeEvent, "addListener");
+    const newTestModel = new MutableTreeModel();
     const newModelSourceMock = moq.Mock.ofType<TreeModelSource>();
     newModelSourceMock.setup((x) => x.onModelChanged).returns(() => newOnModelChangeEvent);
+    newModelSourceMock.setup((x) => x.getModel()).returns(() => newTestModel);
 
     rerender({ modelSource: newModelSourceMock.object });
-
+    expect(result.all.length).to.eq(2);
+    expect(result.all[1]).to.eq(newTestModel);
     expect(firstModelEventRemoveSpy).to.have.been.calledOnce;
     expect(newModelEventAddSpy).to.have.been.calledOnce;
   });