From db4ec4ae0fc5b62aa6a200a38cac0c4724b836dd Mon Sep 17 00:00:00 2001
From: Grigas <35135765+grigasp@users.noreply.github.com>
Date: Wed, 18 Dec 2024 10:48:49 +0200
Subject: [PATCH] Tree widget: Add `selectionPredicate` to models tree building
 blocks (#1124)

---
 apps/test-viewer/src/UiProvidersConfig.tsx    | 29 ++++++++-----
 .../src/components/ViewerOptions.tsx          | 38 +++++++++++------
 ...-8d0932d5-74e8-4425-8fc3-727561809066.json |  7 ++++
 .../tree-widget/api/tree-widget-react.api.md  | 11 +++--
 .../trees/common/components/Tree.tsx          | 41 +++++++++++++++++--
 .../trees/models-tree/ModelsTree.tsx          |  2 +
 .../trees/models-tree/ModelsTreeComponent.tsx |  1 +
 .../trees/models-tree/UseModelsTree.tsx       | 20 ++++++++-
 .../models-tree/internal/ModelsTreeNode.ts    | 31 +++++++++++---
 .../src/e2e-tests/ModelsTree.test.ts          | 28 +++++++++++++
 .../tree-widget/src/test/trees/Common.ts      |  4 +-
 11 files changed, 173 insertions(+), 39 deletions(-)
 create mode 100644 change/@itwin-tree-widget-react-8d0932d5-74e8-4425-8fc3-727561809066.json

diff --git a/apps/test-viewer/src/UiProvidersConfig.tsx b/apps/test-viewer/src/UiProvidersConfig.tsx
index 7da365419..0ccf4fc2f 100644
--- a/apps/test-viewer/src/UiProvidersConfig.tsx
+++ b/apps/test-viewer/src/UiProvidersConfig.tsx
@@ -114,16 +114,21 @@ const configuredUiItems = new Map<string, UiItem>([
               {
                 id: ModelsTreeComponent.id,
                 getLabel: () => ModelsTreeComponent.getLabel(),
-                render: (props) => (
-                  <ModelsTreeComponent
-                    getSchemaContext={getSchemaContext}
-                    density={props.density}
-                    selectionStorage={unifiedSelectionStorage}
-                    selectionMode={"extended"}
-                    onPerformanceMeasured={props.onPerformanceMeasured}
-                    onFeatureUsed={props.onFeatureUsed}
-                  />
-                ),
+                render: (props) => {
+                  // eslint-disable-next-line react-hooks/rules-of-hooks
+                  const { disableNodesSelection } = useViewerOptionsContext();
+                  return (
+                    <ModelsTreeComponent
+                      getSchemaContext={getSchemaContext}
+                      density={props.density}
+                      selectionStorage={unifiedSelectionStorage}
+                      selectionMode={"extended"}
+                      selectionPredicate={disableNodesSelection ? disabledSelectionPredicate : undefined}
+                      onPerformanceMeasured={props.onPerformanceMeasured}
+                      onFeatureUsed={props.onFeatureUsed}
+                    />
+                  );
+                },
               },
               {
                 id: CategoriesTreeComponent.id,
@@ -309,3 +314,7 @@ function TreeWidgetWithOptions(props: { trees: SelectableTreeDefinition[] }) {
     />
   );
 }
+
+function disabledSelectionPredicate() {
+  return false;
+}
diff --git a/apps/test-viewer/src/components/ViewerOptions.tsx b/apps/test-viewer/src/components/ViewerOptions.tsx
index fc09012c6..83f593bf4 100644
--- a/apps/test-viewer/src/components/ViewerOptions.tsx
+++ b/apps/test-viewer/src/components/ViewerOptions.tsx
@@ -5,7 +5,7 @@
 
 import { createContext, useContext, useState } from "react";
 import { StatusBarSection } from "@itwin/appui-react";
-import { SvgVisibilityShow } from "@itwin/itwinui-icons-react";
+import { SvgSelection, SvgVisibilityShow } from "@itwin/itwinui-icons-react";
 import { IconButton } from "@itwin/itwinui-react";
 
 import type { PropsWithChildren } from "react";
@@ -13,10 +13,12 @@ import type { UiItemsProvider } from "@itwin/appui-react";
 
 export interface ViewerOptionsContext {
   density: "default" | "enlarged";
+  disableNodesSelection: boolean;
 }
 
 export interface ViewerActionsContext {
   setDensity: React.Dispatch<React.SetStateAction<"default" | "enlarged">>;
+  setDisableNodesSelection: React.Dispatch<React.SetStateAction<boolean>>;
 }
 
 const viewerOptionsContext = createContext<ViewerOptionsContext>({} as ViewerOptionsContext);
@@ -24,10 +26,10 @@ const viewerActionsContext = createContext<ViewerActionsContext>({} as ViewerAct
 
 export function ViewerOptionsProvider(props: PropsWithChildren<unknown>) {
   const [density, setDensity] = useState<"default" | "enlarged">("default");
-
+  const [disableNodesSelection, setDisableNodesSelection] = useState<boolean>(false);
   return (
-    <viewerActionsContext.Provider value={{ setDensity }}>
-      <viewerOptionsContext.Provider value={{ density }}>{props.children}</viewerOptionsContext.Provider>
+    <viewerActionsContext.Provider value={{ setDensity, setDisableNodesSelection }}>
+      <viewerOptionsContext.Provider value={{ density, disableNodesSelection }}>{props.children}</viewerOptionsContext.Provider>
     </viewerActionsContext.Provider>
   );
 }
@@ -44,24 +46,34 @@ export const statusBarActionsProvider: UiItemsProvider = {
   id: "ViewerOptionsUiItemsProvider",
   getStatusBarItems: () => [
     {
-      id: `expandedLayoutButton`,
-      content: <ToggleLayoutButton />,
+      id: `toggleExpandedLayoutButton`,
+      content: <ToggleExpandedLayoutButton />,
       itemPriority: 1,
       section: StatusBarSection.Left,
     },
+    {
+      id: `toggleTreeNodesSelectionButton`,
+      content: <ToggleTreeNodesSelectionButton />,
+      itemPriority: 2,
+      section: StatusBarSection.Left,
+    },
   ],
 };
 
-function ToggleLayoutButton() {
+function ToggleExpandedLayoutButton() {
   const { setDensity } = useViewerActionsContext();
   return (
-    <IconButton
-      // aria-label="Toggle expanded layout"
-      label="Toggle expanded layout"
-      styleType="borderless"
-      onClick={() => setDensity((prev) => (prev === "default" ? "enlarged" : "default"))}
-    >
+    <IconButton label="Toggle expanded layout" styleType="borderless" onClick={() => setDensity((prev) => (prev === "default" ? "enlarged" : "default"))}>
       <SvgVisibilityShow />
     </IconButton>
   );
 }
+
+function ToggleTreeNodesSelectionButton() {
+  const { setDisableNodesSelection } = useViewerActionsContext();
+  return (
+    <IconButton label="Toggle tree nodes' selection" styleType="borderless" onClick={() => setDisableNodesSelection((prev) => !prev)}>
+      <SvgSelection />
+    </IconButton>
+  );
+}
diff --git a/change/@itwin-tree-widget-react-8d0932d5-74e8-4425-8fc3-727561809066.json b/change/@itwin-tree-widget-react-8d0932d5-74e8-4425-8fc3-727561809066.json
new file mode 100644
index 000000000..5680aeb1a
--- /dev/null
+++ b/change/@itwin-tree-widget-react-8d0932d5-74e8-4425-8fc3-727561809066.json
@@ -0,0 +1,7 @@
+{
+  "type": "minor",
+  "comment": "Added an optional `selectionPredicate` function prop to `ModelsTreeComponent`, `ModelsTree`, `useModelsTree` and `Tree` components. When provided, it allows consumers to conditionally enable/disable selection of tree nodes.",
+  "packageName": "@itwin/tree-widget-react",
+  "email": "35135765+grigasp@users.noreply.github.com",
+  "dependentChangeType": "patch"
+}
diff --git a/packages/itwin/tree-widget/api/tree-widget-react.api.md b/packages/itwin/tree-widget/api/tree-widget-react.api.md
index 07ee3a24b..b61f4bdff 100644
--- a/packages/itwin/tree-widget/api/tree-widget-react.api.md
+++ b/packages/itwin/tree-widget/api/tree-widget-react.api.md
@@ -243,7 +243,7 @@ export const ModelsTreeComponent: {
 };
 
 // @public (undocumented)
-interface ModelsTreeComponentProps extends Pick<ModelsTreeProps, "getSchemaContext" | "selectionStorage" | "density" | "hierarchyLevelConfig" | "selectionMode" | "hierarchyConfig" | "visibilityHandlerOverrides" | "getFilteredPaths"> {
+interface ModelsTreeComponentProps extends Pick<ModelsTreeProps, "getSchemaContext" | "selectionStorage" | "density" | "hierarchyLevelConfig" | "selectionMode" | "selectionPredicate" | "hierarchyConfig" | "visibilityHandlerOverrides" | "getFilteredPaths"> {
     headerButtons?: Array<(props: ModelsTreeHeaderButtonProps) => React.ReactNode>;
     // (undocumented)
     onFeatureUsed?: (feature: string) => void;
@@ -403,6 +403,7 @@ type TreeProps = Pick<FunctionProps<typeof useIModelTree>, "getFilteredPaths" |
     getSchemaContext: (imodel: IModelConnection) => SchemaContext;
     treeName: string;
     selectionStorage: SelectionStorage;
+    selectionPredicate?: (node: PresentationHierarchyNode) => boolean;
     treeRenderer: (treeProps: Required<Pick<TreeRendererProps, "rootNodes" | "expandNode" | "onNodeClick" | "onNodeKeyDown" | "onFilterClick" | "isNodeSelected" | "getHierarchyLevelDetails" | "size" | "getLabel">>) => ReactNode;
     imodelAccess?: FunctionProps<typeof useIModelTree>["imodelAccess"];
     hierarchyLevelSizeLimit?: number;
@@ -485,7 +486,7 @@ interface UseCategoriesTreeResult {
 }
 
 // @beta
-export function useModelsTree({ activeView, filter, hierarchyConfig, visibilityHandlerOverrides, getFilteredPaths, onModelsFiltered, }: UseModelsTreeProps): UseModelsTreeResult;
+export function useModelsTree({ activeView, filter, hierarchyConfig, visibilityHandlerOverrides, getFilteredPaths, onModelsFiltered, selectionPredicate: nodeTypeSelectionPredicate, }: UseModelsTreeProps): UseModelsTreeResult;
 
 // @public
 export function useModelsTreeButtonProps({ imodel, viewport }: {
@@ -514,6 +515,10 @@ interface UseModelsTreeProps {
     hierarchyConfig?: Partial<ModelsTreeHierarchyConfiguration>;
     // (undocumented)
     onModelsFiltered?: (modelIds: Id64String[] | undefined) => void;
+    selectionPredicate?: (props: {
+        node: PresentationHierarchyNode;
+        type: "subject" | "model" | "category" | "element" | "elements-class-group";
+    }) => boolean;
     // (undocumented)
     visibilityHandlerOverrides?: ModelsTreeVisibilityHandlerOverrides;
 }
@@ -521,7 +526,7 @@ interface UseModelsTreeProps {
 // @beta (undocumented)
 interface UseModelsTreeResult {
     // (undocumented)
-    modelsTreeProps: Pick<VisibilityTreeProps, "treeName" | "getHierarchyDefinition" | "getFilteredPaths" | "visibilityHandlerFactory" | "highlight" | "noDataMessage">;
+    modelsTreeProps: Pick<VisibilityTreeProps, "treeName" | "getHierarchyDefinition" | "getFilteredPaths" | "visibilityHandlerFactory" | "highlight" | "noDataMessage" | "selectionPredicate">;
     // (undocumented)
     rendererProps: Required<Pick<VisibilityTreeRendererProps, "getIcon" | "onNodeDoubleClick">>;
 }
diff --git a/packages/itwin/tree-widget/src/components/trees/common/components/Tree.tsx b/packages/itwin/tree-widget/src/components/trees/common/components/Tree.tsx
index 30efea80e..c6fb4eac6 100644
--- a/packages/itwin/tree-widget/src/components/trees/common/components/Tree.tsx
+++ b/packages/itwin/tree-widget/src/components/trees/common/components/Tree.tsx
@@ -23,7 +23,7 @@ import type { FunctionProps } from "../Utils";
 import type { ReactNode } from "react";
 import type { IModelConnection } from "@itwin/core-frontend";
 import type { SchemaContext } from "@itwin/ecschema-metadata";
-import type { SelectionStorage, useIModelTree } from "@itwin/presentation-hierarchies-react";
+import type { PresentationHierarchyNode, SelectionStorage, useIModelTree } from "@itwin/presentation-hierarchies-react";
 import type { HighlightInfo } from "../UseNodeHighlighting";
 import type { TreeRendererProps } from "./TreeRenderer";
 
@@ -38,6 +38,11 @@ export type TreeProps = Pick<FunctionProps<typeof useIModelTree>, "getFilteredPa
     treeName: string;
     /** Unified selection storage that should be used by tree to handle tree selection changes. */
     selectionStorage: SelectionStorage;
+    /**
+     * An optional predicate to allow or prohibit selection of a node.
+     * When not supplied, all nodes are selectable.
+     */
+    selectionPredicate?: (node: PresentationHierarchyNode) => boolean;
     /** Tree renderer that should be used to render tree data. */
     treeRenderer: (
       treeProps: Required<
@@ -89,6 +94,7 @@ function TreeImpl({
   getFilteredPaths,
   defaultHierarchyLevelSizeLimit,
   getHierarchyDefinition,
+  selectionPredicate,
   selectionMode,
   onReload,
   treeRenderer,
@@ -100,8 +106,9 @@ function TreeImpl({
   const [imodelChanged] = useState(new BeEvent<() => void>());
   const {
     rootNodes,
+    getNode,
     isLoading,
-    selectNodes,
+    selectNodes: selectNodesAction,
     setFormatter: _setFormatter,
     expandNode,
     ...treeProps
@@ -123,8 +130,12 @@ function TreeImpl({
   });
   useIModelChangeListener({ imodel, action: useCallback(() => imodelChanged.raiseEvent(), [imodelChanged]) });
 
-  const reportingSelectNodes = useReportingAction({ action: selectNodes });
-  const { onNodeClick, onNodeKeyDown } = useSelectionHandler({ rootNodes, selectNodes: reportingSelectNodes, selectionMode: selectionMode ?? "single" });
+  const selectNodes = useSelectionPredicate({
+    action: useReportingAction({ action: selectNodesAction }),
+    predicate: selectionPredicate,
+    getNode,
+  });
+  const { onNodeClick, onNodeKeyDown } = useSelectionHandler({ rootNodes, selectNodes, selectionMode: selectionMode ?? "single" });
   const { filteringDialog, onFilterClick } = useHierarchyLevelFiltering({
     imodel,
     defaultHierarchyLevelSizeLimit,
@@ -174,3 +185,25 @@ function TreeImpl({
     </div>
   );
 }
+
+function useSelectionPredicate({
+  action,
+  predicate,
+  getNode,
+}: {
+  action: (...args: any[]) => void;
+  predicate?: (node: PresentationHierarchyNode) => boolean;
+  getNode: (nodeId: string) => PresentationHierarchyNode | undefined;
+}): ReturnType<typeof useIModelUnifiedSelectionTree>["selectNodes"] {
+  return useCallback(
+    (nodeIds, changeType) =>
+      action(
+        nodeIds.filter((nodeId) => {
+          const node = getNode(nodeId);
+          return node && (!predicate || predicate(node));
+        }),
+        changeType,
+      ),
+    [action, getNode, predicate],
+  );
+}
diff --git a/packages/itwin/tree-widget/src/components/trees/models-tree/ModelsTree.tsx b/packages/itwin/tree-widget/src/components/trees/models-tree/ModelsTree.tsx
index f262443ba..668e2f1b2 100644
--- a/packages/itwin/tree-widget/src/components/trees/models-tree/ModelsTree.tsx
+++ b/packages/itwin/tree-widget/src/components/trees/models-tree/ModelsTree.tsx
@@ -29,6 +29,7 @@ export function ModelsTree({
   hierarchyLevelConfig,
   hierarchyConfig,
   selectionMode,
+  selectionPredicate,
   visibilityHandlerOverrides,
   getFilteredPaths,
   onModelsFiltered,
@@ -40,6 +41,7 @@ export function ModelsTree({
     visibilityHandlerOverrides,
     getFilteredPaths,
     onModelsFiltered,
+    selectionPredicate,
   });
 
   return (
diff --git a/packages/itwin/tree-widget/src/components/trees/models-tree/ModelsTreeComponent.tsx b/packages/itwin/tree-widget/src/components/trees/models-tree/ModelsTreeComponent.tsx
index 67755730d..30b35d99b 100644
--- a/packages/itwin/tree-widget/src/components/trees/models-tree/ModelsTreeComponent.tsx
+++ b/packages/itwin/tree-widget/src/components/trees/models-tree/ModelsTreeComponent.tsx
@@ -37,6 +37,7 @@ interface ModelsTreeComponentProps
     | "density"
     | "hierarchyLevelConfig"
     | "selectionMode"
+    | "selectionPredicate"
     | "hierarchyConfig"
     | "visibilityHandlerOverrides"
     | "getFilteredPaths"
diff --git a/packages/itwin/tree-widget/src/components/trees/models-tree/UseModelsTree.tsx b/packages/itwin/tree-widget/src/components/trees/models-tree/UseModelsTree.tsx
index 02ee4e9a7..226194e7d 100644
--- a/packages/itwin/tree-widget/src/components/trees/models-tree/UseModelsTree.tsx
+++ b/packages/itwin/tree-widget/src/components/trees/models-tree/UseModelsTree.tsx
@@ -15,6 +15,7 @@ import { FilterLimitExceededError } from "../common/TreeErrors";
 import { useIModelChangeListener } from "../common/UseIModelChangeListener";
 import { useTelemetryContext } from "../common/UseTelemetryContext";
 import { ModelsTreeIdsCache } from "./internal/ModelsTreeIdsCache";
+import { ModelsTreeNode } from "./internal/ModelsTreeNode";
 import { createModelsTreeVisibilityHandler } from "./internal/ModelsTreeVisibilityHandler";
 import { defaultHierarchyConfiguration, ModelsTreeDefinition } from "./ModelsTreeDefinition";
 
@@ -41,13 +42,18 @@ export interface UseModelsTreeProps {
     createInstanceKeyPaths: (props: { targetItems: Array<InstanceKey | ElementsGroupInfo> } | { label: string }) => Promise<HierarchyFilteringPath[]>;
   }) => Promise<HierarchyFilteringPath[]>;
   onModelsFiltered?: (modelIds: Id64String[] | undefined) => void;
+  /**
+   * An optional predicate to allow or prohibit selection of a node.
+   * When not supplied, all nodes are selectable.
+   */
+  selectionPredicate?: (props: { node: PresentationHierarchyNode; type: "subject" | "model" | "category" | "element" | "elements-class-group" }) => boolean;
 }
 
 /** @beta */
 interface UseModelsTreeResult {
   modelsTreeProps: Pick<
     VisibilityTreeProps,
-    "treeName" | "getHierarchyDefinition" | "getFilteredPaths" | "visibilityHandlerFactory" | "highlight" | "noDataMessage"
+    "treeName" | "getHierarchyDefinition" | "getFilteredPaths" | "visibilityHandlerFactory" | "highlight" | "noDataMessage" | "selectionPredicate"
   >;
   rendererProps: Required<Pick<VisibilityTreeRendererProps, "getIcon" | "onNodeDoubleClick">>;
 }
@@ -63,6 +69,7 @@ export function useModelsTree({
   visibilityHandlerOverrides,
   getFilteredPaths,
   onModelsFiltered,
+  selectionPredicate: nodeTypeSelectionPredicate,
 }: UseModelsTreeProps): UseModelsTreeResult {
   const [filteringError, setFilteringError] = useState<ModelsTreeFilteringError | undefined>(undefined);
   const hierarchyConfiguration = useMemo<ModelsTreeHierarchyConfiguration>(
@@ -195,6 +202,16 @@ export function useModelsTree({
     return undefined;
   }, [filter, loadFocusedItems, getModelsTreeIdsCache, onFeatureUsed, getFilteredPaths, hierarchyConfiguration, onModelsFiltered, onFilteredPathsChanged]);
 
+  const nodeSelectionPredicate = useCallback<NonNullable<VisibilityTreeProps["selectionPredicate"]>>(
+    (node) => {
+      if (!nodeTypeSelectionPredicate) {
+        return true;
+      }
+      return nodeTypeSelectionPredicate({ node, type: ModelsTreeNode.getType(node.nodeData) });
+    },
+    [nodeTypeSelectionPredicate],
+  );
+
   return {
     modelsTreeProps: {
       treeName: "models-tree-v2",
@@ -203,6 +220,7 @@ export function useModelsTree({
       getFilteredPaths: getPaths,
       noDataMessage: getNoDataMessage(filter, filteringError),
       highlight: filter ? { text: filter } : undefined,
+      selectionPredicate: nodeSelectionPredicate,
     },
     rendererProps: {
       onNodeDoubleClick,
diff --git a/packages/itwin/tree-widget/src/components/trees/models-tree/internal/ModelsTreeNode.ts b/packages/itwin/tree-widget/src/components/trees/models-tree/internal/ModelsTreeNode.ts
index c09b3df9b..0fa8ba90a 100644
--- a/packages/itwin/tree-widget/src/components/trees/models-tree/internal/ModelsTreeNode.ts
+++ b/packages/itwin/tree-widget/src/components/trees/models-tree/internal/ModelsTreeNode.ts
@@ -4,9 +4,11 @@
  *--------------------------------------------------------------------------------------------*/
 
 import type { Id64String } from "@itwin/core-bentley";
+import { HierarchyNodeKey } from "@itwin/presentation-hierarchies";
 
 interface ModelsTreeNode {
-  extendedData?: any;
+  key: HierarchyNodeKey;
+  extendedData?: { [id: string]: any };
 }
 
 /**
@@ -16,22 +18,39 @@ export namespace ModelsTreeNode {
   /**
    * Determines if a node represents a subject.
    */
-  export const isSubjectNode = (node: ModelsTreeNode) => !!node.extendedData?.isSubject;
+  export const isSubjectNode = (node: Pick<ModelsTreeNode, "extendedData">) => !!node.extendedData?.isSubject;
 
   /**
    * Determines if a node represents a model.
    */
-  export const isModelNode = (node: ModelsTreeNode) => !!node.extendedData?.isModel;
+  export const isModelNode = (node: Pick<ModelsTreeNode, "extendedData">) => !!node.extendedData?.isModel;
 
   /**
    * Determines if a node represents a category.
    */
-  export const isCategoryNode = (node: ModelsTreeNode) => !!node.extendedData?.isCategory;
+  export const isCategoryNode = (node: Pick<ModelsTreeNode, "extendedData">) => !!node.extendedData?.isCategory;
+
+  /** Returns type of the node. */
+  export const getType = (node: ModelsTreeNode): "subject" | "model" | "category" | "element" | "elements-class-group" => {
+    if (HierarchyNodeKey.isClassGrouping(node.key)) {
+      return "elements-class-group";
+    }
+    if (isSubjectNode(node)) {
+      return "subject";
+    }
+    if (isModelNode(node)) {
+      return "model";
+    }
+    if (isCategoryNode(node)) {
+      return "category";
+    }
+    return "element";
+  };
 
   /**
    * Retrieves model ID from node's extended data.
    */
-  export const getModelId = (node: ModelsTreeNode): Id64String | undefined => {
+  export const getModelId = (node: Pick<ModelsTreeNode, "extendedData">): Id64String | undefined => {
     if (node.extendedData?.modelId) {
       return node.extendedData?.modelId;
     }
@@ -43,5 +62,5 @@ export namespace ModelsTreeNode {
   /**
    * Retrieves category ID from node's extended data.
    */
-  export const getCategoryId = (node: ModelsTreeNode): Id64String | undefined => node.extendedData?.categoryId;
+  export const getCategoryId = (node: Pick<ModelsTreeNode, "extendedData">): Id64String | undefined => node.extendedData?.categoryId;
 }
diff --git a/packages/itwin/tree-widget/src/e2e-tests/ModelsTree.test.ts b/packages/itwin/tree-widget/src/e2e-tests/ModelsTree.test.ts
index 347ce4c30..b2f33b205 100644
--- a/packages/itwin/tree-widget/src/e2e-tests/ModelsTree.test.ts
+++ b/packages/itwin/tree-widget/src/e2e-tests/ModelsTree.test.ts
@@ -28,6 +28,34 @@ test.describe("Models tree", () => {
     await locateNode(treeWidget, "BayTown").getByRole("checkbox", { name: "Visible: All models are visible", exact: true }).waitFor();
   });
 
+  test("disabled selection", async ({ page }) => {
+    // disable nodes' selection
+    await page.getByRole("button", { name: "Toggle tree nodes' selection" }).click();
+
+    const subjectNode = locateNode(treeWidget, "BayTown");
+    await subjectNode.click();
+    await expect(subjectNode).toHaveAttribute("aria-selected", "false");
+
+    const physicalModelNode = locateNode(treeWidget, "ProcessPhysicalModel");
+    await physicalModelNode.click();
+    await expect(physicalModelNode).toHaveAttribute("aria-selected", "false");
+
+    await physicalModelNode.getByLabel("Expand").click();
+    const equipmentNode = locateNode(treeWidget, "Equipment");
+    await equipmentNode.click();
+    await expect(equipmentNode).toHaveAttribute("aria-selected", "false");
+
+    await equipmentNode.getByLabel("Expand").click();
+    const parReboilerGroupingNode = locateNode(treeWidget, "Par. Reboiler");
+    await parReboilerGroupingNode.click();
+    await expect(parReboilerGroupingNode).toHaveAttribute("aria-selected", "false");
+
+    await parReboilerGroupingNode.getByLabel("Expand").click();
+    const parReboilerInstanceNode = locateNode(treeWidget, "EX-302 [4-106]");
+    await parReboilerInstanceNode.click();
+    await expect(parReboilerInstanceNode).toHaveAttribute("aria-selected", "false");
+  });
+
   withDifferentDensities((density) => {
     test("initial tree", async ({ page }) => {
       // wait for element to be visible in the tree
diff --git a/packages/itwin/tree-widget/src/test/trees/Common.ts b/packages/itwin/tree-widget/src/test/trees/Common.ts
index be1f206df..e9359277f 100644
--- a/packages/itwin/tree-widget/src/test/trees/Common.ts
+++ b/packages/itwin/tree-widget/src/test/trees/Common.ts
@@ -28,7 +28,7 @@ export function createIModelMock(queryHandler?: (query: string, params?: QueryBi
 
 export function createFakeSinonViewport(
   props?: Partial<Omit<Viewport, "view" | "perModelCategoryVisibility">> & {
-    view?: Partial<ViewState>;
+    view?: Partial<Omit<ViewState, "isSpatialView">> & { isSpatialView?: () => boolean; };
     perModelCategoryVisibility?: Partial<PerModelCategoryVisibility.Overrides>;
     queryHandler?: Parameters<typeof createIModelMock>[0];
   },
@@ -44,7 +44,7 @@ export function createFakeSinonViewport(
     ...props?.perModelCategoryVisibility,
   };
 
-  const view: Partial<ViewState> = {
+  const view: NonNullable<typeof props>["view"] = {
     isSpatialView: sinon.fake.returns(true),
     viewsCategory: sinon.fake.returns(true),
     viewsModel: sinon.fake.returns(true),