Presentation: Fix hierarchy auto update

common/api/presentation-components.api.md

@@ -571,7 +571,7 @@ export abstract class PropertyRecordsBuilder implements IContentVisitor {
 // @internal (undocumented)
 export interface ReloadedHierarchyPart {
     // (undocumented)
-    nodeItems: TreeNodeItem[];
+    nodeItems: DelayLoadedTreeNodeItem[];
     // (undocumented)
     offset: number;
     // (undocumented)

common/changes/@itwin/presentation-components/presentation-fix_hierarchy_auto_update_2022-05-17-10-19.json

@@ -0,0 +1,10 @@
+{
+  "changes": [
+    {
+      "packageName": "@itwin/presentation-components",
+      "comment": "Fix: set correct node's children count after hierarchy update",
+      "type": "none"
+    }
+  ],
+  "packageName": "@itwin/presentation-components"
+}

presentation/components/src/presentation-components/tree/controlled/TreeHooks.ts

@@ -8,12 +8,12 @@
 
 import * as React from "react";
 import { Subscription } from "rxjs/internal/Subscription";
-import { HierarchyUpdateRecord, PageOptions, UPDATE_FULL } from "@itwin/presentation-common";
-import { IModelHierarchyChangeEventArgs, Presentation } from "@itwin/presentation-frontend";
 import {
-  computeVisibleNodes, isTreeModelNode, isTreeModelNodePlaceholder, MutableTreeModel, MutableTreeModelNode, PagedTreeNodeLoader, RenderedItemsRange,
-  TreeModelNode, TreeModelNodeInput, TreeModelSource, TreeNodeItem, usePagedTreeNodeLoader, VisibleTreeNodes,
+  computeVisibleNodes, DelayLoadedTreeNodeItem, isTreeModelNode, isTreeModelNodePlaceholder, MutableTreeModel, MutableTreeModelNode,
+  PagedTreeNodeLoader, RenderedItemsRange, TreeModelNode, TreeModelNodeInput, TreeModelSource, TreeNodeItem, usePagedTreeNodeLoader, VisibleTreeNodes,
 } from "@itwin/components-react";
+import { HierarchyUpdateRecord, PageOptions, UPDATE_FULL } from "@itwin/presentation-common";
+import { IModelHierarchyChangeEventArgs, Presentation } from "@itwin/presentation-frontend";
 import { RulesetRegistrationHelper } from "../../common/RulesetRegistrationHelper";
 import { PresentationTreeDataProvider, PresentationTreeDataProviderProps } from "../DataProvider";
 import { IPresentationTreeDataProvider } from "../IPresentationTreeDataProvider";
@@ -330,7 +330,7 @@ async function updateModelSourceAfterIModelChange(
 /** @internal */
 export interface ReloadedHierarchyPart {
   parentId: string | undefined;
-  nodeItems: TreeNodeItem[];
+  nodeItems: DelayLoadedTreeNodeItem[];
   offset: number;
 }
 
@@ -398,15 +398,15 @@ function rebuildSubTree(oldModel: MutableTreeModel, newModel: MutableTreeModel,
   }
 }
 
-function createModelNodeInput(oldNode: MutableTreeModelNode | undefined, newNode: TreeNodeItem): TreeModelNodeInput {
+function createModelNodeInput(oldNode: MutableTreeModelNode | undefined, newNode: DelayLoadedTreeNodeItem): TreeModelNodeInput {
   const newInput = {
     description: newNode.description,
     isExpanded: !!newNode.autoExpand,
     id: newNode.id,
     item: newNode,
     label: newNode.label,
     isLoading: false,
-    numChildren: 0,
+    numChildren: !newNode.hasChildren ? 0 : undefined,
     isSelected: false,
   };
   if (!oldNode) {

presentation/components/src/test/tree/controlled/TreeHooks.test.ts

@@ -603,6 +603,35 @@ describe("applyHierarchyUpdateRecords", () => {
     expect(updatedTree.getNode("root1")?.description).to.be.eq("updated-description");
   });
 
+  it("updates node children count to 'undefined' when it is collapsed and has children", () => {
+    const initialTree = createTreeModel([
+      {
+        root1: [
+          "child1",
+          "child2",
+        ],
+      },
+    ]);
+
+    const updatedTree = applyHierarchyChanges(
+      initialTree,
+      [{
+        parent: undefined,
+        nodesCount: 1,
+        expandedNodes: [{
+          node: { ...createNode("root1"), description: "updated-description", hasChildren: true },
+          position: 0,
+        }],
+      }],
+      [],
+      {}
+    );
+
+    expectTree(updatedTree, ["root1"]);
+    expect(updatedTree.getNode("root1")?.description).to.be.eq("updated-description");
+    expect(updatedTree.getNode("root1")?.numChildren).to.be.undefined;
+  });
+
   it("updates root node and adds reloaded siblings", () => {
     const initialTree = createTreeModel(["root1", "root2", "root3"]);