Use new Fluent-based Resource Tree for all environments (#1841)

Co-authored-by: Laurent Nguyen <laurent.nguyen@microsoft.com>

.eslintignore

@@ -125,7 +125,7 @@ src/Explorer/Controls/InputTypeahead/InputTypeaheadComponent.tsx
 src/Explorer/Controls/Notebook/NotebookTerminalComponent.test.tsx
 src/Explorer/Controls/Notebook/NotebookTerminalComponent.tsx
 src/Explorer/Controls/NotebookViewer/NotebookViewerComponent.tsx
-src/Explorer/Controls/TreeComponent/TreeComponent.tsx
+src/Explorer/Controls/TreeComponent/LegacyTreeComponent.tsx
 src/Explorer/Graph/GraphExplorerComponent/GraphExplorer.test.tsx
 src/Explorer/Graph/GraphExplorerComponent/NodePropertiesComponent.tsx
 src/Explorer/Graph/GraphExplorerComponent/ReadOnlyNodePropertiesComponent.test.tsx

src/Common/ResourceTreeContainer.tsx

@@ -1,14 +1,10 @@
-import { ResourceTree } from "Explorer/Tree/ResourceTree";
 import React, { FunctionComponent, MutableRefObject, useEffect, useRef } from "react";
 import arrowLeftImg from "../../images/imgarrowlefticon.svg";
 import refreshImg from "../../images/refresh-cosmos.svg";
-import { AuthType } from "../AuthType";
 import Explorer from "../Explorer/Explorer";
-import { ResourceTokenTree } from "../Explorer/Tree/ResourceTokenTree";
+import { ResourceTree } from "../Explorer/Tree/ResourceTree";
-import { ResourceTree2 } from "../Explorer/Tree2/ResourceTree";
 import { userContext } from "../UserContext";
 import { getApiShortDisplayName } from "../Utils/APITypeUtils";
-import { Platform, configContext } from "./../ConfigContext";
 import { NormalizedEventKey } from "./Constants";
 
 export interface ResourceTreeContainerProps {
@@ -74,12 +70,8 @@ export const ResourceTreeContainer: FunctionComponent<ResourceTreeContainerProps
             </div>
           </div>
         </div>
-        {userContext.authType === AuthType.ResourceToken ? (
+        {userContext.features.enableKoResourceTree ? (
-          <ResourceTokenTree />
-        ) : userContext.features.enableKoResourceTree ? (
           <div style={{ overflowY: "auto" }} data-bind="react:resourceTree" />
-        ) : configContext.platform === Platform.Fabric ? (
-          <ResourceTree2 container={container} />
         ) : (
           <ResourceTree container={container} />
         )}

src/Contracts/ViewModels.ts

@@ -176,6 +176,11 @@ export interface Collection extends CollectionBase {
   loadTriggers(): Promise<any>;
   loadOffer(): Promise<void>;
 
+  showStoredProcedures: ko.Observable<boolean>;
+  showTriggers: ko.Observable<boolean>;
+  showUserDefinedFunctions: ko.Observable<boolean>;
+  showConflicts: ko.Observable<boolean>;
+
   createStoredProcedureNode(data: StoredProcedureDefinition & Resource): StoredProcedure;
   createUserDefinedFunctionNode(data: UserDefinedFunctionDefinition & Resource): UserDefinedFunction;
   createTriggerNode(data: TriggerDefinition | SqlTriggerResource): Trigger;

src/Explorer/ContextMenuButtonFactory.tsx

@@ -1,3 +1,4 @@
+import { TreeNodeMenuItem } from "Explorer/Controls/TreeComponent/TreeNodeComponent";
 import { useDatabases } from "Explorer/useDatabases";
 import { Action } from "Shared/Telemetry/TelemetryConstants";
 import { traceOpen } from "Shared/Telemetry/TelemetryProcessor";
@@ -19,7 +20,6 @@ import { userContext } from "../UserContext";
 import { getCollectionName, getDatabaseName } from "../Utils/APITypeUtils";
 import { useSidePanel } from "../hooks/useSidePanel";
 import { Platform, configContext } from "./../ConfigContext";
-import { TreeNodeMenuItem } from "./Controls/TreeComponent/TreeComponent";
 import Explorer from "./Explorer";
 import { useNotebook } from "./Notebook/useNotebook";
 import { DeleteCollectionConfirmationPane } from "./Panes/DeleteCollectionConfirmationPane/DeleteCollectionConfirmationPane";

src/Explorer/Controls/TreeComponent/LegacyTreeComponent.test.tsx

@@ -1,48 +1,48 @@
-import React from "react";
 import { shallow } from "enzyme";
-import { TreeComponent, TreeNode, TreeNodeComponent } from "./TreeComponent";
+import React from "react";
+import { LegacyTreeComponent, LegacyTreeNode, LegacyTreeNodeComponent } from "./LegacyTreeComponent";
 
-const buildChildren = (): TreeNode[] => {
+const buildChildren = (): LegacyTreeNode[] => {
-  const grandChild11: TreeNode = {
+  const grandChild11: LegacyTreeNode = {
     label: "ZgrandChild11",
   };
-  const grandChild12: TreeNode = {
+  const grandChild12: LegacyTreeNode = {
     label: "AgrandChild12",
   };
-  const child1: TreeNode = {
+  const child1: LegacyTreeNode = {
     label: "Bchild1",
     children: [grandChild11, grandChild12],
   };
 
-  const child2: TreeNode = {
+  const child2: LegacyTreeNode = {
     label: "2child2",
   };
 
   return [child1, child2];
 };
 
-const buildChildren2 = (): TreeNode[] => {
+const buildChildren2 = (): LegacyTreeNode[] => {
-  const grandChild11: TreeNode = {
+  const grandChild11: LegacyTreeNode = {
     label: "ZgrandChild11",
   };
-  const grandChild12: TreeNode = {
+  const grandChild12: LegacyTreeNode = {
     label: "AgrandChild12",
   };
 
-  const child1: TreeNode = {
+  const child1: LegacyTreeNode = {
     label: "aChild",
   };
 
-  const child2: TreeNode = {
+  const child2: LegacyTreeNode = {
     label: "bchild",
     children: [grandChild11, grandChild12],
   };
 
-  const child3: TreeNode = {
+  const child3: LegacyTreeNode = {
     label: "cchild",
   };
 
-  const child4: TreeNode = {
+  const child4: LegacyTreeNode = {
     label: "dchild",
     children: [grandChild11, grandChild12],
   };
@@ -50,7 +50,7 @@ const buildChildren2 = (): TreeNode[] => {
   return [child1, child2, child3, child4];
 };
 
-describe("TreeComponent", () => {
+describe("LegacyTreeComponent", () => {
   it("renders a simple tree", () => {
     const root = {
       label: "root",
@@ -62,14 +62,14 @@ describe("TreeComponent", () => {
       className: "tree",
     };
 
-    const wrapper = shallow(<TreeComponent {...props} />);
+    const wrapper = shallow(<LegacyTreeComponent {...props} />);
     expect(wrapper).toMatchSnapshot();
   });
 });
 
-describe("TreeNodeComponent", () => {
+describe("LegacyTreeNodeComponent", () => {
   it("renders a simple node (sorted children, expanded)", () => {
-    const node: TreeNode = {
+    const node: LegacyTreeNode = {
       label: "label",
       id: "id",
       children: buildChildren(),
@@ -98,12 +98,12 @@ describe("TreeNodeComponent", () => {
       generation: 12,
       paddingLeft: 23,
     };
-    const wrapper = shallow(<TreeNodeComponent {...props} />);
+    const wrapper = shallow(<LegacyTreeNodeComponent {...props} />);
     expect(wrapper).toMatchSnapshot();
   });
 
   it("renders unsorted children by default", () => {
-    const node: TreeNode = {
+    const node: LegacyTreeNode = {
       label: "label",
       children: buildChildren(),
       isExpanded: true,
@@ -113,12 +113,12 @@ describe("TreeNodeComponent", () => {
       generation: 2,
       paddingLeft: 9,
     };
-    const wrapper = shallow(<TreeNodeComponent {...props} />);
+    const wrapper = shallow(<LegacyTreeNodeComponent {...props} />);
     expect(wrapper).toMatchSnapshot();
   });
 
   it("does not render children by default", () => {
-    const node: TreeNode = {
+    const node: LegacyTreeNode = {
       label: "label",
       children: buildChildren(),
       isAlphaSorted: false,
@@ -128,12 +128,12 @@ describe("TreeNodeComponent", () => {
       generation: 2,
       paddingLeft: 9,
     };
-    const wrapper = shallow(<TreeNodeComponent {...props} />);
+    const wrapper = shallow(<LegacyTreeNodeComponent {...props} />);
     expect(wrapper).toMatchSnapshot();
   });
 
   it("renders sorted children, expanded, leaves and parents separated", () => {
-    const node: TreeNode = {
+    const node: LegacyTreeNode = {
       label: "label",
       id: "id",
       children: buildChildren2(),
@@ -156,12 +156,12 @@ describe("TreeNodeComponent", () => {
       generation: 12,
       paddingLeft: 23,
     };
-    const wrapper = shallow(<TreeNodeComponent {...props} />);
+    const wrapper = shallow(<LegacyTreeNodeComponent {...props} />);
     expect(wrapper).toMatchSnapshot();
   });
 
   it("renders loading icon", () => {
-    const node: TreeNode = {
+    const node: LegacyTreeNode = {
       label: "label",
       children: [],
       isExpanded: true,
@@ -172,7 +172,7 @@ describe("TreeNodeComponent", () => {
       generation: 2,
       paddingLeft: 9,
     };
-    const wrapper = shallow(<TreeNodeComponent {...props} />);
+    const wrapper = shallow(<LegacyTreeNodeComponent {...props} />);
     expect(wrapper).toMatchSnapshot();
   });
 });

src/Explorer/Controls/TreeComponent/LegacyTreeComponent.tsx

@@ -12,6 +12,7 @@ import {
   IContextualMenuItemProps,
   IContextualMenuProps,
 } from "@fluentui/react";
+import { TreeNodeMenuItem } from "Explorer/Controls/TreeComponent/TreeNodeComponent";
 import * as React from "react";
 import AnimateHeight from "react-animate-height";
 import LoadingIndicator_3Squares from "../../../../images/LoadingIndicator_3Squares.gif";
@@ -22,18 +23,10 @@ import { StyleConstants } from "../../../Common/StyleConstants";
 import { Action, ActionModifiers } from "../../../Shared/Telemetry/TelemetryConstants";
 import * as TelemetryProcessor from "../../../Shared/Telemetry/TelemetryProcessor";
 
-export interface TreeNodeMenuItem {
+export interface LegacyTreeNode {
-  label: string;
-  onClick: () => void;
-  iconSrc?: string;
-  isDisabled?: boolean;
-  styleClass?: string;
-}
-
-export interface TreeNode {
   label: string;
   id?: string;
-  children?: TreeNode[];
+  children?: LegacyTreeNode[];
   contextMenu?: TreeNodeMenuItem[];
   iconSrc?: string;
   isExpanded?: boolean;
@@ -50,34 +43,37 @@ export interface TreeNode {
   onContextMenuOpen?: () => void;
 }
 
-export interface TreeComponentProps {
+export interface LegacyTreeComponentProps {
-  rootNode: TreeNode;
+  rootNode: LegacyTreeNode;
   style?: any;
   className?: string;
 }
 
-export class TreeComponent extends React.Component<TreeComponentProps> {
+export class LegacyTreeComponent extends React.Component<LegacyTreeComponentProps> {
   public render(): JSX.Element {
     return (
       <div style={this.props.style} className={`treeComponent ${this.props.className}`} role="tree">
-        <TreeNodeComponent paddingLeft={0} node={this.props.rootNode} generation={0} />
+        <LegacyTreeNodeComponent paddingLeft={0} node={this.props.rootNode} generation={0} />
       </div>
     );
   }
 }
 
 /* Tree node is a react component */
-interface TreeNodeComponentProps {
+interface LegacyTreeNodeComponentProps {
-  node: TreeNode;
+  node: LegacyTreeNode;
   generation: number;
   paddingLeft: number;
 }
 
-interface TreeNodeComponentState {
+interface LegacyTreeNodeComponentState {
   isExpanded: boolean;
   isMenuShowing: boolean;
 }
-export class TreeNodeComponent extends React.Component<TreeNodeComponentProps, TreeNodeComponentState> {
+export class LegacyTreeNodeComponent extends React.Component<
+  LegacyTreeNodeComponentProps,
+  LegacyTreeNodeComponentState
+> {
   private static readonly paddingPerGenerationPx = 16;
   private static readonly iconOffset = 22;
   private static readonly transitionDurationMS = 200;
@@ -85,7 +81,7 @@ export class TreeNodeComponent extends React.Component<TreeNodeComponentProps, T
   private contextMenuRef = React.createRef<HTMLDivElement>();
   private isExpanded: boolean;
 
-  constructor(props: TreeNodeComponentProps) {
+  constructor(props: LegacyTreeNodeComponentProps) {
     super(props);
     this.isExpanded = props.node.isExpanded;
     this.state = {
@@ -94,13 +90,13 @@ export class TreeNodeComponent extends React.Component<TreeNodeComponentProps, T
     };
   }
 
-  componentDidUpdate(prevProps: TreeNodeComponentProps, prevState: TreeNodeComponentState) {
+  componentDidUpdate(prevProps: LegacyTreeNodeComponentProps, prevState: LegacyTreeNodeComponentState) {
     // Only call when expand has actually changed
     if (this.state.isExpanded !== prevState.isExpanded) {
       if (this.state.isExpanded) {
-        this.props.node.onExpanded && setTimeout(this.props.node.onExpanded, TreeNodeComponent.callbackDelayMS);
+        this.props.node.onExpanded && setTimeout(this.props.node.onExpanded, LegacyTreeNodeComponent.callbackDelayMS);
       } else {
-        this.props.node.onCollapsed && setTimeout(this.props.node.onCollapsed, TreeNodeComponent.callbackDelayMS);
+        this.props.node.onCollapsed && setTimeout(this.props.node.onCollapsed, LegacyTreeNodeComponent.callbackDelayMS);
       }
     }
     if (this.props.node.isExpanded !== this.isExpanded) {
@@ -115,18 +111,18 @@ export class TreeNodeComponent extends React.Component<TreeNodeComponentProps, T
     return this.renderNode(this.props.node, this.props.generation);
   }
 
-  private static getSortedChildren(treeNode: TreeNode): TreeNode[] {
+  private static getSortedChildren(treeNode: LegacyTreeNode): LegacyTreeNode[] {
     if (!treeNode || !treeNode.children) {
       return undefined;
     }
 
-    const compareFct = (a: TreeNode, b: TreeNode) => a.label.localeCompare(b.label);
+    const compareFct = (a: LegacyTreeNode, b: LegacyTreeNode) => a.label.localeCompare(b.label);
 
     let unsortedChildren;
     if (treeNode.isLeavesParentsSeparate) {
       // Separate parents and leave
-      const parents: TreeNode[] = treeNode.children.filter((node) => node.children);
+      const parents: LegacyTreeNode[] = treeNode.children.filter((node) => node.children);
-      const leaves: TreeNode[] = treeNode.children.filter((node) => !node.children);
+      const leaves: LegacyTreeNode[] = treeNode.children.filter((node) => !node.children);
 
       if (treeNode.isAlphaSorted) {
         parents.sort(compareFct);
@@ -141,18 +137,18 @@ export class TreeNodeComponent extends React.Component<TreeNodeComponentProps, T
     return unsortedChildren;
   }
 
-  private static isNodeHeaderBlank(node: TreeNode): boolean {
+  private static isNodeHeaderBlank(node: LegacyTreeNode): boolean {
     return (node.label === undefined || node.label === null) && !node.contextMenu;
   }
 
-  private renderNode(node: TreeNode, generation: number): JSX.Element {
+  private renderNode(node: LegacyTreeNode, generation: number): JSX.Element {
-    let paddingLeft = generation * TreeNodeComponent.paddingPerGenerationPx;
+    const paddingLeft = generation * LegacyTreeNodeComponent.paddingPerGenerationPx;
     let additionalOffsetPx = 15;
 
     if (node.children) {
-      const childrenWithSubChildren = node.children.filter((child: TreeNode) => !!child.children);
+      const childrenWithSubChildren = node.children.filter((child: LegacyTreeNode) => !!child.children);
       if (childrenWithSubChildren.length > 0) {
-        additionalOffsetPx = TreeNodeComponent.iconOffset;
+        additionalOffsetPx = LegacyTreeNodeComponent.iconOffset;
       }
     }
 
@@ -160,10 +156,10 @@ export class TreeNodeComponent extends React.Component<TreeNodeComponentProps, T
     const showSelected =
       this.props.node.isSelected &&
       this.props.node.isSelected() &&
-      !TreeNodeComponent.isAnyDescendantSelected(this.props.node);
+      !LegacyTreeNodeComponent.isAnyDescendantSelected(this.props.node);
 
     const headerStyle: React.CSSProperties = { paddingLeft: this.props.paddingLeft };
-    if (TreeNodeComponent.isNodeHeaderBlank(node)) {
+    if (LegacyTreeNodeComponent.isNodeHeaderBlank(node)) {
       headerStyle.height = 0;
       headerStyle.padding = 0;
     }
@@ -195,10 +191,13 @@ export class TreeNodeComponent extends React.Component<TreeNodeComponentProps, T
           <img className="loadingIcon" src={LoadingIndicator_3Squares} hidden={!this.props.node.isLoading} />
         </div>
         {node.children && (
-          <AnimateHeight duration={TreeNodeComponent.transitionDurationMS} height={this.state.isExpanded ? "auto" : 0}>
+          <AnimateHeight
+            duration={LegacyTreeNodeComponent.transitionDurationMS}
+            height={this.state.isExpanded ? "auto" : 0}
+          >
             <div className="nodeChildren" data-test={node.label} role="group">
-              {TreeNodeComponent.getSortedChildren(node).map((childNode: TreeNode) => (
+              {LegacyTreeNodeComponent.getSortedChildren(node).map((childNode: LegacyTreeNode) => (
-                <TreeNodeComponent
+                <LegacyTreeNodeComponent
                   key={`${childNode.label}-${generation + 1}-${childNode.timestamp}`}
                   node={childNode}
                   generation={generation + 1}
@@ -216,12 +215,14 @@ export class TreeNodeComponent extends React.Component<TreeNodeComponentProps, T
    * Recursive: is the node or any descendant selected
    * @param node
    */
-  private static isAnyDescendantSelected(node: TreeNode): boolean {
+  private static isAnyDescendantSelected(node: LegacyTreeNode): boolean {
     return (
       node.children &&
       node.children.reduce(
-        (previous: boolean, child: TreeNode) =>
+        (previous: boolean, child: LegacyTreeNode) =>
-          previous || (child.isSelected && child.isSelected()) || TreeNodeComponent.isAnyDescendantSelected(child),
+          previous ||
+          (child.isSelected && child.isSelected()) ||
+          LegacyTreeNodeComponent.isAnyDescendantSelected(child),
         false,
       )
     );
@@ -232,10 +233,10 @@ export class TreeNodeComponent extends React.Component<TreeNodeComponentProps, T
   }
 
   private onRightClick = (): void => {
-    this.contextMenuRef.current.firstChild.dispatchEvent(TreeNodeComponent.createClickEvent());
+    this.contextMenuRef.current.firstChild.dispatchEvent(LegacyTreeNodeComponent.createClickEvent());
   };
 
-  private renderContextMenuButton(node: TreeNode): JSX.Element {
+  private renderContextMenuButton(node: LegacyTreeNode): JSX.Element {
     const menuItemLabel = "More";
     const buttonStyles: Partial<IButtonStyles> = {
       rootFocused: { outline: `1px dashed ${StyleConstants.FocusColor}` },
@@ -265,7 +266,7 @@ export class TreeNodeComponent extends React.Component<TreeNodeComponentProps, T
               <div
                 data-test={`treeComponentMenuItemContainer`}
                 className="treeComponentMenuItemContainer"
-                onContextMenu={(e) => e.target.dispatchEvent(TreeNodeComponent.createClickEvent())}
+                onContextMenu={(e) => e.target.dispatchEvent(LegacyTreeNodeComponent.createClickEvent())}
               >
                 {props.item.onRenderIcon()}
                 <span
@@ -297,7 +298,7 @@ export class TreeNodeComponent extends React.Component<TreeNodeComponentProps, T
     );
   }
 
-  private renderCollapseExpandIcon(node: TreeNode): JSX.Element {
+  private renderCollapseExpandIcon(node: LegacyTreeNode): JSX.Element {
     if (!node.children || !node.label) {
       return <></>;
     }
@@ -314,12 +315,12 @@ export class TreeNodeComponent extends React.Component<TreeNodeComponentProps, T
     );
   }
 
-  private onNodeClick = (event: React.MouseEvent<HTMLDivElement>, node: TreeNode): void => {
+  private onNodeClick = (event: React.MouseEvent<HTMLDivElement>, node: LegacyTreeNode): void => {
     event.stopPropagation();
     if (node.children) {
       const isExpanded = !this.state.isExpanded;
       // Prevent collapsing if node header is blank
-      if (!(TreeNodeComponent.isNodeHeaderBlank(node) && !isExpanded)) {
+      if (!(LegacyTreeNodeComponent.isNodeHeaderBlank(node) && !isExpanded)) {
         this.setState({ isExpanded });
       }
     }
@@ -327,14 +328,14 @@ export class TreeNodeComponent extends React.Component<TreeNodeComponentProps, T
     this.props.node.onClick && this.props.node.onClick(this.state.isExpanded);
   };
 
-  private onNodeKeyPress = (event: React.KeyboardEvent<HTMLDivElement>, node: TreeNode): void => {
+  private onNodeKeyPress = (event: React.KeyboardEvent<HTMLDivElement>, node: LegacyTreeNode): void => {
     if (event.charCode === Constants.KeyCodes.Space || event.charCode === Constants.KeyCodes.Enter) {
       event.stopPropagation();
       this.props.node.onClick && this.props.node.onClick(this.state.isExpanded);
     }
   };
 
-  private onCollapseExpandIconKeyPress = (event: React.KeyboardEvent<HTMLDivElement>, node: TreeNode): void => {
+  private onCollapseExpandIconKeyPress = (event: React.KeyboardEvent<HTMLDivElement>, node: LegacyTreeNode): void => {
     if (event.charCode === Constants.KeyCodes.Space || event.charCode === Constants.KeyCodes.Enter) {
       event.stopPropagation();
       if (node.children) {

src/Explorer/Controls/TreeComponent/TreeNodeComponent.test.tsx

@@ -0,0 +1,180 @@
+import { TreeItem, TreeItemLayout, tokens } from "@fluentui/react-components";
+import PromiseSource from "Utils/PromiseSource";
+import { mount, shallow } from "enzyme";
+import React from "react";
+import { act } from "react-dom/test-utils";
+import { TreeNode, TreeNodeComponent } from "./TreeNodeComponent";
+
+function generateTestNode(id: string, additionalProps?: Partial<TreeNode>): TreeNode {
+  const node: TreeNode = {
+    id,
+    label: `${id}Label`,
+    className: `${id}Class`,
+    iconSrc: `${id}Icon`,
+    onClick: jest.fn().mockName(`${id}Click`),
+    ...additionalProps,
+  };
+  return node;
+}
+
+describe("TreeNodeComponent", () => {
+  it("renders a single node", () => {
+    const node = generateTestNode("root");
+    const component = shallow(<TreeNodeComponent node={node} treeNodeId={node.id} />);
+    expect(component).toMatchSnapshot();
+
+    // The "click" handler is actually attached to onOpenChange, with a type of "Click".
+    component
+      .find(TreeItem)
+      .props()
+      .onOpenChange(null!, { open: true, value: "borp", target: null!, event: null!, type: "Click" });
+    expect(node.onClick).toHaveBeenCalled();
+  });
+
+  it("renders a node with a menu", () => {
+    const node = generateTestNode("root", {
+      contextMenu: [
+        {
+          label: "enabledItem",
+          onClick: jest.fn().mockName("enabledItemClick"),
+        },
+        {
+          label: "disabledItem",
+          onClick: jest.fn().mockName("disabledItemClick"),
+          isDisabled: true,
+        },
+      ],
+    });
+    const component = shallow(<TreeNodeComponent node={node} treeNodeId={node.id} />);
+    expect(component).toMatchSnapshot();
+  });
+
+  it("renders a loading spinner if the node is loading", async () => {
+    const loading = new PromiseSource();
+    const node = generateTestNode("root", {
+      onExpanded: () => loading.promise,
+    });
+    const component = shallow(<TreeNodeComponent node={node} treeNodeId={node.id} />);
+
+    act(() => {
+      component
+        .find(TreeItem)
+        .props()
+        .onOpenChange(null!, { open: true, value: "borp", target: null!, event: null!, type: "ExpandIconClick" });
+    });
+
+    expect(component).toMatchSnapshot("loading");
+    await loading.resolveAndWait();
+    expect(component).toMatchSnapshot("loaded");
+  });
+
+  it("renders single selected leaf node as selected", () => {
+    const node = generateTestNode("root", {
+      isSelected: () => true,
+    });
+    const component = shallow(<TreeNodeComponent node={node} treeNodeId={node.id} />);
+    expect(component.find(TreeItemLayout).props().style?.backgroundColor).toStrictEqual(
+      tokens.colorNeutralBackground1Selected,
+    );
+    expect(component).toMatchSnapshot();
+  });
+
+  it("renders selected parent node as selected if no descendant nodes are selected", () => {
+    const node = generateTestNode("root", {
+      isSelected: () => true,
+      children: [
+        generateTestNode("child1", {
+          children: [generateTestNode("grandchild1"), generateTestNode("grandchild2")],
+        }),
+        generateTestNode("child2"),
+      ],
+    });
+    const component = shallow(<TreeNodeComponent node={node} treeNodeId={node.id} />);
+    expect(component.find(TreeItemLayout).props().style?.backgroundColor).toStrictEqual(
+      tokens.colorNeutralBackground1Selected,
+    );
+    expect(component).toMatchSnapshot();
+  });
+
+  it("renders selected parent node as unselected if any descendant node is selected", () => {
+    const node = generateTestNode("root", {
+      isSelected: () => true,
+      children: [
+        generateTestNode("child1", {
+          children: [
+            generateTestNode("grandchild1", {
+              isSelected: () => true,
+            }),
+            generateTestNode("grandchild2"),
+          ],
+        }),
+        generateTestNode("child2"),
+      ],
+    });
+    const component = shallow(<TreeNodeComponent node={node} treeNodeId={node.id} />);
+    expect(component.find(TreeItemLayout).props().style?.backgroundColor).toBeUndefined();
+    expect(component).toMatchSnapshot();
+  });
+
+  it("renders an icon if the node has one", () => {
+    const node = generateTestNode("root", {
+      iconSrc: "the-icon.svg",
+    });
+    const component = shallow(<TreeNodeComponent node={node} treeNodeId={node.id} />);
+    expect(component).toMatchSnapshot();
+  });
+
+  it("does not render children if the node is loading", () => {
+    const node = generateTestNode("root", {
+      isLoading: true,
+      children: [
+        generateTestNode("child1", {
+          children: [generateTestNode("grandchild1"), generateTestNode("grandchild2")],
+        }),
+        generateTestNode("child2"),
+      ],
+    });
+    const component = shallow(<TreeNodeComponent node={node} treeNodeId={node.id} />);
+    expect(component).toMatchSnapshot();
+  });
+
+  it("fully renders a tree", () => {
+    const child3Loading = new PromiseSource();
+    const node = generateTestNode("root", {
+      isSelected: () => true,
+      children: [
+        generateTestNode("child1", {
+          children: [
+            generateTestNode("grandchild1", {
+              iconSrc: "grandchild1Icon.svg",
+              isSelected: () => true,
+            }),
+            generateTestNode("grandchild2"),
+          ],
+        }),
+        generateTestNode("child2Loading", {
+          isLoading: true,
+          children: [generateTestNode("grandchild3NotRendered")],
+        }),
+        generateTestNode("child3Expanding", {
+          onExpanded: () => child3Loading.promise,
+        }),
+      ],
+    });
+    const component = mount(<TreeNodeComponent node={node} treeNodeId={node.id} />);
+
+    // Find and expand the child3Expanding node
+    const expandingChild = component.find(TreeItem).filterWhere((n) => n.props().value === "root/child3ExpandingLabel");
+    act(() => {
+      expandingChild.props().onOpenChange(null!, {
+        open: true,
+        value: "root/child3ExpandingLabel",
+        target: null!,
+        event: null!,
+        type: "Click",
+      });
+    });
+
+    expect(component).toMatchSnapshot();
+  });
+});

src/Explorer/Controls/TreeComponent/TreeNodeComponent.tsx

@@ -0,0 +1,206 @@
+import {
+  Button,
+  Menu,
+  MenuItem,
+  MenuList,
+  MenuOpenChangeData,
+  MenuOpenEvent,
+  MenuPopover,
+  MenuTrigger,
+  Spinner,
+  Tree,
+  TreeItem,
+  TreeItemLayout,
+  TreeOpenChangeData,
+  TreeOpenChangeEvent,
+} from "@fluentui/react-components";
+import { MoreHorizontal20Regular } from "@fluentui/react-icons";
+import { tokens } from "@fluentui/react-theme";
+import * as React from "react";
+import { useCallback } from "react";
+
+export interface TreeNodeMenuItem {
+  label: string;
+  onClick: () => void;
+  iconSrc?: string;
+  isDisabled?: boolean;
+  styleClass?: string;
+}
+
+export interface TreeNode {
+  label: string;
+  id?: string;
+  children?: TreeNode[];
+  contextMenu?: TreeNodeMenuItem[];
+  iconSrc?: string;
+  isExpanded?: boolean;
+  className?: string;
+  isAlphaSorted?: boolean;
+  // data?: any; // Piece of data corresponding to this node
+  timestamp?: number;
+  isLeavesParentsSeparate?: boolean; // Display parents together first, then leaves
+  isLoading?: boolean;
+  isScrollable?: boolean;
+  isSelected?: () => boolean;
+  onClick?: () => void; // Only if a leaf, other click will expand/collapse
+  onExpanded?: () => Promise<void>;
+  onCollapsed?: () => void;
+  onContextMenuOpen?: () => void;
+}
+
+export interface TreeNodeComponentProps {
+  node: TreeNode;
+  className?: string;
+  treeNodeId: string;
+}
+
+/** Function that returns true if any descendant (at any depth) of this node is selected. */
+function isAnyDescendantSelected(node: TreeNode): boolean {
+  return (
+    node.children &&
+    node.children.reduce(
+      (previous: boolean, child: TreeNode) =>
+        previous || (child.isSelected && child.isSelected()) || isAnyDescendantSelected(child),
+      false,
+    )
+  );
+}
+
+const getTreeIcon = (iconSrc: string): JSX.Element => <img src={iconSrc} alt="" style={{ width: 16, height: 16 }} />;
+
+export const TreeNodeComponent: React.FC<TreeNodeComponentProps> = ({
+  node,
+  treeNodeId,
+}: TreeNodeComponentProps): JSX.Element => {
+  const [isLoading, setIsLoading] = React.useState<boolean>(false);
+
+  const getSortedChildren = (treeNode: TreeNode): TreeNode[] => {
+    if (!treeNode || !treeNode.children) {
+      return undefined;
+    }
+
+    const compareFct = (a: TreeNode, b: TreeNode) => a.label.localeCompare(b.label);
+
+    let unsortedChildren;
+    if (treeNode.isLeavesParentsSeparate) {
+      // Separate parents and leave
+      const parents: TreeNode[] = treeNode.children.filter((node) => node.children);
+      const leaves: TreeNode[] = treeNode.children.filter((node) => !node.children);
+
+      if (treeNode.isAlphaSorted) {
+        parents.sort(compareFct);
+        leaves.sort(compareFct);
+      }
+
+      unsortedChildren = parents.concat(leaves);
+    } else {
+      unsortedChildren = treeNode.isAlphaSorted ? treeNode.children.sort(compareFct) : treeNode.children;
+    }
+
+    return unsortedChildren;
+  };
+
+  const isBranch = node.children?.length > 0;
+
+  const onOpenChange = useCallback(
+    (_: TreeOpenChangeEvent, data: TreeOpenChangeData) => {
+      if (data.type === "Click" && !isBranch && node.onClick) {
+        node.onClick();
+      }
+      if (!node.isExpanded && data.open && node.onExpanded) {
+        // Catch the transition non-expanded to expanded
+        setIsLoading(true);
+        node.onExpanded?.().then(() => setIsLoading(false));
+      } else if (node.isExpanded && !data.open && node.onCollapsed) {
+        // Catch the transition expanded to non-expanded
+        node.onCollapsed?.();
+      }
+    },
+    [isBranch, node, setIsLoading],
+  );
+
+  const onMenuOpenChange = useCallback(
+    (e: MenuOpenEvent, data: MenuOpenChangeData) => {
+      if (data.open) {
+        node.onContextMenuOpen?.();
+      }
+    },
+    [node],
+  );
+
+  // We show a node as selected if it is selected AND no descendant is selected.
+  // We want to show only the deepest selected node as selected.
+  const isCurrentNodeSelected = node.isSelected && node.isSelected();
+  const shouldShowAsSelected = isCurrentNodeSelected && !isAnyDescendantSelected(node);
+
+  const contextMenuItems = (node.contextMenu ?? []).map((menuItem) => (
+    <MenuItem
+      data-test={`TreeNode/ContextMenuItem:${menuItem.label}`}
+      disabled={menuItem.isDisabled}
+      key={menuItem.label}
+      onClick={menuItem.onClick}
+    >
+      {menuItem.label}
+    </MenuItem>
+  ));
+
+  const treeItem = (
+    <TreeItem
+      value={treeNodeId}
+      itemType={isBranch ? "branch" : "leaf"}
+      style={{ height: "100%" }}
+      onOpenChange={onOpenChange}
+    >
+      <TreeItemLayout
+        className={node.className}
+        data-test={`TreeNode:${treeNodeId}`}
+        actions={
+          contextMenuItems.length > 0 && (
+            <Menu onOpenChange={onMenuOpenChange}>
+              <MenuTrigger disableButtonEnhancement>
+                <Button
+                  aria-label="More options"
+                  data-test="TreeNode/ContextMenuTrigger"
+                  appearance="subtle"
+                  icon={<MoreHorizontal20Regular />}
+                />
+              </MenuTrigger>
+              <MenuPopover data-test={`TreeNode/ContextMenu:${treeNodeId}`}>
+                <MenuList>{contextMenuItems}</MenuList>
+              </MenuPopover>
+            </Menu>
+          )
+        }
+        expandIcon={isLoading ? <Spinner size="extra-tiny" /> : undefined}
+        iconBefore={node.iconSrc && getTreeIcon(node.iconSrc)}
+        style={{
+          backgroundColor: shouldShowAsSelected ? tokens.colorNeutralBackground1Selected : undefined,
+        }}
+      >
+        {node.label}
+      </TreeItemLayout>
+      {!node.isLoading && node.children?.length > 0 && (
+        <Tree style={{ overflow: node.isScrollable ? "auto" : undefined }}>
+          {getSortedChildren(node).map((childNode: TreeNode) => (
+            <TreeNodeComponent key={childNode.label} node={childNode} treeNodeId={`${treeNodeId}/${childNode.label}`} />
+          ))}
+        </Tree>
+      )}
+    </TreeItem>
+  );
+
+  if (contextMenuItems.length === 0) {
+    return treeItem;
+  }
+
+  // For accessibility, it's highly recommended that any 'actions' also be available in the context menu.
+  // See https://react.fluentui.dev/?path=/docs/components-tree--default#actions
+  return (
+    <Menu positioning="below-end" openOnContext onOpenChange={onMenuOpenChange}>
+      <MenuTrigger disableButtonEnhancement>{treeItem}</MenuTrigger>
+      <MenuPopover>
+        <MenuList>{contextMenuItems}</MenuList>
+      </MenuPopover>
+    </Menu>
+  );
+};

src/Explorer/Controls/TreeComponent/__snapshots__/LegacyTreeComponent.test.tsx.snap

@@ -1,11 +1,11 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
-exports[`TreeComponent renders a simple tree 1`] = `
+exports[`LegacyTreeComponent renders a simple tree 1`] = `
 <div
   className="treeComponent tree"
   role="tree"
 >
-  <TreeNodeComponent
+  <LegacyTreeNodeComponent
     generation={0}
     node={
       Object {
@@ -33,7 +33,7 @@ exports[`TreeComponent renders a simple tree 1`] = `
 </div>
 `;
 
-exports[`TreeNodeComponent does not render children by default 1`] = `
+exports[`LegacyTreeNodeComponent does not render children by default 1`] = `
 <div
   className=" main2 nodeItem "
   onClick={[Function]}
@@ -102,7 +102,7 @@ exports[`TreeNodeComponent does not render children by default 1`] = `
       data-test="label"
       role="group"
     >
-      <TreeNodeComponent
+      <LegacyTreeNodeComponent
         generation={3}
         key="Bchild1-3-undefined"
         node={
@@ -120,7 +120,7 @@ exports[`TreeNodeComponent does not render children by default 1`] = `
         }
         paddingLeft={32}
       />
-      <TreeNodeComponent
+      <LegacyTreeNodeComponent
         generation={3}
         key="2child2-3-undefined"
         node={
@@ -135,7 +135,7 @@ exports[`TreeNodeComponent does not render children by default 1`] = `
 </div>
 `;
 
-exports[`TreeNodeComponent renders a simple node (sorted children, expanded) 1`] = `
+exports[`LegacyTreeNodeComponent renders a simple node (sorted children, expanded) 1`] = `
 <div
   className="nodeClassname main12 nodeItem "
   id="id"
@@ -254,7 +254,7 @@ exports[`TreeNodeComponent renders a simple node (sorted children, expanded) 1`]
       data-test="label"
       role="group"
     >
-      <TreeNodeComponent
+      <LegacyTreeNodeComponent
         generation={13}
         key="2child2-13-undefined"
         node={
@@ -264,7 +264,7 @@ exports[`TreeNodeComponent renders a simple node (sorted children, expanded) 1`]
         }
         paddingLeft={214}
       />
-      <TreeNodeComponent
+      <LegacyTreeNodeComponent
         generation={13}
         key="Bchild1-13-undefined"
         node={
@@ -287,7 +287,7 @@ exports[`TreeNodeComponent renders a simple node (sorted children, expanded) 1`]
 </div>
 `;
 
-exports[`TreeNodeComponent renders loading icon 1`] = `
+exports[`LegacyTreeNodeComponent renders loading icon 1`] = `
 <div
   className=" main2 nodeItem "
   onClick={[Function]}
@@ -360,7 +360,7 @@ exports[`TreeNodeComponent renders loading icon 1`] = `
 </div>
 `;
 
-exports[`TreeNodeComponent renders sorted children, expanded, leaves and parents separated 1`] = `
+exports[`LegacyTreeNodeComponent renders sorted children, expanded, leaves and parents separated 1`] = `
 <div
   className="nodeClassname main12 nodeItem "
   id="id"
@@ -470,7 +470,7 @@ exports[`TreeNodeComponent renders sorted children, expanded, leaves and parents
       data-test="label"
       role="group"
     >
-      <TreeNodeComponent
+      <LegacyTreeNodeComponent
         generation={13}
         key="bchild-13-undefined"
         node={
@@ -488,7 +488,7 @@ exports[`TreeNodeComponent renders sorted children, expanded, leaves and parents
         }
         paddingLeft={192}
       />
-      <TreeNodeComponent
+      <LegacyTreeNodeComponent
         generation={13}
         key="dchild-13-undefined"
         node={
@@ -506,7 +506,7 @@ exports[`TreeNodeComponent renders sorted children, expanded, leaves and parents
         }
         paddingLeft={192}
       />
-      <TreeNodeComponent
+      <LegacyTreeNodeComponent
         generation={13}
         key="aChild-13-undefined"
         node={
@@ -516,7 +516,7 @@ exports[`TreeNodeComponent renders sorted children, expanded, leaves and parents
         }
         paddingLeft={214}
       />
-      <TreeNodeComponent
+      <LegacyTreeNodeComponent
         generation={13}
         key="cchild-13-undefined"
         node={
@@ -531,7 +531,7 @@ exports[`TreeNodeComponent renders sorted children, expanded, leaves and parents
 </div>
 `;
 
-exports[`TreeNodeComponent renders unsorted children by default 1`] = `
+exports[`LegacyTreeNodeComponent renders unsorted children by default 1`] = `
 <div
   className=" main2 nodeItem "
   onClick={[Function]}
@@ -600,7 +600,7 @@ exports[`TreeNodeComponent renders unsorted children by default 1`] = `
       data-test="label"
       role="group"
     >
-      <TreeNodeComponent
+      <LegacyTreeNodeComponent
         generation={3}
         key="Bchild1-3-undefined"
         node={
@@ -618,7 +618,7 @@ exports[`TreeNodeComponent renders unsorted children by default 1`] = `
         }
         paddingLeft={32}
       />
-      <TreeNodeComponent
+      <LegacyTreeNodeComponent
         generation={3}
         key="2child2-3-undefined"
         node={

src/Explorer/Controls/TreeComponent2/TreeNode2Component.tsx

@@ -1,147 +0,0 @@
-import {
-  Button,
-  Menu,
-  MenuItem,
-  MenuList,
-  MenuPopover,
-  MenuTrigger,
-  Spinner,
-  Tree,
-  TreeItem,
-  TreeItemLayout,
-  TreeOpenChangeData,
-  TreeOpenChangeEvent,
-} from "@fluentui/react-components";
-import { MoreHorizontal20Regular } from "@fluentui/react-icons";
-import { tokens } from "@fluentui/react-theme";
-import * as React from "react";
-
-export interface TreeNode2MenuItem {
-  label: string;
-  onClick: () => void;
-  iconSrc?: string;
-  isDisabled?: boolean;
-  styleClass?: string;
-}
-
-export interface TreeNode2 {
-  label: string;
-  id?: string;
-  children?: TreeNode2[];
-  contextMenu?: TreeNode2MenuItem[];
-  iconSrc?: string;
-  isExpanded?: boolean;
-  className?: string;
-  isAlphaSorted?: boolean;
-  // data?: any; // Piece of data corresponding to this node
-  timestamp?: number;
-  isLeavesParentsSeparate?: boolean; // Display parents together first, then leaves
-  isLoading?: boolean;
-  isScrollable?: boolean;
-  isSelected?: () => boolean;
-  onClick?: () => void; // Only if a leaf, other click will expand/collapse
-  onExpanded?: () => Promise<void>;
-  onCollapsed?: () => void;
-  onContextMenuOpen?: () => void;
-}
-
-export interface TreeNode2ComponentProps {
-  node: TreeNode2;
-  className?: string;
-  treeNodeId: string;
-}
-
-const getTreeIcon = (iconSrc: string): JSX.Element => <img src={iconSrc} alt="" style={{ width: 20, height: 20 }} />;
-
-export const TreeNode2Component: React.FC<TreeNode2ComponentProps> = ({
-  node,
-  treeNodeId,
-}: TreeNode2ComponentProps): JSX.Element => {
-  const [isLoading, setIsLoading] = React.useState<boolean>(false);
-
-  const getSortedChildren = (treeNode: TreeNode2): TreeNode2[] => {
-    if (!treeNode || !treeNode.children) {
-      return undefined;
-    }
-
-    const compareFct = (a: TreeNode2, b: TreeNode2) => a.label.localeCompare(b.label);
-
-    let unsortedChildren;
-    if (treeNode.isLeavesParentsSeparate) {
-      // Separate parents and leave
-      const parents: TreeNode2[] = treeNode.children.filter((node) => node.children);
-      const leaves: TreeNode2[] = treeNode.children.filter((node) => !node.children);
-
-      if (treeNode.isAlphaSorted) {
-        parents.sort(compareFct);
-        leaves.sort(compareFct);
-      }
-
-      unsortedChildren = parents.concat(leaves);
-    } else {
-      unsortedChildren = treeNode.isAlphaSorted ? treeNode.children.sort(compareFct) : treeNode.children;
-    }
-
-    return unsortedChildren;
-  };
-
-  const onOpenChange = (_: TreeOpenChangeEvent, data: TreeOpenChangeData) => {
-    if (!node.isExpanded && data.open && node.onExpanded) {
-      // Catch the transition non-expanded to expanded
-      setIsLoading(true);
-      node.onExpanded?.().then(() => setIsLoading(false));
-    } else if (node.isExpanded && !data.open && node.onCollapsed) {
-      // Catch the transition expanded to non-expanded
-      node.onCollapsed?.();
-    }
-  };
-
-  return (
-    <TreeItem
-      value={treeNodeId}
-      itemType={node.children !== undefined ? "branch" : "leaf"}
-      style={{ height: "100%" }}
-      onOpenChange={onOpenChange}
-    >
-      <TreeItemLayout
-        className={node.className}
-        actions={
-          node.contextMenu && (
-            <Menu>
-              <MenuTrigger disableButtonEnhancement>
-                <Button aria-label="More options" appearance="subtle" icon={<MoreHorizontal20Regular />} />
-              </MenuTrigger>
-              <MenuPopover>
-                <MenuList>
-                  {node.contextMenu.map((menuItem) => (
-                    <MenuItem disabled={menuItem.isDisabled} key={menuItem.label} onClick={menuItem.onClick}>
-                      {menuItem.label}
-                    </MenuItem>
-                  ))}
-                </MenuList>
-              </MenuPopover>
-            </Menu>
-          )
-        }
-        expandIcon={isLoading ? <Spinner size="extra-tiny" /> : undefined}
-        iconBefore={node.iconSrc && getTreeIcon(node.iconSrc)}
-        style={{
-          backgroundColor: node.isSelected && node.isSelected() ? tokens.colorNeutralBackground1Selected : undefined,
-        }}
-      >
-        <span onClick={() => node.onClick?.()}>{node.label}</span>
-      </TreeItemLayout>
-      {!node.isLoading && node.children?.length > 0 && (
-        <Tree style={{ overflow: node.isScrollable ? "auto" : undefined }}>
-          {getSortedChildren(node).map((childNode: TreeNode2) => (
-            <TreeNode2Component
-              key={childNode.label}
-              node={childNode}
-              treeNodeId={`${treeNodeId}/${childNode.label}`}
-            />
-          ))}
-        </Tree>
-      )}
-    </TreeItem>
-  );
-};

src/Explorer/Explorer.tsx

@@ -136,14 +136,6 @@ export default class Explorer {
 
     this.isTabsContentExpanded = ko.observable(false);
 
-    document.addEventListener(
-      "contextmenu",
-      (e) => {
-        e.preventDefault();
-      },
-      false,
-    );
-
     $(() => {
       $(document.body).click(() => $(".commandDropdownContainer").hide());
     });

src/Explorer/Panes/AddCollectionPanel.tsx

@@ -832,8 +832,7 @@ export class AddCollectionPanel extends React.Component<AddCollectionPanelProps,
                     iconName="Info"
                     className="panelInfoIcon"
                     tabIndex={0}
-                    ariaLabel="Enable analytical store capability to perform near real-time analytics on your operational data, without
+                    ariaLabel="Enable analytical store capability to perform near real-time analytics on your operational data, without impacting the performance of transactional workloads."
-        impacting the performance of transactional workloads."
                   />
                 </TooltipHost>
               </Stack>

src/Explorer/Panes/PanelContainerComponent.tsx

@@ -53,6 +53,7 @@ export class PanelContainerComponent extends React.Component<PanelContainerProps
 
     return (
       <Panel
+        data-test={`Panel:${this.props.headerText}`}
         headerText={this.props.headerText}
         isOpen={this.props.isOpen}
         onDismiss={this.onDissmiss}

src/Explorer/Panes/__snapshots__/PanelContainerComponent.test.tsx.snap

@@ -4,6 +4,7 @@ exports[`PaneContainerComponent test should be resize if notification console is
 <StyledPanelBase
   closeButtonAriaLabel="Close test"
   customWidth="440px"
+  data-test="Panel:test"
   headerClassName="panelHeader"
   headerText="test"
   isFooterAtBottom={true}
@@ -44,6 +45,7 @@ exports[`PaneContainerComponent test should render with panel content and header
 <StyledPanelBase
   closeButtonAriaLabel="Close test"
   customWidth="440px"
+  data-test="Panel:test"
   headerClassName="panelHeader"
   headerText="test"
   isFooterAtBottom={true}

src/Explorer/Tree/ResourceTokenTree.tsx

@@ -1,65 +0,0 @@
-import React from "react";
-import CollectionIcon from "../../../images/tree-collection.svg";
-import * as ViewModels from "../../Contracts/ViewModels";
-import { useTabs } from "../../hooks/useTabs";
-import { userContext } from "../../UserContext";
-import { TreeComponent, TreeNode } from "../Controls/TreeComponent/TreeComponent";
-import { useCommandBar } from "../Menus/CommandBar/CommandBarComponentAdapter";
-import { mostRecentActivity } from "../MostRecentActivity/MostRecentActivity";
-import { useDatabases } from "../useDatabases";
-import { useSelectedNode } from "../useSelectedNode";
-
-export const ResourceTokenTree: React.FC = (): JSX.Element => {
-  const collection = useDatabases((state) => state.resourceTokenCollection);
-
-  const buildCollectionNode = (): TreeNode => {
-    if (!collection) {
-      return {
-        label: undefined,
-        isExpanded: true,
-        children: [],
-      };
-    }
-
-    const children: TreeNode[] = [];
-    children.push({
-      label: "Items",
-      onClick: () => {
-        collection.onDocumentDBDocumentsClick();
-        // push to most recent
-        mostRecentActivity.collectionWasOpened(userContext.databaseAccount?.id, collection);
-      },
-      isSelected: () =>
-        useSelectedNode
-          .getState()
-          .isDataNodeSelected(collection.databaseId, collection.id(), [ViewModels.CollectionTabKind.Documents]),
-    });
-
-    const collectionNode: TreeNode = {
-      label: collection.id(),
-      iconSrc: CollectionIcon,
-      isExpanded: true,
-      children,
-      className: "collectionHeader",
-      onClick: () => {
-        // Rewritten version of expandCollapseCollection
-        useSelectedNode.getState().setSelectedNode(collection);
-        useCommandBar.getState().setContextButtons([]);
-        useTabs
-          .getState()
-          .refreshActiveTab(
-            (tab) => tab.collection?.id() === collection.id() && tab.collection.databaseId === collection.databaseId,
-          );
-      },
-      isSelected: () => useSelectedNode.getState().isDataNodeSelected(collection.databaseId, collection.id()),
-    };
-
-    return {
-      label: undefined,
-      isExpanded: true,
-      children: [collectionNode],
-    };
-  };
-
-  return <TreeComponent className="dataResourceTree" rootNode={buildCollectionNode()} />;
-};

src/Explorer/Tree/ResourceTree.tsx

@@ -1,43 +1,29 @@
-import { SampleDataTree } from "Explorer/Tree/SampleDataTree";
+import {
-import { getItemName } from "Utils/APITypeUtils";
+  BrandVariants,
+  FluentProvider,
+  Theme,
+  Tree,
+  TreeItemValue,
+  TreeOpenChangeData,
+  TreeOpenChangeEvent,
+  createLightTheme,
+} from "@fluentui/react-components";
+import { AuthType } from "AuthType";
+import { TreeNode, TreeNodeComponent } from "Explorer/Controls/TreeComponent/TreeNodeComponent";
+import {
+  createDatabaseTreeNodes,
+  createResourceTokenTreeNodes,
+  createSampleDataTreeNodes,
+} from "Explorer/Tree/treeNodeUtil";
+import { useDatabases } from "Explorer/useDatabases";
+import { userContext } from "UserContext";
 import { useQueryCopilot } from "hooks/useQueryCopilot";
+import { useTabs } from "hooks/useTabs";
 import * as React from "react";
+import { useEffect, useMemo } from "react";
 import shallow from "zustand/shallow";
-import CosmosDBIcon from "../../../images/Azure-Cosmos-DB.svg";
-import DeleteIcon from "../../../images/delete.svg";
-import CopyIcon from "../../../images/notebook/Notebook-copy.svg";
-import NewNotebookIcon from "../../../images/notebook/Notebook-new.svg";
-import NotebookIcon from "../../../images/notebook/Notebook-resource.svg";
-import FileIcon from "../../../images/notebook/file-cosmos.svg";
-import PublishIcon from "../../../images/notebook/publish_content.svg";
-import RefreshIcon from "../../../images/refresh-cosmos.svg";
-import CollectionIcon from "../../../images/tree-collection.svg";
-import { isPublicInternetAccessAllowed } from "../../Common/DatabaseAccountUtility";
-import * as DataModels from "../../Contracts/DataModels";
-import * as ViewModels from "../../Contracts/ViewModels";
-import { Action, ActionModifiers, Source } from "../../Shared/Telemetry/TelemetryConstants";
-import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor";
-import { userContext } from "../../UserContext";
-import { isServerlessAccount } from "../../Utils/CapabilityUtils";
-import * as GitHubUtils from "../../Utils/GitHubUtils";
-import { useTabs } from "../../hooks/useTabs";
-import * as ResourceTreeContextMenuButtonFactory from "../ContextMenuButtonFactory";
-import { AccordionComponent, AccordionItemComponent } from "../Controls/Accordion/AccordionComponent";
-import { useDialog } from "../Controls/Dialog";
-import { TreeComponent, TreeNode, TreeNodeMenuItem } from "../Controls/TreeComponent/TreeComponent";
 import Explorer from "../Explorer";
-import { useCommandBar } from "../Menus/CommandBar/CommandBarComponentAdapter";
-import { mostRecentActivity } from "../MostRecentActivity/MostRecentActivity";
-import { NotebookContentItem, NotebookContentItemType } from "../Notebook/NotebookContentItem";
-import { NotebookUtil } from "../Notebook/NotebookUtil";
 import { useNotebook } from "../Notebook/useNotebook";
-import TabsBase from "../Tabs/TabsBase";
-import { useDatabases } from "../useDatabases";
-import { useSelectedNode } from "../useSelectedNode";
-import { Platform, configContext } from "./../../ConfigContext";
-import StoredProcedure from "./StoredProcedure";
-import Trigger from "./Trigger";
-import UserDefinedFunction from "./UserDefinedFunction";
 
 export const MyNotebooksTitle = "My Notebooks";
 export const GitHubReposTitle = "GitHub repos";
@@ -46,611 +32,153 @@ interface ResourceTreeProps {
   container: Explorer;
 }
 
+const cosmosdb: BrandVariants = {
+  10: "#020305",
+  20: "#111723",
+  30: "#16263D",
+  40: "#193253",
+  50: "#1B3F6A",
+  60: "#1B4C82",
+  70: "#18599B",
+  80: "#1267B4",
+  90: "#3174C2",
+  100: "#4F82C8",
+  110: "#6790CF",
+  120: "#7D9ED5",
+  130: "#92ACDC",
+  140: "#A6BAE2",
+  150: "#BAC9E9",
+  160: "#CDD8EF",
+};
+
+const lightTheme: Theme = {
+  ...createLightTheme(cosmosdb),
+};
+
+export const DATA_TREE_LABEL = "DATA";
+export const MY_DATA_TREE_LABEL = "MY DATA";
+export const SAMPLE_DATA_TREE_LABEL = "SAMPLE DATA";
+
+/**
+ * Top-level tree that has no label, but contains all subtrees
+ */
 export const ResourceTree: React.FC<ResourceTreeProps> = ({ container }: ResourceTreeProps): JSX.Element => {
-  const databases = useDatabases((state) => state.databases);
+  const [openItems, setOpenItems] = React.useState<Iterable<TreeItemValue>>([DATA_TREE_LABEL]);
-  const {
+
-    isNotebookEnabled,
+  const { isNotebookEnabled } = useNotebook(
-    myNotebooksContentRoot,
-    galleryContentRoot,
-    gitHubNotebooksContentRoot,
-    updateNotebookItem,
-  } = useNotebook(
     (state) => ({
       isNotebookEnabled: state.isNotebookEnabled,
-      myNotebooksContentRoot: state.myNotebooksContentRoot,
-      galleryContentRoot: state.galleryContentRoot,
-      gitHubNotebooksContentRoot: state.gitHubNotebooksContentRoot,
-      updateNotebookItem: state.updateNotebookItem,
     }),
     shallow,
   );
-  const { activeTab, refreshActiveTab } = useTabs();
-  const showScriptNodes =
-    configContext.platform !== Platform.Fabric && (userContext.apiType === "SQL" || userContext.apiType === "Gremlin");
-  const pseudoDirPath = "PsuedoDir";
 
-  const buildChildNodes = (
+  // We intentionally avoid using a state selector here because we want to re-render the tree if the active tab changes.
-    item: NotebookContentItem,
+  const { refreshActiveTab } = useTabs();
-    onFileClick: (item: NotebookContentItem) => void,
-    isGithubTree?: boolean,
-  ): TreeNode[] => {
-    if (!item || !item.children) {
-      return [];
-    } else {
-      return item.children.map((item) => {
-        const result =
-          item.type === NotebookContentItemType.Directory
-            ? buildNotebookDirectoryNode(item, onFileClick, isGithubTree)
-            : buildNotebookFileNode(item, onFileClick, isGithubTree);
-        result.timestamp = item.timestamp;
-        return result;
-      });
-    }
-  };
 
-  const buildNotebookFileNode = (
+  const { databases, resourceTokenCollection, sampleDataResourceTokenCollection } = useDatabases((state) => ({
-    item: NotebookContentItem,
+    databases: state.databases,
-    onFileClick: (item: NotebookContentItem) => void,
+    resourceTokenCollection: state.resourceTokenCollection,
-    isGithubTree?: boolean,
+    sampleDataResourceTokenCollection: state.sampleDataResourceTokenCollection,
-  ): TreeNode => {
+  }));
-    return {
+  const { isCopilotEnabled, isCopilotSampleDBEnabled } = useQueryCopilot((state) => ({
-      label: item.name,
+    isCopilotEnabled: state.copilotEnabled,
-      iconSrc: NotebookUtil.isNotebookFile(item.path) ? NotebookIcon : FileIcon,
+    isCopilotSampleDBEnabled: state.copilotSampleDBEnabled,
-      className: "notebookHeader",
+  }));
-      onClick: () => onFileClick(item),
-      isSelected: () => {
-        return (
-          activeTab &&
-          activeTab.tabKind === ViewModels.CollectionTabKind.NotebookV2 &&
-          /* TODO Redesign Tab interface so that resource tree doesn't need to know about NotebookV2Tab.
-                     NotebookV2Tab could be dynamically imported, but not worth it to just get this type right.
-                   */
-          (activeTab as any).notebookPath() === item.path
-        );
-      },
-      contextMenu: createFileContextMenu(container, item, isGithubTree),
-      data: item,
-    };
-  };
 
-  const createFileContextMenu = (
+  const databaseTreeNodes =
-    container: Explorer,
+    userContext.authType === AuthType.ResourceToken
-    item: NotebookContentItem,
+      ? createResourceTokenTreeNodes(resourceTokenCollection)
-    isGithubTree?: boolean,
+      : createDatabaseTreeNodes(container, isNotebookEnabled, databases, refreshActiveTab);
-  ): TreeNodeMenuItem[] => {
-    let items: TreeNodeMenuItem[] = [
-      {
-        label: "Rename",
-        iconSrc: NotebookIcon,
-        onClick: () => container.renameNotebook(item, isGithubTree),
-      },
-      {
-        label: "Delete",
-        iconSrc: DeleteIcon,
-        onClick: () => {
-          useDialog
-            .getState()
-            .showOkCancelModalDialog(
-              "Confirm delete",
-              `Are you sure you want to delete "${item.name}"`,
-              "Delete",
-              () => container.deleteNotebookFile(item, isGithubTree),
-              "Cancel",
-              undefined,
-            );
-        },
-      },
-      {
-        label: "Copy to ...",
-        iconSrc: CopyIcon,
-        onClick: () => copyNotebook(container, item),
-      },
-      {
-        label: "Download",
-        iconSrc: NotebookIcon,
-        onClick: () => container.downloadFile(item),
-      },
-    ];
 
-    if (item.type === NotebookContentItemType.Notebook && userContext.features.publicGallery) {
-      items.push({
-        label: "Publish to gallery",
-        iconSrc: PublishIcon,
-        onClick: async () => {
-          TelemetryProcessor.trace(Action.NotebooksGalleryClickPublishToGallery, ActionModifiers.Mark, {
-            source: Source.ResourceTreeMenu,
-          });
-
-          const content = await container.readFile(item);
-          if (content) {
-            await container.publishNotebook(item.name, content);
-          }
-        },
-      });
-    }
-
-    // "Copy to ..." isn't needed if github locations are not available
-    if (!container.notebookManager?.gitHubOAuthService.isLoggedIn()) {
-      items = items.filter((item) => item.label !== "Copy to ...");
-    }
-
-    return items;
-  };
-
-  const copyNotebook = async (container: Explorer, item: NotebookContentItem) => {
-    const content = await container.readFile(item);
-    if (content) {
-      container.copyNotebook(item.name, content);
-    }
-  };
-
-  const createDirectoryContextMenu = (
-    container: Explorer,
-    item: NotebookContentItem,
-    isGithubTree?: boolean,
-  ): TreeNodeMenuItem[] => {
-    let items: TreeNodeMenuItem[] = [
-      {
-        label: "Refresh",
-        iconSrc: RefreshIcon,
-        onClick: () => loadSubitems(item, isGithubTree),
-      },
-      {
-        label: "Delete",
-        iconSrc: DeleteIcon,
-        onClick: () => {
-          useDialog
-            .getState()
-            .showOkCancelModalDialog(
-              "Confirm delete",
-              `Are you sure you want to delete "${item.name}?"`,
-              "Delete",
-              () => container.deleteNotebookFile(item, isGithubTree),
-              "Cancel",
-              undefined,
-            );
-        },
-      },
-      {
-        label: "Rename",
-        iconSrc: NotebookIcon,
-        onClick: () => container.renameNotebook(item, isGithubTree),
-      },
-      {
-        label: "New Directory",
-        iconSrc: NewNotebookIcon,
-        onClick: () => container.onCreateDirectory(item, isGithubTree),
-      },
-    ];
-
-    //disallow renaming of temporary notebook workspace
-    if (item?.path === useNotebook.getState().notebookBasePath) {
-      items = items.filter((item) => item.label !== "Rename");
-    }
-
-    // For GitHub paths remove "Delete", "Rename", "New Directory", "Upload File"
-    if (GitHubUtils.fromContentUri(item.path)) {
-      items = items.filter(
-        (item) =>
-          item.label !== "Delete" &&
-          item.label !== "Rename" &&
-          item.label !== "New Directory" &&
-          item.label !== "Upload File",
-      );
-    }
-
-    return items;
-  };
-
-  const buildNotebookDirectoryNode = (
-    item: NotebookContentItem,
-    onFileClick: (item: NotebookContentItem) => void,
-    isGithubTree?: boolean,
-  ): TreeNode => {
-    return {
-      label: item.name,
-      iconSrc: undefined,
-      className: "notebookHeader",
-      isAlphaSorted: true,
-      isLeavesParentsSeparate: true,
-      onClick: () => {
-        if (!item.children) {
-          loadSubitems(item, isGithubTree);
-        }
-      },
-      isSelected: () => {
-        return (
-          activeTab &&
-          activeTab.tabKind === ViewModels.CollectionTabKind.NotebookV2 &&
-          /* TODO Redesign Tab interface so that resource tree doesn't need to know about NotebookV2Tab.
-                     NotebookV2Tab could be dynamically imported, but not worth it to just get this type right.
-                   */
-          (activeTab as any).notebookPath() === item.path
-        );
-      },
-      contextMenu: item.path !== pseudoDirPath ? createDirectoryContextMenu(container, item, isGithubTree) : undefined,
-      data: item,
-      children: buildChildNodes(item, onFileClick, isGithubTree),
-    };
-  };
-
-  const buildDataTree = (): TreeNode => {
-    const databaseTreeNodes: TreeNode[] = databases.map((database: ViewModels.Database) => {
-      const databaseNode: TreeNode = {
-        label: database.id(),
-        iconSrc: CosmosDBIcon,
-        isExpanded: database.isDatabaseExpanded(),
-        className: "databaseHeader",
-        children: [],
-        isSelected: () => useSelectedNode.getState().isDataNodeSelected(database.id()),
-        contextMenu: ResourceTreeContextMenuButtonFactory.createDatabaseContextMenu(container, database.id()),
-        onClick: async (isExpanded) => {
-          useSelectedNode.getState().setSelectedNode(database);
-          // Rewritten version of expandCollapseDatabase():
-          if (isExpanded) {
-            database.collapseDatabase();
-          } else {
-            if (databaseNode.children?.length === 0) {
-              databaseNode.isLoading = true;
-            }
-            await database.expandDatabase();
-          }
-          databaseNode.isLoading = false;
-          useCommandBar.getState().setContextButtons([]);
-          refreshActiveTab((tab: TabsBase) => tab.collection?.databaseId === database.id());
-        },
-        onContextMenuOpen: () => useSelectedNode.getState().setSelectedNode(database),
-      };
-
-      if (database.isDatabaseShared()) {
-        databaseNode.children.push({
-          id: database.isSampleDB ? "sampleScaleSettings" : "",
-          label: "Scale",
-          isSelected: () =>
-            useSelectedNode
-              .getState()
-              .isDataNodeSelected(database.id(), undefined, [ViewModels.CollectionTabKind.DatabaseSettingsV2]),
-          onClick: database.onSettingsClick.bind(database),
-        });
-      }
-
-      // Find collections
-      database
-        .collections()
-        .forEach((collection: ViewModels.Collection) =>
-          databaseNode.children.push(buildCollectionNode(database, collection)),
-        );
-
-      if (database.collectionsContinuationToken) {
-        const loadMoreNode: TreeNode = {
-          label: "load more",
-          className: "loadMoreHeader",
-          onClick: async () => {
-            await database.loadCollections();
-            useDatabases.getState().updateDatabase(database);
-          },
-        };
-        databaseNode.children.push(loadMoreNode);
-      }
-
-      database.collections.subscribe((collections: ViewModels.Collection[]) => {
-        collections.forEach((collection: ViewModels.Collection) =>
-          databaseNode.children.push(buildCollectionNode(database, collection)),
-        );
-      });
-
-      return databaseNode;
-    });
-
-    return {
-      label: undefined,
-      isExpanded: true,
-      children: databaseTreeNodes,
-    };
-  };
-
-  const buildCollectionNode = (database: ViewModels.Database, collection: ViewModels.Collection): TreeNode => {
-    const children: TreeNode[] = [];
-    children.push({
-      label: getItemName(),
-      id: collection.isSampleCollection ? "sampleItems" : "",
-      onClick: () => {
-        collection.openTab();
-        // push to most recent
-        mostRecentActivity.collectionWasOpened(userContext.databaseAccount?.id, collection);
-      },
-      isSelected: () =>
-        useSelectedNode
-          .getState()
-          .isDataNodeSelected(collection.databaseId, collection.id(), [
-            ViewModels.CollectionTabKind.Documents,
-            ViewModels.CollectionTabKind.Graph,
-          ]),
-      contextMenu: ResourceTreeContextMenuButtonFactory.createCollectionContextMenuButton(container, collection),
-    });
-
-    if (
-      isNotebookEnabled &&
-      userContext.apiType === "Mongo" &&
-      isPublicInternetAccessAllowed() &&
-      useNotebook.getState().isPhoenixFeatures
-    ) {
-      children.push({
-        label: "Schema (Preview)",
-        onClick: collection.onSchemaAnalyzerClick.bind(collection),
-        isSelected: () =>
-          useSelectedNode
-            .getState()
-            .isDataNodeSelected(collection.databaseId, collection.id(), [ViewModels.CollectionTabKind.SchemaAnalyzer]),
-      });
-    }
-
-    if (userContext.apiType !== "Cassandra" || !isServerlessAccount()) {
-      let id = "";
-      if (collection.isSampleCollection) {
-        id = database.isDatabaseShared() ? "sampleSettings" : "sampleScaleSettings";
-      }
-
-      children.push({
-        id,
-        label: database.isDatabaseShared() || isServerlessAccount() ? "Settings" : "Scale & Settings",
-        onClick: collection.onSettingsClick.bind(collection),
-        isSelected: () =>
-          useSelectedNode
-            .getState()
-            .isDataNodeSelected(collection.databaseId, collection.id(), [
-              ViewModels.CollectionTabKind.CollectionSettingsV2,
-            ]),
-      });
-    }
-
-    const schemaNode: TreeNode = buildSchemaNode(collection);
-    if (schemaNode) {
-      children.push(schemaNode);
-    }
-
-    if (showScriptNodes) {
-      children.push(buildStoredProcedureNode(collection));
-      children.push(buildUserDefinedFunctionsNode(collection));
-      children.push(buildTriggerNode(collection));
-    }
-
-    // This is a rewrite of showConflicts
-    const showConflicts =
-      userContext?.databaseAccount?.properties.enableMultipleWriteLocations &&
-      collection.rawDataModel &&
-      !!collection.rawDataModel.conflictResolutionPolicy;
-
-    if (showConflicts) {
-      children.push({
-        label: "Conflicts",
-        onClick: collection.onConflictsClick.bind(collection),
-        isSelected: () =>
-          useSelectedNode
-            .getState()
-            .isDataNodeSelected(collection.databaseId, collection.id(), [ViewModels.CollectionTabKind.Conflicts]),
-      });
-    }
-
-    return {
-      label: collection.id(),
-      iconSrc: CollectionIcon,
-      isExpanded: collection.isCollectionExpanded(),
-      children: children,
-      className: "collectionHeader",
-      contextMenu: ResourceTreeContextMenuButtonFactory.createCollectionContextMenuButton(container, collection),
-      onClick: () => {
-        // Rewritten version of expandCollapseCollection
-        useSelectedNode.getState().setSelectedNode(collection);
-        useCommandBar.getState().setContextButtons([]);
-        refreshActiveTab(
-          (tab: TabsBase) =>
-            tab.collection?.id() === collection.id() && tab.collection.databaseId === collection.databaseId,
-        );
-      },
-      onExpanded: () => {
-        if (showScriptNodes) {
-          collection.loadStoredProcedures();
-          collection.loadUserDefinedFunctions();
-          collection.loadTriggers();
-        }
-      },
-      isSelected: () => useSelectedNode.getState().isDataNodeSelected(collection.databaseId, collection.id()),
-      onContextMenuOpen: () => useSelectedNode.getState().setSelectedNode(collection),
-    };
-  };
-
-  const buildStoredProcedureNode = (collection: ViewModels.Collection): TreeNode => {
-    return {
-      label: "Stored Procedures",
-      children: collection.storedProcedures().map((sp: StoredProcedure) => ({
-        label: sp.id(),
-        onClick: sp.open.bind(sp),
-        isSelected: () =>
-          useSelectedNode
-            .getState()
-            .isDataNodeSelected(collection.databaseId, collection.id(), [
-              ViewModels.CollectionTabKind.StoredProcedures,
-            ]),
-        contextMenu: ResourceTreeContextMenuButtonFactory.createStoreProcedureContextMenuItems(container, sp),
-      })),
-      onClick: async () => {
-        await collection.loadStoredProcedures();
-        collection.selectedSubnodeKind(ViewModels.CollectionTabKind.StoredProcedures);
-        refreshActiveTab(
-          (tab: TabsBase) =>
-            tab.collection?.id() === collection.id() && tab.collection.databaseId === collection.databaseId,
-        );
-      },
-    };
-  };
-
-  const buildUserDefinedFunctionsNode = (collection: ViewModels.Collection): TreeNode => {
-    return {
-      label: "User Defined Functions",
-      children: collection.userDefinedFunctions().map((udf: UserDefinedFunction) => ({
-        label: udf.id(),
-        onClick: udf.open.bind(udf),
-        isSelected: () =>
-          useSelectedNode
-            .getState()
-            .isDataNodeSelected(collection.databaseId, collection.id(), [
-              ViewModels.CollectionTabKind.UserDefinedFunctions,
-            ]),
-        contextMenu: ResourceTreeContextMenuButtonFactory.createUserDefinedFunctionContextMenuItems(container, udf),
-      })),
-      onClick: async () => {
-        await collection.loadUserDefinedFunctions();
-        collection.selectedSubnodeKind(ViewModels.CollectionTabKind.UserDefinedFunctions);
-        refreshActiveTab(
-          (tab: TabsBase) =>
-            tab.collection?.id() === collection.id() && tab.collection.databaseId === collection.databaseId,
-        );
-      },
-    };
-  };
-
-  const buildTriggerNode = (collection: ViewModels.Collection): TreeNode => {
-    return {
-      label: "Triggers",
-      children: collection.triggers().map((trigger: Trigger) => ({
-        label: trigger.id(),
-        onClick: trigger.open.bind(trigger),
-        isSelected: () =>
-          useSelectedNode
-            .getState()
-            .isDataNodeSelected(collection.databaseId, collection.id(), [ViewModels.CollectionTabKind.Triggers]),
-        contextMenu: ResourceTreeContextMenuButtonFactory.createTriggerContextMenuItems(container, trigger),
-      })),
-      onClick: async () => {
-        await collection.loadTriggers();
-        collection.selectedSubnodeKind(ViewModels.CollectionTabKind.Triggers);
-        refreshActiveTab(
-          (tab: TabsBase) =>
-            tab.collection?.id() === collection.id() && tab.collection.databaseId === collection.databaseId,
-        );
-      },
-    };
-  };
-
-  const buildSchemaNode = (collection: ViewModels.Collection): TreeNode => {
-    if (collection.analyticalStorageTtl() === undefined) {
-      return undefined;
-    }
-
-    if (!collection.schema || !collection.schema.fields) {
-      return undefined;
-    }
-
-    return {
-      label: "Schema",
-      children: getSchemaNodes(collection.schema.fields),
-      onClick: () => {
-        collection.selectedSubnodeKind(ViewModels.CollectionTabKind.Schema);
-        refreshActiveTab((tab: TabsBase) => tab.collection && tab.collection.rid === collection.rid);
-      },
-    };
-  };
-
-  const getSchemaNodes = (fields: DataModels.IDataField[]): TreeNode[] => {
-    const schema: any = {};
-
-    //unflatten
-    fields.forEach((field: DataModels.IDataField) => {
-      const path: string[] = field.path.split(".");
-      const fieldProperties = [field.dataType.name, `HasNulls: ${field.hasNulls}`];
-      let current: any = {};
-      path.forEach((name: string, pathIndex: number) => {
-        if (pathIndex === 0) {
-          if (schema[name] === undefined) {
-            if (pathIndex === path.length - 1) {
-              schema[name] = fieldProperties;
-            } else {
-              schema[name] = {};
-            }
-          }
-          current = schema[name];
-        } else {
-          if (current[name] === undefined) {
-            if (pathIndex === path.length - 1) {
-              current[name] = fieldProperties;
-            } else {
-              current[name] = {};
-            }
-          }
-          current = current[name];
-        }
-      });
-    });
-
-    const traverse = (obj: any): TreeNode[] => {
-      const children: TreeNode[] = [];
-
-      if (obj !== undefined && !Array.isArray(obj) && typeof obj === "object") {
-        Object.entries(obj).forEach(([key, value]) => {
-          children.push({ label: key, children: traverse(value) });
-        });
-      } else if (Array.isArray(obj)) {
-        return [{ label: obj[0] }, { label: obj[1] }];
-      }
-
-      return children;
-    };
-
-    return traverse(schema);
-  };
-
-  const loadSubitems = async (item: NotebookContentItem, isGithubTree?: boolean): Promise<void> => {
-    const updatedItem = await container.notebookManager?.notebookContentClient?.updateItemChildren(item);
-    updateNotebookItem(updatedItem, isGithubTree);
-  };
-
-  const dataRootNode = buildDataTree();
   const isSampleDataEnabled =
-    useQueryCopilot().copilotEnabled &&
+    isCopilotEnabled &&
-    useQueryCopilot().copilotSampleDBEnabled &&
+    isCopilotSampleDBEnabled &&
     userContext.sampleDataConnectionInfo &&
     userContext.apiType === "SQL";
-  const sampleDataResourceTokenCollection = useDatabases((state) => state.sampleDataResourceTokenCollection);
+
+  const sampleDataNodes = useMemo<TreeNode[]>(() => {
+    return isSampleDataEnabled && sampleDataResourceTokenCollection
+      ? createSampleDataTreeNodes(sampleDataResourceTokenCollection)
+      : [];
+  }, [isSampleDataEnabled, sampleDataResourceTokenCollection]);
+
+  const rootNodes: TreeNode[] = useMemo(() => {
+    if (sampleDataNodes.length > 0) {
+      return [
+        {
+          id: "data",
+          label: MY_DATA_TREE_LABEL,
+          className: "accordionItemHeader",
+          children: databaseTreeNodes,
+          isScrollable: true,
+        },
+        {
+          id: "sampleData",
+          label: SAMPLE_DATA_TREE_LABEL,
+          className: "accordionItemHeader",
+          children: sampleDataNodes,
+        },
+      ];
+    } else {
+      return [
+        {
+          id: "data",
+          label: DATA_TREE_LABEL,
+          className: "accordionItemHeader",
+          children: databaseTreeNodes,
+          isScrollable: true,
+        },
+      ];
+    }
+  }, [databaseTreeNodes, sampleDataNodes]);
+
+  useEffect(() => {
+    // Compute open items based on node.isExpanded
+    const updateOpenItems = (node: TreeNode, parentNodeId: string): void => {
+      // This will look for ANY expanded node, event if its parent node isn't expanded
+      // and add it to the openItems list
+      const globalId = parentNodeId === undefined ? node.label : `${parentNodeId}/${node.label}`;
+
+      if (node.isExpanded) {
+        let found = false;
+        for (const id of openItems) {
+          if (id === globalId) {
+            found = true;
+            break;
+          }
+        }
+        if (!found) {
+          setOpenItems((prevOpenItems) => [...prevOpenItems, globalId]);
+        }
+      }
+
+      if (node.children) {
+        for (const child of node.children) {
+          updateOpenItems(child, globalId);
+        }
+      }
+    };
+
+    rootNodes.forEach((n) => updateOpenItems(n, undefined));
+  }, [rootNodes, openItems, setOpenItems]);
+
+  const handleOpenChange = (event: TreeOpenChangeEvent, data: TreeOpenChangeData) => setOpenItems(data.openItems);
 
   return (
     <>
-      {!isNotebookEnabled && !isSampleDataEnabled && (
+      <FluentProvider theme={lightTheme} style={{ overflow: "hidden" }}>
-        <TreeComponent className="dataResourceTree" rootNode={dataRootNode} />
+        <Tree
-      )}
+          aria-label="CosmosDB resources"
-      {isNotebookEnabled && !isSampleDataEnabled && (
+          openItems={openItems}
-        <>
+          onOpenChange={handleOpenChange}
-          <AccordionComponent>
+          size="small"
-            <AccordionItemComponent title={"DATA"} isExpanded={!gitHubNotebooksContentRoot}>
+          style={{ height: "100%", minWidth: "290px" }}
-              <TreeComponent className="dataResourceTree" rootNode={dataRootNode} />
+        >
-            </AccordionItemComponent>
+          {rootNodes.map((node) => (
-          </AccordionComponent>
+            <TreeNodeComponent key={node.label} className="dataResourceTree" node={node} treeNodeId={node.label} />
-        </>
+          ))}
-      )}
+        </Tree>
-      {!isNotebookEnabled && isSampleDataEnabled && (
+      </FluentProvider>
-        <>
-          <AccordionComponent>
-            <AccordionItemComponent title={"MY DATA"} isExpanded={!gitHubNotebooksContentRoot}>
-              <TreeComponent className="dataResourceTree" rootNode={dataRootNode} />
-            </AccordionItemComponent>
-            <AccordionItemComponent title={"SAMPLE DATA"} containerStyles={{ display: "table" }}>
-              <SampleDataTree sampleDataResourceTokenCollection={sampleDataResourceTokenCollection} />
-            </AccordionItemComponent>
-          </AccordionComponent>
-        </>
-      )}
-      {isNotebookEnabled && isSampleDataEnabled && (
-        <>
-          <AccordionComponent>
-            <AccordionItemComponent title={"MY DATA"} isExpanded={!gitHubNotebooksContentRoot}>
-              <TreeComponent className="dataResourceTree" rootNode={dataRootNode} />
-            </AccordionItemComponent>
-            <AccordionItemComponent title={"SAMPLE DATA"} containerStyles={{ display: "table" }}>
-              <SampleDataTree sampleDataResourceTokenCollection={sampleDataResourceTokenCollection} />
-            </AccordionItemComponent>
-          </AccordionComponent>
-        </>
-      )}
     </>
   );
 };

src/Explorer/Tree/ResourceTreeAdapter.test.tsx

@@ -2,7 +2,11 @@ import { shallow } from "enzyme";
 import React from "react";
 import * as DataModels from "../../Contracts/DataModels";
 import * as ViewModels from "../../Contracts/ViewModels";
-import { TreeComponent, TreeComponentProps, TreeNode } from "../Controls/TreeComponent/TreeComponent";
+import {
+  LegacyTreeComponent,
+  LegacyTreeComponentProps,
+  LegacyTreeNode,
+} from "../Controls/TreeComponent/LegacyTreeComponent";
 import Explorer from "../Explorer";
 import Collection from "./Collection";
 import { ResourceTreeAdapter } from "./ResourceTreeAdapter";
@@ -225,12 +229,12 @@ describe("Resource tree for schema", () => {
   const resourceTree = new ResourceTreeAdapter(mockContainer);
 
   it("should render", () => {
-    const rootNode: TreeNode = resourceTree.buildSchemaNode(createMockCollection());
+    const rootNode: LegacyTreeNode = resourceTree.buildSchemaNode(createMockCollection());
-    const props: TreeComponentProps = {
+    const props: LegacyTreeComponentProps = {
       rootNode,
       className: "dataResourceTree",
     };
-    const wrapper = shallow(<TreeComponent {...props} />);
+    const wrapper = shallow(<LegacyTreeComponent {...props} />);
     expect(wrapper).toMatchSnapshot();
   });
 });

src/Explorer/Tree/ResourceTreeAdapter.tsx

@@ -1,3 +1,5 @@
+import { TreeNodeMenuItem } from "Explorer/Controls/TreeComponent/TreeNodeComponent";
+import { shouldShowScriptNodes } from "Explorer/Tree/treeNodeUtil";
 import { getItemName } from "Utils/APITypeUtils";
 import * as ko from "knockout";
 import * as React from "react";
@@ -23,7 +25,7 @@ import * as GitHubUtils from "../../Utils/GitHubUtils";
 import { useTabs } from "../../hooks/useTabs";
 import * as ResourceTreeContextMenuButtonFactory from "../ContextMenuButtonFactory";
 import { useDialog } from "../Controls/Dialog";
-import { TreeComponent, TreeNode, TreeNodeMenuItem } from "../Controls/TreeComponent/TreeComponent";
+import { LegacyTreeComponent, LegacyTreeNode } from "../Controls/TreeComponent/LegacyTreeComponent";
 import Explorer from "../Explorer";
 import { useCommandBar } from "../Menus/CommandBar/CommandBarComponentAdapter";
 import { mostRecentActivity } from "../MostRecentActivity/MostRecentActivity";
@@ -33,7 +35,6 @@ import { useNotebook } from "../Notebook/useNotebook";
 import TabsBase from "../Tabs/TabsBase";
 import { useDatabases } from "../useDatabases";
 import { useSelectedNode } from "../useSelectedNode";
-import { Platform, configContext } from "./../../ConfigContext";
 import StoredProcedure from "./StoredProcedure";
 import Trigger from "./Trigger";
 import UserDefinedFunction from "./UserDefinedFunction";
@@ -95,7 +96,7 @@ export class ResourceTreeAdapter implements ReactAdapter {
 
   public renderComponent(): JSX.Element {
     const dataRootNode = this.buildDataTree();
-    return <TreeComponent className="dataResourceTree" rootNode={dataRootNode} />;
+    return <LegacyTreeComponent className="dataResourceTree" rootNode={dataRootNode} />;
   }
 
   public async initialize(): Promise<void[]> {
@@ -157,9 +158,11 @@ export class ResourceTreeAdapter implements ReactAdapter {
     }
   }
 
-  private buildDataTree(): TreeNode {
+  private buildDataTree(): LegacyTreeNode {
-    const databaseTreeNodes: TreeNode[] = useDatabases.getState().databases.map((database: ViewModels.Database) => {
+    const databaseTreeNodes: LegacyTreeNode[] = useDatabases
-      const databaseNode: TreeNode = {
+      .getState()
+      .databases.map((database: ViewModels.Database) => {
+        const databaseNode: LegacyTreeNode = {
           label: database.id(),
           iconSrc: CosmosDBIcon,
           isExpanded: false,
@@ -219,18 +222,8 @@ export class ResourceTreeAdapter implements ReactAdapter {
     };
   }
 
-  /**
+  private buildCollectionNode(database: ViewModels.Database, collection: ViewModels.Collection): LegacyTreeNode {
-   * This is a rewrite of Collection.ts : showScriptsMenu, showStoredProcedures, showTriggers, showUserDefinedFunctions
+    const children: LegacyTreeNode[] = [];
-   * @param container
-   */
-  private static showScriptNodes(container: Explorer): boolean {
-    return (
-      configContext.platform !== Platform.Fabric && (userContext.apiType === "SQL" || userContext.apiType === "Gremlin")
-    );
-  }
-
-  private buildCollectionNode(database: ViewModels.Database, collection: ViewModels.Collection): TreeNode {
-    const children: TreeNode[] = [];
     children.push({
       label: getItemName(),
       onClick: () => {
@@ -274,12 +267,12 @@ export class ResourceTreeAdapter implements ReactAdapter {
       });
     }
 
-    const schemaNode: TreeNode = this.buildSchemaNode(collection);
+    const schemaNode: LegacyTreeNode = this.buildSchemaNode(collection);
     if (schemaNode) {
       children.push(schemaNode);
     }
 
-    if (ResourceTreeAdapter.showScriptNodes(this.container)) {
+    if (shouldShowScriptNodes()) {
       children.push(this.buildStoredProcedureNode(collection));
       children.push(this.buildUserDefinedFunctionsNode(collection));
       children.push(this.buildTriggerNode(collection));
@@ -321,7 +314,7 @@ export class ResourceTreeAdapter implements ReactAdapter {
           );
       },
       onExpanded: () => {
-        if (ResourceTreeAdapter.showScriptNodes(this.container)) {
+        if (shouldShowScriptNodes()) {
           collection.loadStoredProcedures();
           collection.loadUserDefinedFunctions();
           collection.loadTriggers();
@@ -332,7 +325,7 @@ export class ResourceTreeAdapter implements ReactAdapter {
     };
   }
 
-  private buildStoredProcedureNode(collection: ViewModels.Collection): TreeNode {
+  private buildStoredProcedureNode(collection: ViewModels.Collection): LegacyTreeNode {
     return {
       label: "Stored Procedures",
       children: collection.storedProcedures().map((sp: StoredProcedure) => ({
@@ -358,7 +351,7 @@ export class ResourceTreeAdapter implements ReactAdapter {
     };
   }
 
-  private buildUserDefinedFunctionsNode(collection: ViewModels.Collection): TreeNode {
+  private buildUserDefinedFunctionsNode(collection: ViewModels.Collection): LegacyTreeNode {
     return {
       label: "User Defined Functions",
       children: collection.userDefinedFunctions().map((udf: UserDefinedFunction) => ({
@@ -387,7 +380,7 @@ export class ResourceTreeAdapter implements ReactAdapter {
     };
   }
 
-  private buildTriggerNode(collection: ViewModels.Collection): TreeNode {
+  private buildTriggerNode(collection: ViewModels.Collection): LegacyTreeNode {
     return {
       label: "Triggers",
       children: collection.triggers().map((trigger: Trigger) => ({
@@ -411,7 +404,7 @@ export class ResourceTreeAdapter implements ReactAdapter {
     };
   }
 
-  public buildSchemaNode(collection: ViewModels.Collection): TreeNode {
+  public buildSchemaNode(collection: ViewModels.Collection): LegacyTreeNode {
     if (collection.analyticalStorageTtl() == undefined) {
       return undefined;
     }
@@ -430,7 +423,7 @@ export class ResourceTreeAdapter implements ReactAdapter {
     };
   }
 
-  private getSchemaNodes(fields: DataModels.IDataField[]): TreeNode[] {
+  private getSchemaNodes(fields: DataModels.IDataField[]): LegacyTreeNode[] {
     const schema: any = {};
 
     //unflatten
@@ -461,8 +454,8 @@ export class ResourceTreeAdapter implements ReactAdapter {
       });
     });
 
-    const traverse = (obj: any): TreeNode[] => {
+    const traverse = (obj: any): LegacyTreeNode[] => {
-      const children: TreeNode[] = [];
+      const children: LegacyTreeNode[] = [];
 
       if (obj !== null && !Array.isArray(obj) && typeof obj === "object") {
         Object.entries(obj).forEach(([key, value]) => {
@@ -483,7 +476,7 @@ export class ResourceTreeAdapter implements ReactAdapter {
     onFileClick: (item: NotebookContentItem) => void,
     createDirectoryContextMenu: boolean,
     createFileContextMenu: boolean,
-  ): TreeNode[] {
+  ): LegacyTreeNode[] {
     if (!item || !item.children) {
       return [];
     } else {
@@ -502,7 +495,7 @@ export class ResourceTreeAdapter implements ReactAdapter {
     item: NotebookContentItem,
     onFileClick: (item: NotebookContentItem) => void,
     createFileContextMenu: boolean,
-  ): TreeNode {
+  ): LegacyTreeNode {
     return {
       label: item.name,
       iconSrc: NotebookUtil.isNotebookFile(item.path) ? NotebookIcon : FileIcon,
@@ -650,7 +643,7 @@ export class ResourceTreeAdapter implements ReactAdapter {
     onFileClick: (item: NotebookContentItem) => void,
     createDirectoryContextMenu: boolean,
     createFileContextMenu: boolean,
-  ): TreeNode {
+  ): LegacyTreeNode {
     return {
       label: item.name,
       iconSrc: undefined,

src/Explorer/Tree/SampleDataTree.tsx

@@ -7,15 +7,15 @@ import CosmosDBIcon from "../../../images/Azure-Cosmos-DB.svg";
 import CollectionIcon from "../../../images/tree-collection.svg";
 import * as ViewModels from "../../Contracts/ViewModels";
 import * as ResourceTreeContextMenuButtonFactory from "../ContextMenuButtonFactory";
-import { TreeComponent, TreeNode } from "../Controls/TreeComponent/TreeComponent";
+import { LegacyTreeComponent, LegacyTreeNode } from "../Controls/TreeComponent/LegacyTreeComponent";
 
 export const SampleDataTree = ({
   sampleDataResourceTokenCollection,
 }: {
   sampleDataResourceTokenCollection: ViewModels.CollectionBase;
 }): JSX.Element => {
-  const buildSampleDataTree = (): TreeNode => {
+  const buildSampleDataTree = (): LegacyTreeNode => {
-    const updatedSampleTree: TreeNode = {
+    const updatedSampleTree: LegacyTreeNode = {
       label: sampleDataResourceTokenCollection.databaseId,
       isExpanded: false,
       iconSrc: CosmosDBIcon,
@@ -70,7 +70,7 @@ export const SampleDataTree = ({
   };
 
   return (
-    <TreeComponent
+    <LegacyTreeComponent
       className="dataResourceTree"
       rootNode={sampleDataResourceTokenCollection ? buildSampleDataTree() : { label: "Sample data not initialized." }}
     />

src/Explorer/Tree/__snapshots__/ResourceTreeAdapter.test.tsx.snap

@@ -5,7 +5,7 @@ exports[`Resource tree for schema should render 1`] = `
   className="treeComponent dataResourceTree"
   role="tree"
 >
-  <TreeNodeComponent
+  <LegacyTreeNodeComponent
     generation={0}
     node={
       Object {

src/Explorer/Tree/treeNodeUtil.test.ts

@@ -0,0 +1,602 @@
+import { CapabilityNames } from "Common/Constants";
+import { Platform, updateConfigContext } from "ConfigContext";
+import { TreeNode } from "Explorer/Controls/TreeComponent/TreeNodeComponent";
+import Explorer from "Explorer/Explorer";
+import { useCommandBar } from "Explorer/Menus/CommandBar/CommandBarComponentAdapter";
+import { useNotebook } from "Explorer/Notebook/useNotebook";
+import { DeleteDatabaseConfirmationPanel } from "Explorer/Panes/DeleteDatabaseConfirmationPanel";
+import TabsBase from "Explorer/Tabs/TabsBase";
+import StoredProcedure from "Explorer/Tree/StoredProcedure";
+import Trigger from "Explorer/Tree/Trigger";
+import UserDefinedFunction from "Explorer/Tree/UserDefinedFunction";
+import {
+  createDatabaseTreeNodes,
+  createResourceTokenTreeNodes,
+  createSampleDataTreeNodes,
+} from "Explorer/Tree/treeNodeUtil";
+import { useDatabases } from "Explorer/useDatabases";
+import { useSelectedNode } from "Explorer/useSelectedNode";
+import { updateUserContext } from "UserContext";
+import PromiseSource from "Utils/PromiseSource";
+import { useSidePanel } from "hooks/useSidePanel";
+import { useTabs } from "hooks/useTabs";
+import ko from "knockout";
+import * as DataModels from "../../Contracts/DataModels";
+import * as ViewModels from "../../Contracts/ViewModels";
+
+jest.mock("Explorer/Explorer", () => {
+  // eslint-disable-next-line @typescript-eslint/no-extraneous-class
+  class MockExplorer {
+    onNewCollectionClicked = jest.fn();
+  }
+
+  return MockExplorer;
+});
+
+jest.mock("Explorer/Tree/StoredProcedure", () => {
+  let counter = 0;
+  class MockStoredProcedure {
+    id: () => string;
+    open = jest.fn();
+    delete = jest.fn();
+    constructor() {
+      this.id = () => `mockSproc${counter}`;
+      counter++;
+    }
+  }
+
+  return MockStoredProcedure;
+});
+
+jest.mock("Explorer/Tree/UserDefinedFunction", () => {
+  let counter = 0;
+  class MockUserDefinedFunction {
+    id: () => string;
+    open = jest.fn();
+    delete = jest.fn();
+    constructor() {
+      this.id = () => `mockUdf${counter}`;
+      counter++;
+    }
+  }
+
+  return MockUserDefinedFunction;
+});
+
+jest.mock("Explorer/Tree/Trigger", () => {
+  let counter = 0;
+  class MockTrigger {
+    id: () => string;
+    open = jest.fn();
+    delete = jest.fn();
+    constructor() {
+      this.id = () => `mockTrigger${counter}`;
+      counter++;
+    }
+  }
+
+  return MockTrigger;
+});
+
+jest.mock("Common/DatabaseAccountUtility", () => {
+  return {
+    isPublicInternetAccessAllowed: () => true,
+  };
+});
+
+// Defining this value outside the mock, AND prefixing the name with 'mock' is required by Jest's mocking logic.
+let nextTabIndex = 1;
+class MockTab extends TabsBase {
+  constructor(tabOptions: Pick<ViewModels.TabOptions, "tabKind"> & Partial<ViewModels.TabOptions>) {
+    super({
+      title: `Mock Tab ${nextTabIndex}`,
+      tabPath: `mockTabs/tab${nextTabIndex}`,
+      ...tabOptions,
+    } as ViewModels.TabOptions);
+    nextTabIndex++;
+  }
+
+  onActivate = jest.fn();
+}
+
+/** A basic test collection that can be expanded on in tests. */
+const baseCollection = {
+  container: new Explorer(),
+  databaseId: "testDatabase",
+  id: ko.observable<string>("testCollection"),
+  defaultTtl: ko.observable<number>(5),
+  analyticalStorageTtl: ko.observable<number>(undefined),
+  selectedSubnodeKind: ko.observable<ViewModels.CollectionTabKind>(),
+  indexingPolicy: ko.observable<DataModels.IndexingPolicy>({
+    automatic: true,
+    indexingMode: "consistent",
+    includedPaths: [],
+    excludedPaths: [],
+  }),
+  uniqueKeyPolicy: {} as DataModels.UniqueKeyPolicy,
+  usageSizeInKB: ko.observable(100),
+  offer: ko.observable<DataModels.Offer>({
+    autoscaleMaxThroughput: undefined,
+    manualThroughput: 10000,
+    minimumThroughput: 6000,
+    id: "offer",
+    offerReplacePending: false,
+  }),
+  conflictResolutionPolicy: ko.observable<DataModels.ConflictResolutionPolicy>(
+    {} as DataModels.ConflictResolutionPolicy,
+  ),
+  changeFeedPolicy: ko.observable<DataModels.ChangeFeedPolicy>({} as DataModels.ChangeFeedPolicy),
+  geospatialConfig: ko.observable<DataModels.GeospatialConfig>({} as DataModels.GeospatialConfig),
+  getDatabase: () => {},
+  partitionKey: {
+    paths: [],
+    kind: "hash",
+    version: 2,
+  },
+  storedProcedures: ko.observableArray([]),
+  userDefinedFunctions: ko.observableArray([]),
+  triggers: ko.observableArray([]),
+  partitionKeyProperties: ["testPartitionKey"],
+  readSettings: () => {},
+  isCollectionExpanded: ko.observable(true),
+  onSettingsClick: jest.fn(),
+  onDocumentDBDocumentsClick: jest.fn(),
+  onNewQueryClick: jest.fn(),
+  onConflictsClick: jest.fn(),
+  onSchemaAnalyzerClick: jest.fn(),
+} as unknown as ViewModels.Collection;
+
+/** A basic test database that can be expanded on in tests */
+const baseDatabase = {
+  container: new Explorer(),
+  id: ko.observable<string>("testDatabase"),
+  collections: ko.observableArray<ViewModels.Collection>([]),
+  isDatabaseShared: ko.pureComputed(() => false),
+  isDatabaseExpanded: ko.observable(true),
+  selectedSubnodeKind: ko.observable<ViewModels.CollectionTabKind>(),
+  expandDatabase: jest.fn().mockResolvedValue({}),
+  collapseDatabase: jest.fn(),
+  onSettingsClick: jest.fn(),
+} as unknown as ViewModels.Database;
+
+/** Configures app state so that useSelectedNode.getState().isDataNodeSelected() returns true for the provided arguments. */
+function selectDataNode(
+  node: ViewModels.Database | ViewModels.CollectionBase,
+  subnodeKind?: ViewModels.CollectionTabKind,
+) {
+  useSelectedNode.getState().setSelectedNode(node);
+
+  if (subnodeKind !== undefined) {
+    node.selectedSubnodeKind(subnodeKind);
+    useTabs.getState().activateNewTab(new MockTab({ tabKind: subnodeKind, node }));
+  }
+}
+
+describe("createSampleDataTreeNodes", () => {
+  let sampleDataResourceTokenCollection: ViewModels.Collection;
+  let nodes: TreeNode[];
+
+  beforeEach(() => {
+    jest.resetAllMocks();
+    const collection = { ...baseCollection };
+    useDatabases.setState({ sampleDataResourceTokenCollection: collection });
+    useSelectedNode.setState({ selectedNode: undefined });
+
+    sampleDataResourceTokenCollection = collection;
+    nodes = createSampleDataTreeNodes(sampleDataResourceTokenCollection);
+  });
+
+  it("creates the expected tree nodes", () => {
+    expect(nodes).toMatchSnapshot();
+  });
+});
+
+describe("createResourceTokenTreeNodes", () => {
+  let resourceTokenCollection: ViewModels.Collection;
+  let nodes: TreeNode[];
+
+  beforeEach(() => {
+    jest.resetAllMocks();
+    const collection = { ...baseCollection };
+    useDatabases.setState({ resourceTokenCollection: collection });
+    useSelectedNode.setState({ selectedNode: undefined });
+
+    resourceTokenCollection = collection;
+    nodes = createResourceTokenTreeNodes(resourceTokenCollection);
+  });
+
+  it("returns an empty node when collection is undefined or null", () => {
+    const snapshot = `
+Array [
+  Object {
+    "children": Array [],
+    "isExpanded": true,
+    "label": "",
+  },
+]
+`;
+    expect(createResourceTokenTreeNodes(undefined)).toMatchInlineSnapshot(snapshot);
+    expect(createResourceTokenTreeNodes(null)).toMatchInlineSnapshot(snapshot);
+  });
+
+  it("creates the expected tree nodes", () => {
+    expect(nodes).toMatchSnapshot();
+  });
+});
+
+describe("createDatabaseTreeNodes", () => {
+  let explorer: Explorer;
+  let standardDb: ViewModels.Database;
+  let sharedDb: ViewModels.Database;
+  let giganticDb: ViewModels.Database;
+  let standardCollection: ViewModels.Collection;
+  let sampleItemsCollection: ViewModels.Collection;
+  let schemaCollection: ViewModels.Collection;
+  let conflictsCollection: ViewModels.Collection;
+  let sproc: StoredProcedure;
+  let udf: UserDefinedFunction;
+  let trigger: Trigger;
+  let refreshActiveTab: (comparator: (tab: TabsBase) => boolean) => void;
+
+  beforeEach(() => {
+    jest.resetAllMocks();
+    explorer = new Explorer();
+    standardDb = {
+      ...baseDatabase,
+      id: ko.observable("standardDb"),
+      container: explorer,
+    } as ViewModels.Database;
+    sharedDb = {
+      ...baseDatabase,
+      id: ko.observable("sharedDatabase"),
+      container: explorer,
+      isDatabaseShared: ko.pureComputed(() => true),
+    } as ViewModels.Database;
+    giganticDb = {
+      ...baseDatabase,
+      id: ko.observable("giganticDatabase"),
+      container: explorer,
+      collectionsContinuationToken: "continuationToken",
+    } as ViewModels.Database;
+
+    standardCollection = {
+      ...baseCollection,
+      id: ko.observable("standardCollection"),
+      container: explorer,
+      databaseId: standardDb.id(),
+    } as ViewModels.Collection;
+
+    // These classes are mocked, so the constructor args don't matter
+    sproc = new StoredProcedure(explorer, standardCollection, {} as never);
+    standardCollection.storedProcedures = ko.pureComputed(() => [sproc]);
+    udf = new UserDefinedFunction(explorer, standardCollection, {} as never);
+    standardCollection.userDefinedFunctions = ko.pureComputed(() => [udf]);
+    trigger = new Trigger(explorer, standardCollection, {} as never);
+    standardCollection.triggers = ko.pureComputed(() => [trigger]);
+
+    sampleItemsCollection = {
+      ...baseCollection,
+      id: ko.observable("sampleItemsCollection"),
+      container: explorer,
+      databaseId: sharedDb.id(),
+      isSampleCollection: true,
+    } as ViewModels.Collection;
+
+    schemaCollection = {
+      ...baseCollection,
+      id: ko.observable("schemaCollection"),
+      container: explorer,
+      databaseId: sharedDb.id(),
+      analyticalStorageTtl: ko.observable<number>(5),
+      schema: {
+        fields: [
+          {
+            path: "address.street",
+            dataType: { name: "string" },
+            hasNulls: false,
+          },
+          {
+            path: "address.line2",
+            dataType: { name: "string" },
+            hasNulls: true,
+          },
+          {
+            path: "address.zip",
+            dataType: { name: "number" },
+            hasNulls: false,
+          },
+          {
+            path: "orderId",
+            dataType: { name: "string" },
+            hasNulls: false,
+          },
+        ],
+      } as unknown,
+    } as ViewModels.Collection;
+
+    conflictsCollection = {
+      ...baseCollection,
+      id: ko.observable("conflictsCollection"),
+      rawDataModel: {
+        conflictResolutionPolicy: {
+          mode: "Custom",
+          conflictResolutionPath: "path",
+          conflictResolutionProcedure: "proc",
+        },
+      },
+    } as ViewModels.Collection;
+
+    standardDb.collections = ko.observableArray([standardCollection, conflictsCollection]);
+    sharedDb.collections = ko.observableArray([sampleItemsCollection]);
+    giganticDb.collections = ko.observableArray([schemaCollection]);
+
+    useDatabases.setState({
+      databases: [standardDb, sharedDb, giganticDb],
+      updateDatabase: jest.fn(),
+    });
+    useSelectedNode.setState({ selectedNode: undefined });
+
+    refreshActiveTab = jest.fn();
+  });
+
+  describe("using NoSQL API on Hosted Platform", () => {
+    let nodes: TreeNode[];
+    beforeEach(() => {
+      updateConfigContext({
+        platform: Platform.Hosted,
+      });
+      updateUserContext({
+        databaseAccount: {
+          properties: {
+            capabilities: [],
+          },
+        } as never,
+      });
+      nodes = createDatabaseTreeNodes(explorer, false, useDatabases.getState().databases, refreshActiveTab);
+    });
+
+    it("creates expected tree", () => {
+      expect(nodes).toMatchSnapshot();
+    });
+  });
+
+  it.each<[string, Platform, boolean, Partial<DataModels.DatabaseAccountExtendedProperties>]>([
+    ["the SQL API, on Fabric", Platform.Fabric, false, { capabilities: [], enableMultipleWriteLocations: true }],
+    ["the SQL API, on Portal", Platform.Portal, false, { capabilities: [], enableMultipleWriteLocations: true }],
+    [
+      "the Cassandra API, serverless, on Hosted",
+      Platform.Hosted,
+      false,
+      {
+        capabilities: [
+          { name: CapabilityNames.EnableCassandra, description: "" },
+          { name: CapabilityNames.EnableServerless, description: "" },
+        ],
+      },
+    ],
+    [
+      "the Mongo API, with Notebooks and Phoenix features, on Emulator",
+      Platform.Emulator,
+      true,
+      {
+        capabilities: [{ name: CapabilityNames.EnableMongo, description: "" }],
+      },
+    ],
+  ])("generates the correct tree structure for %s", (_, platform, isNotebookEnabled, dbAccountProperties) => {
+    useNotebook.setState({ isPhoenixFeatures: isNotebookEnabled });
+    updateConfigContext({ platform });
+    updateUserContext({
+      databaseAccount: {
+        properties: {
+          enableMultipleWriteLocations: true,
+          ...dbAccountProperties,
+        },
+      } as unknown as DataModels.DatabaseAccount,
+    });
+    const nodes = createDatabaseTreeNodes(
+      explorer,
+      isNotebookEnabled,
+      useDatabases.getState().databases,
+      refreshActiveTab,
+    );
+    expect(nodes).toMatchSnapshot();
+  });
+
+  // The above tests focused on the tree structure. The below tests focus on some core behaviors of the nodes.
+  // They are not exhaustive, because exhaustive tests here require a lot of mocking and can become very brittle.
+  // The goal is to cover some key behaviors like loading child nodes, opening tabs/side panels, etc.
+
+  it("adds new collections to database as they appear", () => {
+    const nodes = createDatabaseTreeNodes(explorer, false, useDatabases.getState().databases, refreshActiveTab);
+    const giganticDbNode = nodes.find((node) => node.label === giganticDb.id());
+    expect(giganticDbNode).toBeDefined();
+    expect(giganticDbNode.children.map((node) => node.label)).toStrictEqual(["schemaCollection", "load more"]);
+
+    giganticDb.collections.push({
+      ...baseCollection,
+      id: ko.observable("addedCollection"),
+    });
+
+    expect(giganticDbNode.children.map((node) => node.label)).toStrictEqual([
+      "schemaCollection",
+      "addedCollection",
+      "load more",
+    ]);
+  });
+
+  describe("the database node", () => {
+    let nodes: TreeNode[];
+    let standardDbNode: TreeNode;
+    let sharedDbNode: TreeNode;
+    let giganticDbNode: TreeNode;
+
+    beforeEach(() => {
+      updateConfigContext({ platform: Platform.Hosted });
+      updateUserContext({
+        databaseAccount: {
+          properties: {
+            capabilities: [],
+          },
+        } as unknown as DataModels.DatabaseAccount,
+      });
+      nodes = createDatabaseTreeNodes(explorer, false, useDatabases.getState().databases, refreshActiveTab);
+      standardDbNode = nodes.find((node) => node.label === standardDb.id());
+      sharedDbNode = nodes.find((node) => node.label === sharedDb.id());
+      giganticDbNode = nodes.find((node) => node.label === giganticDb.id());
+    });
+
+    it("loads child nodes when expanded", async () => {
+      // Temporarily clear the child nodes to trigger the loading behavior
+      standardDbNode.children = [];
+
+      const expanding = new PromiseSource();
+      let expandCalled = false;
+      standardDb.expandDatabase = () => {
+        expandCalled = true;
+        return expanding.promise;
+      };
+
+      standardDbNode.onExpanded();
+      expect(useSelectedNode.getState().selectedNode).toBe(standardDb);
+      expect(standardDbNode.isLoading).toStrictEqual(true);
+      expect(expandCalled).toStrictEqual(true);
+
+      await expanding.resolveAndWait();
+
+      expect(standardDbNode.isLoading).toStrictEqual(false);
+      expect(useCommandBar.getState().contextButtons).toStrictEqual([]);
+      expect(refreshActiveTab).toHaveBeenCalled();
+      expect(useDatabases.getState().updateDatabase).toHaveBeenCalledWith(standardDb);
+    });
+
+    it("opens a New Container panel when 'New Container' option in context menu is clicked", () => {
+      const newContainerMenuItem = standardDbNode.contextMenu.find((item) => item.label === "New Container");
+      newContainerMenuItem.onClick();
+      expect(explorer.onNewCollectionClicked).toHaveBeenCalled();
+    });
+
+    it("opens a Delete Database panel when 'Delete Database' option in context menu is clicked", () => {
+      const deleteDatabaseMenuItem = standardDbNode.contextMenu.find((item) => item.label === "Delete Database");
+      deleteDatabaseMenuItem.onClick();
+      expect(useSidePanel.getState().headerText).toStrictEqual("Delete Database");
+      expect(useSidePanel.getState().panelContent.type).toStrictEqual(DeleteDatabaseConfirmationPanel);
+    });
+
+    describe("the Scale subnode", () => {
+      let scaleNode: TreeNode;
+      beforeEach(() => {
+        scaleNode = sharedDbNode.children.find((node) => node.label === "Scale");
+      });
+
+      it("is selected when Scale tab is open", () => {
+        expect(scaleNode.isSelected()).toStrictEqual(false);
+        selectDataNode(sharedDb, ViewModels.CollectionTabKind.DatabaseSettingsV2);
+        expect(scaleNode.isSelected()).toStrictEqual(true);
+      });
+
+      it("opens settings tab when clicked", () => {
+        expect(sharedDb.onSettingsClick).not.toHaveBeenCalled();
+        scaleNode.onClick();
+        expect(sharedDb.onSettingsClick).toHaveBeenCalled();
+      });
+    });
+
+    describe("the load more node", () => {
+      it("loads more collections when clicked", async () => {
+        const loadCollections = new PromiseSource();
+        let loadCalled = false;
+        giganticDb.loadCollections = () => {
+          loadCalled = true;
+          return loadCollections.promise;
+        };
+
+        const loadMoreNode = giganticDbNode.children.find((node) => node.label === "load more");
+        loadMoreNode.onClick();
+        expect(loadCalled).toStrictEqual(true);
+        await loadCollections.resolveAndWait();
+        expect(useDatabases.getState().updateDatabase).toHaveBeenCalledWith(giganticDb);
+      });
+    });
+
+    describe("the Collection subnode", () => {
+      let standardCollectionNode: TreeNode;
+      beforeEach(() => {
+        standardCollectionNode = standardDbNode.children.find((node) => node.label === standardCollection.id());
+      });
+
+      it.each([
+        [
+          "for SQL API",
+          () => updateUserContext({ databaseAccount: { properties: {} } as unknown as DataModels.DatabaseAccount }),
+        ],
+        [
+          "for Gremlin API",
+          () =>
+            updateUserContext({
+              databaseAccount: {
+                properties: { capabilities: [{ name: CapabilityNames.EnableGremlin, description: "" }] },
+              } as unknown as DataModels.DatabaseAccount,
+            }),
+        ],
+      ])("loads sprocs/udfs/triggers when expanded, %s", async () => {
+        standardCollection.loadStoredProcedures = jest.fn(() => Promise.resolve());
+        standardCollection.loadUserDefinedFunctions = jest.fn(() => Promise.resolve());
+        standardCollection.loadTriggers = jest.fn(() => Promise.resolve());
+
+        await standardCollectionNode.onExpanded();
+
+        expect(standardCollection.loadStoredProcedures).toHaveBeenCalled();
+        expect(standardCollection.loadUserDefinedFunctions).toHaveBeenCalled();
+        expect(standardCollection.loadTriggers).toHaveBeenCalled();
+      });
+
+      it.each([
+        ["in Fabric", () => updateConfigContext({ platform: Platform.Fabric })],
+        [
+          "for Cassandra API",
+          () =>
+            updateUserContext({
+              databaseAccount: {
+                properties: { capabilities: [{ name: CapabilityNames.EnableCassandra, description: "" }] },
+              } as unknown as DataModels.DatabaseAccount,
+            }),
+        ],
+        [
+          "for Mongo API",
+          () =>
+            updateUserContext({
+              databaseAccount: {
+                properties: { capabilities: [{ name: CapabilityNames.EnableMongo, description: "" }] },
+              } as unknown as DataModels.DatabaseAccount,
+            }),
+        ],
+        [
+          "for Tables API",
+          () =>
+            updateUserContext({
+              databaseAccount: {
+                properties: { capabilities: [{ name: CapabilityNames.EnableTable, description: "" }] },
+              } as unknown as DataModels.DatabaseAccount,
+            }),
+        ],
+      ])("does not load sprocs/udfs/triggers when expanded, %s", async (_, setup) => {
+        setup();
+
+        // Rebuild the nodes after changing the user/config context.
+        nodes = createDatabaseTreeNodes(explorer, false, useDatabases.getState().databases, refreshActiveTab);
+        standardDbNode = nodes.find((node) => node.label === standardDb.id());
+        standardCollectionNode = standardDbNode.children.find((node) => node.label === standardCollection.id());
+
+        standardCollection.loadStoredProcedures = jest.fn(() => Promise.resolve());
+        standardCollection.loadUserDefinedFunctions = jest.fn(() => Promise.resolve());
+        standardCollection.loadTriggers = jest.fn(() => Promise.resolve());
+
+        await standardCollectionNode.onExpanded();
+
+        expect(standardCollection.loadStoredProcedures).not.toHaveBeenCalled();
+        expect(standardCollection.loadUserDefinedFunctions).not.toHaveBeenCalled();
+        expect(standardCollection.loadTriggers).not.toHaveBeenCalled();
+      });
+    });
+  });
+});

src/Explorer/Tree/treeNodeUtil.ts

@@ -1,4 +1,4 @@
-import { TreeNode2 } from "Explorer/Controls/TreeComponent2/TreeNode2Component";
+import { TreeNode } from "Explorer/Controls/TreeComponent/TreeNodeComponent";
 import TabsBase from "Explorer/Tabs/TabsBase";
 import StoredProcedure from "Explorer/Tree/StoredProcedure";
 import Trigger from "Explorer/Tree/Trigger";
@@ -6,8 +6,11 @@ import UserDefinedFunction from "Explorer/Tree/UserDefinedFunction";
 import { useDatabases } from "Explorer/useDatabases";
 import { getItemName } from "Utils/APITypeUtils";
 import { isServerlessAccount } from "Utils/CapabilityUtils";
+import { useTabs } from "hooks/useTabs";
+import CosmosDBIcon from "../../../images/Azure-Cosmos-DB.svg";
 import CollectionIcon from "../../../images/tree-collection.svg";
 import { isPublicInternetAccessAllowed } from "../../Common/DatabaseAccountUtility";
+import { Platform, configContext } from "../../ConfigContext";
 import * as DataModels from "../../Contracts/DataModels";
 import * as ViewModels from "../../Contracts/ViewModels";
 import { userContext } from "../../UserContext";
@@ -17,7 +20,192 @@ import { useCommandBar } from "../Menus/CommandBar/CommandBarComponentAdapter";
 import { mostRecentActivity } from "../MostRecentActivity/MostRecentActivity";
 import { useNotebook } from "../Notebook/useNotebook";
 import { useSelectedNode } from "../useSelectedNode";
-import { Platform, configContext } from "./../../ConfigContext";
+
+export const shouldShowScriptNodes = (): boolean => {
+  return (
+    configContext.platform !== Platform.Fabric && (userContext.apiType === "SQL" || userContext.apiType === "Gremlin")
+  );
+};
+
+export const createSampleDataTreeNodes = (sampleDataResourceTokenCollection: ViewModels.CollectionBase): TreeNode[] => {
+  const updatedSampleTree: TreeNode = {
+    label: sampleDataResourceTokenCollection.databaseId,
+    isExpanded: false,
+    iconSrc: CosmosDBIcon,
+    className: "databaseHeader",
+    children: [
+      {
+        label: sampleDataResourceTokenCollection.id(),
+        iconSrc: CollectionIcon,
+        isExpanded: false,
+        className: "collectionHeader",
+        contextMenu: ResourceTreeContextMenuButtonFactory.createSampleCollectionContextMenuButton(),
+        onClick: () => {
+          useSelectedNode.getState().setSelectedNode(sampleDataResourceTokenCollection);
+          useCommandBar.getState().setContextButtons([]);
+          useTabs
+            .getState()
+            .refreshActiveTab(
+              (tab: TabsBase) =>
+                tab.collection?.id() === sampleDataResourceTokenCollection.id() &&
+                tab.collection.databaseId === sampleDataResourceTokenCollection.databaseId,
+            );
+        },
+        isSelected: () =>
+          useSelectedNode
+            .getState()
+            .isDataNodeSelected(sampleDataResourceTokenCollection.databaseId, sampleDataResourceTokenCollection.id()),
+        onContextMenuOpen: () => useSelectedNode.getState().setSelectedNode(sampleDataResourceTokenCollection),
+        children: [
+          {
+            label: "Items",
+            onClick: () => sampleDataResourceTokenCollection.onDocumentDBDocumentsClick(),
+            contextMenu: ResourceTreeContextMenuButtonFactory.createSampleCollectionContextMenuButton(),
+            isSelected: () =>
+              useSelectedNode
+                .getState()
+                .isDataNodeSelected(
+                  sampleDataResourceTokenCollection.databaseId,
+                  sampleDataResourceTokenCollection.id(),
+                  [ViewModels.CollectionTabKind.Documents],
+                ),
+          },
+        ],
+      },
+    ],
+  };
+
+  return [updatedSampleTree];
+};
+
+export const createResourceTokenTreeNodes = (collection: ViewModels.CollectionBase): TreeNode[] => {
+  if (!collection) {
+    return [
+      {
+        label: "",
+        isExpanded: true,
+        children: [],
+      },
+    ];
+  }
+
+  const children: TreeNode[] = [];
+  children.push({
+    label: "Items",
+    onClick: () => {
+      collection.onDocumentDBDocumentsClick();
+      // push to most recent
+      mostRecentActivity.collectionWasOpened(userContext.databaseAccount?.id, collection);
+    },
+    isSelected: () =>
+      useSelectedNode
+        .getState()
+        .isDataNodeSelected(collection.databaseId, collection.id(), [ViewModels.CollectionTabKind.Documents]),
+  });
+
+  const collectionNode: TreeNode = {
+    label: collection.id(),
+    iconSrc: CollectionIcon,
+    isExpanded: true,
+    children,
+    className: "collectionHeader",
+    onClick: () => {
+      // Rewritten version of expandCollapseCollection
+      useSelectedNode.getState().setSelectedNode(collection);
+      useCommandBar.getState().setContextButtons([]);
+      useTabs
+        .getState()
+        .refreshActiveTab(
+          (tab) => tab.collection?.id() === collection.id() && tab.collection.databaseId === collection.databaseId,
+        );
+    },
+    isSelected: () => useSelectedNode.getState().isDataNodeSelected(collection.databaseId, collection.id()),
+  };
+
+  return [collectionNode];
+};
+
+export const createDatabaseTreeNodes = (
+  container: Explorer,
+  isNotebookEnabled: boolean,
+  databases: ViewModels.Database[],
+  refreshActiveTab: (comparator: (tab: TabsBase) => boolean) => void,
+): TreeNode[] => {
+  const databaseTreeNodes: TreeNode[] = databases.map((database: ViewModels.Database) => {
+    const buildDatabaseChildNodes = (databaseNode: TreeNode) => {
+      databaseNode.children = [];
+      if (database.isDatabaseShared() && configContext.platform !== Platform.Fabric) {
+        databaseNode.children.push({
+          id: database.isSampleDB ? "sampleScaleSettings" : "",
+          label: "Scale",
+          isSelected: () =>
+            useSelectedNode
+              .getState()
+              .isDataNodeSelected(database.id(), undefined, [ViewModels.CollectionTabKind.DatabaseSettingsV2]),
+          onClick: database.onSettingsClick.bind(database),
+        });
+      }
+
+      // Find collections
+      database
+        .collections()
+        .forEach((collection: ViewModels.Collection) =>
+          databaseNode.children.push(
+            buildCollectionNode(database, collection, isNotebookEnabled, container, refreshActiveTab),
+          ),
+        );
+
+      if (database.collectionsContinuationToken) {
+        const loadMoreNode: TreeNode = {
+          label: "load more",
+          className: "loadMoreHeader",
+          onClick: async () => {
+            await database.loadCollections();
+            useDatabases.getState().updateDatabase(database);
+          },
+        };
+        databaseNode.children.push(loadMoreNode);
+      }
+    };
+
+    const databaseNode: TreeNode = {
+      label: database.id(),
+      iconSrc: CosmosDBIcon,
+      className: "databaseHeader",
+      children: [],
+      isSelected: () => useSelectedNode.getState().isDataNodeSelected(database.id()),
+      contextMenu: ResourceTreeContextMenuButtonFactory.createDatabaseContextMenu(container, database.id()),
+      onExpanded: async () => {
+        useSelectedNode.getState().setSelectedNode(database);
+        if (!databaseNode.children || databaseNode.children?.length === 0) {
+          databaseNode.isLoading = true;
+        }
+        await database.expandDatabase();
+        databaseNode.isLoading = false;
+        useCommandBar.getState().setContextButtons([]);
+        refreshActiveTab((tab: TabsBase) => tab.collection?.databaseId === database.id());
+        useDatabases.getState().updateDatabase(database);
+      },
+      onContextMenuOpen: () => useSelectedNode.getState().setSelectedNode(database),
+      isExpanded: database.isDatabaseExpanded(),
+      onCollapsed: () => {
+        database.collapseDatabase();
+        // useCommandBar.getState().setContextButtons([]);
+        useDatabases.getState().updateDatabase(database);
+      },
+    };
+
+    buildDatabaseChildNodes(databaseNode);
+
+    database.collections.subscribe(() => {
+      buildDatabaseChildNodes(databaseNode);
+    });
+
+    return databaseNode;
+  });
+
+  return databaseTreeNodes;
+};
 
 export const buildCollectionNode = (
   database: ViewModels.Database,
@@ -25,15 +213,14 @@ export const buildCollectionNode = (
   isNotebookEnabled: boolean,
   container: Explorer,
   refreshActiveTab: (comparator: (tab: TabsBase) => boolean) => void,
-): TreeNode2 => {
+): TreeNode => {
-  let children: TreeNode2[];
+  let children: TreeNode[];
-
   // Flat Tree for Fabric
   if (configContext.platform !== Platform.Fabric) {
     children = buildCollectionNodeChildren(database, collection, isNotebookEnabled, container, refreshActiveTab);
   }
 
-  return {
+  const collectionNode: TreeNode = {
     label: collection.id(),
     iconSrc: CollectionIcon,
     children: children,
@@ -54,6 +241,15 @@ export const buildCollectionNode = (
           tab.collection?.id() === collection.id() && tab.collection.databaseId === collection.databaseId,
       );
       useDatabases.getState().updateDatabase(database);
+
+      // If we're showing script nodes, start loading them.
+      if (shouldShowScriptNodes()) {
+        await collection.loadStoredProcedures();
+        await collection.loadUserDefinedFunctions();
+        await collection.loadTriggers();
+      }
+
+      useDatabases.getState().updateDatabase(database);
     },
     isSelected: () => useSelectedNode.getState().isDataNodeSelected(collection.databaseId, collection.id()),
     onContextMenuOpen: () => useSelectedNode.getState().setSelectedNode(collection),
@@ -64,6 +260,8 @@ export const buildCollectionNode = (
     },
     isExpanded: collection.isCollectionExpanded(),
   };
+
+  return collectionNode;
 };
 
 const buildCollectionNodeChildren = (
@@ -72,9 +270,8 @@ const buildCollectionNodeChildren = (
   isNotebookEnabled: boolean,
   container: Explorer,
   refreshActiveTab: (comparator: (tab: TabsBase) => boolean) => void,
-): TreeNode2[] => {
+): TreeNode[] => {
-  const showScriptNodes = userContext.apiType === "SQL" || userContext.apiType === "Gremlin";
+  const children: TreeNode[] = [];
-  const children: TreeNode2[] = [];
   children.push({
     label: getItemName(),
     id: collection.isSampleCollection ? "sampleItems" : "",
@@ -128,14 +325,14 @@ const buildCollectionNodeChildren = (
     });
   }
 
-  const schemaNode: TreeNode2 = buildSchemaNode(collection, container, refreshActiveTab);
+  const schemaNode: TreeNode = buildSchemaNode(collection, container, refreshActiveTab);
   if (schemaNode) {
     children.push(schemaNode);
   }
 
   const onUpdateDatabase = () => useDatabases.getState().updateDatabase(database);
 
-  if (showScriptNodes) {
+  if (shouldShowScriptNodes()) {
     children.push(buildStoredProcedureNode(collection, container, refreshActiveTab, onUpdateDatabase));
     children.push(buildUserDefinedFunctionsNode(collection, container, refreshActiveTab, onUpdateDatabase));
     children.push(buildTriggerNode(collection, container, refreshActiveTab, onUpdateDatabase));
@@ -166,7 +363,7 @@ const buildStoredProcedureNode = (
   container: Explorer,
   refreshActiveTab: (comparator: (tab: TabsBase) => boolean) => void,
   onUpdateDatabase: () => void,
-): TreeNode2 => {
+): TreeNode => {
   return {
     label: "Stored Procedures",
     children: collection.storedProcedures().map((sp: StoredProcedure) => ({
@@ -195,7 +392,7 @@ const buildUserDefinedFunctionsNode = (
   container: Explorer,
   refreshActiveTab: (comparator: (tab: TabsBase) => boolean) => void,
   onUpdateDatabase: () => void,
-): TreeNode2 => {
+): TreeNode => {
   return {
     label: "User Defined Functions",
     children: collection.userDefinedFunctions().map((udf: UserDefinedFunction) => ({
@@ -226,7 +423,7 @@ const buildTriggerNode = (
   container: Explorer,
   refreshActiveTab: (comparator: (tab: TabsBase) => boolean) => void,
   onUpdateDatabase: () => void,
-): TreeNode2 => {
+): TreeNode => {
   return {
     label: "Triggers",
     children: collection.triggers().map((trigger: Trigger) => ({
@@ -254,7 +451,7 @@ const buildSchemaNode = (
   collection: ViewModels.Collection,
   container: Explorer,
   refreshActiveTab: (comparator: (tab: TabsBase) => boolean) => void,
-): TreeNode2 => {
+): TreeNode => {
   if (collection.analyticalStorageTtl() === undefined) {
     return undefined;
   }
@@ -273,7 +470,7 @@ const buildSchemaNode = (
   };
 };
 
-const getSchemaNodes = (fields: DataModels.IDataField[]): TreeNode2[] => {
+const getSchemaNodes = (fields: DataModels.IDataField[]): TreeNode[] => {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   const schema: any = {};
 
@@ -307,8 +504,8 @@ const getSchemaNodes = (fields: DataModels.IDataField[]): TreeNode2[] => {
   });
 
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  const traverse = (obj: any): TreeNode2[] => {
+  const traverse = (obj: any): TreeNode[] => {
-    const children: TreeNode2[] = [];
+    const children: TreeNode[] = [];
 
     if (obj !== undefined && !Array.isArray(obj) && typeof obj === "object") {
       Object.entries(obj).forEach(([key, value]) => {

src/Explorer/Tree2/ResourceTree.tsx

@@ -1,107 +0,0 @@
-import {
-  FluentProvider,
-  Tree,
-  TreeItemValue,
-  TreeOpenChangeData,
-  TreeOpenChangeEvent,
-} from "@fluentui/react-components";
-import { configContext } from "ConfigContext";
-import { TreeNode2, TreeNode2Component } from "Explorer/Controls/TreeComponent2/TreeNode2Component";
-import { getPlatformTheme } from "Explorer/Theme/ThemeUtil";
-import { useDatabaseTreeNodes } from "Explorer/Tree2/useDatabaseTreeNodes";
-import * as React from "react";
-import shallow from "zustand/shallow";
-import Explorer from "../Explorer";
-import { useNotebook } from "../Notebook/useNotebook";
-
-export const MyNotebooksTitle = "My Notebooks";
-export const GitHubReposTitle = "GitHub repos";
-
-interface ResourceTreeProps {
-  container: Explorer;
-}
-
-export const DATA_TREE_LABEL = "DATA";
-
-/**
- * Top-level tree that has no label, but contains all subtrees
- */
-export const ResourceTree2: React.FC<ResourceTreeProps> = ({ container }: ResourceTreeProps): JSX.Element => {
-  const {
-    isNotebookEnabled,
-    // myNotebooksContentRoot,
-    // galleryContentRoot,
-    // gitHubNotebooksContentRoot,
-    // updateNotebookItem,
-  } = useNotebook(
-    (state) => ({
-      isNotebookEnabled: state.isNotebookEnabled,
-      myNotebooksContentRoot: state.myNotebooksContentRoot,
-      galleryContentRoot: state.galleryContentRoot,
-      gitHubNotebooksContentRoot: state.gitHubNotebooksContentRoot,
-      updateNotebookItem: state.updateNotebookItem,
-    }),
-    shallow,
-  );
-  // const { activeTab } = useTabs();
-  const databaseTreeNodes = useDatabaseTreeNodes(container, isNotebookEnabled);
-  const [openItems, setOpenItems] = React.useState<Iterable<TreeItemValue>>([DATA_TREE_LABEL]);
-
-  const dataNodeTree: TreeNode2 = {
-    id: "data",
-    label: DATA_TREE_LABEL,
-    className: "accordionItemHeader",
-    children: databaseTreeNodes,
-    isScrollable: true,
-  };
-
-  React.useEffect(() => {
-    // Compute open items based on node.isExpanded
-    const updateOpenItems = (node: TreeNode2, parentNodeId: string): void => {
-      // This will look for ANY expanded node, event if its parent node isn't expanded
-      // and add it to the openItems list
-      const globalId = parentNodeId === undefined ? node.label : `${parentNodeId}/${node.label}`;
-
-      if (node.isExpanded) {
-        let found = false;
-        for (const id of openItems) {
-          if (id === globalId) {
-            found = true;
-            break;
-          }
-        }
-        if (!found) {
-          setOpenItems((prevOpenItems) => [...prevOpenItems, globalId]);
-        }
-      }
-
-      if (node.children) {
-        for (const child of node.children) {
-          updateOpenItems(child, globalId);
-        }
-      }
-    };
-
-    updateOpenItems(dataNodeTree, undefined);
-  }, [databaseTreeNodes]);
-
-  const handleOpenChange = (event: TreeOpenChangeEvent, data: TreeOpenChangeData) => setOpenItems(data.openItems);
-
-  return (
-    <>
-      <FluentProvider theme={getPlatformTheme(configContext.platform)} style={{ overflow: "hidden" }}>
-        <Tree
-          aria-label="CosmosDB resources"
-          openItems={openItems}
-          onOpenChange={handleOpenChange}
-          size="medium"
-          style={{ height: "100%", width: "290px" }}
-        >
-          {[dataNodeTree].map((node) => (
-            <TreeNode2Component key={node.label} className="dataResourceTree" node={node} treeNodeId={node.label} />
-          ))}
-        </Tree>
-      </FluentProvider>
-    </>
-  );
-};

src/Explorer/Tree2/useDatabaseTreeNodes.ts

@@ -1,90 +0,0 @@
-import { TreeNode2 } from "Explorer/Controls/TreeComponent2/TreeNode2Component";
-import TabsBase from "Explorer/Tabs/TabsBase";
-import { buildCollectionNode } from "Explorer/Tree2/containerTreeNodeUtil";
-import { useDatabases } from "Explorer/useDatabases";
-import { useTabs } from "hooks/useTabs";
-import CosmosDBIcon from "../../../images/Azure-Cosmos-DB.svg";
-import * as ViewModels from "../../Contracts/ViewModels";
-import Explorer from "../Explorer";
-import { useCommandBar } from "../Menus/CommandBar/CommandBarComponentAdapter";
-import { useSelectedNode } from "../useSelectedNode";
-import { Platform, configContext } from "./../../ConfigContext";
-
-export const useDatabaseTreeNodes = (container: Explorer, isNotebookEnabled: boolean): TreeNode2[] => {
-  const databases = useDatabases((state) => state.databases);
-  const { refreshActiveTab } = useTabs();
-
-  const databaseTreeNodes: TreeNode2[] = databases.map((database: ViewModels.Database) => {
-    const databaseNode: TreeNode2 = {
-      label: database.id(),
-      iconSrc: CosmosDBIcon,
-      className: "databaseHeader",
-      children: [],
-      isSelected: () => useSelectedNode.getState().isDataNodeSelected(database.id()),
-      contextMenu: undefined, // TODO Disable this for now as the actions don't work. ResourceTreeContextMenuButtonFactory.createDatabaseContextMenu(container, database.id()),
-      onExpanded: async () => {
-        useSelectedNode.getState().setSelectedNode(database);
-        if (!databaseNode.children || databaseNode.children?.length === 0) {
-          databaseNode.isLoading = true;
-        }
-        await database.expandDatabase();
-        databaseNode.isLoading = false;
-        useCommandBar.getState().setContextButtons([]);
-        refreshActiveTab((tab: TabsBase) => tab.collection?.databaseId === database.id());
-        useDatabases.getState().updateDatabase(database);
-      },
-      onContextMenuOpen: () => useSelectedNode.getState().setSelectedNode(database),
-      isExpanded: database.isDatabaseExpanded(),
-      onCollapsed: () => {
-        database.collapseDatabase();
-        // useCommandBar.getState().setContextButtons([]);
-        useDatabases.getState().updateDatabase(database);
-      },
-    };
-
-    if (database.isDatabaseShared() && configContext.platform !== Platform.Fabric) {
-      databaseNode.children.push({
-        id: database.isSampleDB ? "sampleScaleSettings" : "",
-        label: "Scale",
-        isSelected: () =>
-          useSelectedNode
-            .getState()
-            .isDataNodeSelected(database.id(), undefined, [ViewModels.CollectionTabKind.DatabaseSettingsV2]),
-        onClick: database.onSettingsClick.bind(database),
-      });
-    }
-
-    // Find collections
-    database
-      .collections()
-      .forEach((collection: ViewModels.Collection) =>
-        databaseNode.children.push(
-          buildCollectionNode(database, collection, isNotebookEnabled, container, refreshActiveTab),
-        ),
-      );
-
-    if (database.collectionsContinuationToken) {
-      const loadMoreNode: TreeNode2 = {
-        label: "load more",
-        className: "loadMoreHeader",
-        onClick: async () => {
-          await database.loadCollections();
-          useDatabases.getState().updateDatabase(database);
-        },
-      };
-      databaseNode.children.push(loadMoreNode);
-    }
-
-    database.collections.subscribe((collections: ViewModels.Collection[]) => {
-      collections.forEach((collection: ViewModels.Collection) =>
-        databaseNode.children.push(
-          buildCollectionNode(database, collection, isNotebookEnabled, container, refreshActiveTab),
-        ),
-      );
-    });
-
-    return databaseNode;
-  });
-
-  return databaseTreeNodes;
-};

src/UserContext.ts

@@ -69,7 +69,7 @@ export type AdminFeedbackPolicySettings = {
   [key in AdminFeedbackControlPolicy]: boolean;
 };
 
-interface UserContext {
+export interface UserContext {
   readonly fabricContext?: FabricContext;
   readonly authType?: AuthType;
   readonly masterKey?: string;

src/Utils/PromiseSource.ts

@@ -0,0 +1,18 @@
+export default class PromiseSource<T = void> {
+  promise: Promise<T>;
+  resolve: (value: T) => void;
+  reject: (reason?: unknown) => void;
+
+  constructor() {
+    this.promise = new Promise((resolve, reject) => {
+      this.resolve = resolve;
+      this.reject = reject;
+    });
+  }
+
+  /** Resolves the promise, then gets off the thread and waits until the currently-registered 'then' callback run. */
+  async resolveAndWait(value: T): Promise<T> {
+    this.resolve(value);
+    return await this.promise;
+  }
+}

test/cassandra/container.spec.ts

@@ -1,6 +1,14 @@
 import { jest } from "@jest/globals";
 import "expect-playwright";
-import { generateUniqueName, getAzureCLICredentialsToken } from "../utils/shared";
+import {
+  AccountType,
+  generateUniqueName,
+  getPanelSelector,
+  getTestExplorerUrl,
+  getTreeMenuItemSelector,
+  getTreeNodeSelector,
+  openContextMenu,
+} from "../utils/shared";
 import { waitForExplorer } from "../utils/waitForExplorer";
 jest.setTimeout(120000);
 
@@ -8,30 +16,42 @@ test("Cassandra keyspace and table CRUD", async () => {
   const keyspaceId = generateUniqueName("keyspace");
   const tableId = generateUniqueName("table");
 
-  // We can't retrieve AZ CLI credentials from the browser so we get them here.
-  const token = await getAzureCLICredentialsToken();
   page.setDefaultTimeout(50000);
 
-  await page.goto(`https://localhost:1234/testExplorer.html?accountName=portal-cassandra-runner&token=${token}`);
+  const url = await getTestExplorerUrl(AccountType.Cassandra);
+  await page.goto(url);
   await page.waitForSelector("iframe");
   const explorer = await waitForExplorer();
 
   await explorer.click('[data-test="New Table"]');
+
+  await explorer.waitForSelector(getPanelSelector("Add Table"));
   await explorer.click('[aria-label="Keyspace id"]');
   await explorer.fill('[aria-label="Keyspace id"]', keyspaceId);
   await explorer.click('[aria-label="addCollection-table Id Create table"]');
   await explorer.fill('[aria-label="addCollection-table Id Create table"]', tableId);
+  await explorer.fill('[aria-label="Table max RU/s"]', "1000");
   await explorer.click("#sidePanelOkButton");
-  await explorer.click(`.nodeItem >> text=${keyspaceId}`);
+  await explorer.waitForSelector(getPanelSelector("Add Table"), { state: "detached" });
-  await explorer.click(`[data-test="${tableId}"] [aria-label="More options"]`);
+
-  await explorer.click('button[role="menuitem"]:has-text("Delete Table")');
+  await explorer.click(getTreeNodeSelector(`DATA/${keyspaceId}`));
+  await openContextMenu(explorer, `DATA/${keyspaceId}/${tableId}`);
+  await explorer.click(getTreeMenuItemSelector(`DATA/${keyspaceId}/${tableId}`, "Delete Table"));
+
+  await explorer.waitForSelector(getPanelSelector("Delete Table"));
   await explorer.fill('text=* Confirm by typing the table id >> input[type="text"]', tableId);
   await explorer.click('[aria-label="OK"]');
-  await explorer.click(`[data-test="${keyspaceId}"] [aria-label="More options"]`);
+  await explorer.waitForSelector(getPanelSelector("Delete Table"), { state: "detached" });
-  await explorer.click('button[role="menuitem"]:has-text("Delete Keyspace")');
+
+  await openContextMenu(explorer, `DATA/${keyspaceId}`);
+  await explorer.click(getTreeMenuItemSelector(`DATA/${keyspaceId}`, "Delete Keyspace"));
+
+  await explorer.waitForSelector(getPanelSelector("Delete Keyspace"));
   await explorer.click('text=* Confirm by typing the database id >> input[type="text"]');
   await explorer.fill('text=* Confirm by typing the database id >> input[type="text"]', keyspaceId);
   await explorer.click("#sidePanelOkButton");
+  await explorer.waitForSelector(getPanelSelector("Delete Keyspace"), { state: "detached" });
+
   await expect(explorer).not.toHaveText(".dataResourceTree", keyspaceId);
   await expect(explorer).not.toHaveText(".dataResourceTree", tableId);
 });

test/graph/container.spec.ts

@@ -1,6 +1,14 @@
 import { jest } from "@jest/globals";
 import "expect-playwright";
-import { generateDatabaseNameWithTimestamp, generateUniqueName, getAzureCLICredentialsToken } from "../utils/shared";
+import {
+  generateDatabaseNameWithTimestamp,
+  generateUniqueName,
+  getAzureCLICredentialsToken,
+  getPanelSelector,
+  getTreeMenuItemSelector,
+  getTreeNodeSelector,
+  openContextMenu,
+} from "../utils/shared";
 import { waitForExplorer } from "../utils/waitForExplorer";
 jest.setTimeout(240000);
 
@@ -17,22 +25,34 @@ test("Graph CRUD", async () => {
 
   // Create new database and graph
   await explorer.click('[data-test="New Graph"]');
+
+  await explorer.waitForSelector(getPanelSelector("New Graph"));
   await explorer.fill('[aria-label="New database id, Type a new database id"]', databaseId);
   await explorer.fill('[aria-label="Graph id, Example Graph1"]', containerId);
   await explorer.fill('[aria-label="Partition key"]', "/pk");
   await explorer.click("#sidePanelOkButton");
-  await explorer.click(`.nodeItem >> text=${databaseId}`);
+  await explorer.waitForSelector(getPanelSelector("New Graph"), { state: "detached" });
-  await explorer.click(`.nodeItem >> text=${containerId}`);
+
   // Delete database and graph
-  await explorer.click(`[data-test="${containerId}"] [aria-label="More options"]`);
+
-  await explorer.click('button[role="menuitem"]:has-text("Delete Graph")');
+  await explorer.click(getTreeNodeSelector(`DATA/${databaseId}`));
+  await openContextMenu(explorer, `DATA/${databaseId}/${containerId}`);
+  await explorer.click(getTreeMenuItemSelector(`DATA/${databaseId}/${containerId}`, "Delete Graph"));
+
+  await explorer.waitForSelector(getPanelSelector("Delete Graph"));
   await explorer.fill('text=* Confirm by typing the graph id >> input[type="text"]', containerId);
   await explorer.click('[aria-label="OK"]');
-  await explorer.click(`[data-test="${databaseId}"] [aria-label="More options"]`);
+  await explorer.waitForSelector(getPanelSelector("Delete Graph"), { state: "detached" });
-  await explorer.click('button[role="menuitem"]:has-text("Delete Database")');
+
+  await openContextMenu(explorer, `DATA/${databaseId}`);
+  await explorer.click(getTreeMenuItemSelector(`DATA/${databaseId}`, "Delete Database"));
+
+  await explorer.waitForSelector(getPanelSelector("Delete Database"));
   await explorer.click('text=* Confirm by typing the database id >> input[type="text"]');
   await explorer.fill('text=* Confirm by typing the database id >> input[type="text"]', databaseId);
   await explorer.click("#sidePanelOkButton");
+  await explorer.waitForSelector(getPanelSelector("Delete Database"), { state: "detached" });
+
   await expect(explorer).not.toHaveText(".dataResourceTree", databaseId);
   await expect(explorer).not.toHaveText(".dataResourceTree", containerId);
 });

test/mongo/container.spec.ts

@@ -1,6 +1,15 @@
 import { jest } from "@jest/globals";
 import "expect-playwright";
-import { generateDatabaseNameWithTimestamp, generateUniqueName, getAzureCLICredentialsToken } from "../utils/shared";
+import {
+  AccountType,
+  generateDatabaseNameWithTimestamp,
+  generateUniqueName,
+  getPanelSelector,
+  getTestExplorerUrl,
+  getTreeMenuItemSelector,
+  getTreeNodeSelector,
+  openContextMenu,
+} from "../utils/shared";
 import { waitForExplorer } from "../utils/waitForExplorer";
 jest.setTimeout(240000);
 
@@ -8,42 +17,56 @@ test("Mongo CRUD", async () => {
   const databaseId = generateDatabaseNameWithTimestamp();
   const containerId = generateUniqueName("container");
 
-  // We can't retrieve AZ CLI credentials from the browser so we get them here.
-  const token = await getAzureCLICredentialsToken();
   page.setDefaultTimeout(50000);
 
-  await page.goto(`https://localhost:1234/testExplorer.html?accountName=portal-mongo-runner&token=${token}`);
+  const url = await getTestExplorerUrl(AccountType.Mongo);
+  await page.goto(url);
   const explorer = await waitForExplorer();
 
   // Create new database and collection
   await explorer.click('[data-test="New Collection"]');
+
+  await explorer.waitForSelector(getPanelSelector("New Collection"));
   await explorer.fill('[aria-label="New database id, Type a new database id"]', databaseId);
   await explorer.fill('[aria-label="Collection id, Example Collection1"]', containerId);
   await explorer.fill('[aria-label="Shard key"]', "pk");
   await explorer.click("#sidePanelOkButton");
-  await explorer.click(`.nodeItem >> text=${databaseId}`);
+  await explorer.waitForSelector(getPanelSelector("New Collection"), { state: "detached" });
-  await explorer.click(`.nodeItem >> text=${containerId}`);
+
+  await explorer.click(getTreeNodeSelector(`DATA/${databaseId}`));
+  await explorer.click(getTreeNodeSelector(`DATA/${databaseId}/${containerId}`));
+
   // Create indexing policy
-  await explorer.click(".nodeItem >> text=Settings");
+  await explorer.click(getTreeNodeSelector(`DATA/${databaseId}/${containerId}/Settings`));
   await explorer.click('button[role="tab"]:has-text("Indexing Policy")');
   await explorer.click('[aria-label="Index Field Name 0"]');
   await explorer.fill('[aria-label="Index Field Name 0"]', "foo");
   await explorer.click("text=Select an index type");
   await explorer.click('button[role="option"]:has-text("Single Field")');
   await explorer.click('[data-test="Save"]');
+
   // Remove indexing policy
   await explorer.click('[aria-label="Delete index Button"]');
   await explorer.click('[data-test="Save"]');
+
   // Delete database and collection
-  await explorer.click(`[data-test="${containerId}"] [aria-label="More options"]`);
+  await openContextMenu(explorer, `DATA/${databaseId}/${containerId}`);
-  await explorer.click('button[role="menuitem"]:has-text("Delete Collection")');
+  await explorer.click(getTreeMenuItemSelector(`DATA/${databaseId}/${containerId}`, "Delete Collection"));
+
+  await explorer.waitForSelector(getPanelSelector("Delete Collection"));
   await explorer.fill('text=* Confirm by typing the collection id >> input[type="text"]', containerId);
   await explorer.click('[aria-label="OK"]');
-  await explorer.click(`[data-test="${databaseId}"] [aria-label="More options"]`);
+  await explorer.waitForSelector(getPanelSelector("Delete Collection"), { state: "detached" });
-  await explorer.click('button[role="menuitem"]:has-text("Delete Database")');
+
+  await openContextMenu(explorer, `DATA/${databaseId}`);
+  await explorer.click(getTreeMenuItemSelector(`DATA/${databaseId}`, "Delete Database"));
+
+  await explorer.waitForSelector(getPanelSelector("Delete Database"));
   await explorer.click('text=* Confirm by typing the database id >> input[type="text"]');
   await explorer.fill('text=* Confirm by typing the database id >> input[type="text"]', databaseId);
   await explorer.click("#sidePanelOkButton");
+  await explorer.waitForSelector(getPanelSelector("Delete Database"), { state: "detached" });
+
   await expect(explorer).not.toHaveText(".dataResourceTree", databaseId);
   await expect(explorer).not.toHaveText(".dataResourceTree", containerId);
 });

test/mongo/container32.spec.ts

@@ -1,6 +1,15 @@
 import { jest } from "@jest/globals";
 import "expect-playwright";
-import { generateDatabaseNameWithTimestamp, generateUniqueName, getAzureCLICredentialsToken } from "../utils/shared";
+import {
+  AccountType,
+  generateDatabaseNameWithTimestamp,
+  generateUniqueName,
+  getPanelSelector,
+  getTestExplorerUrl,
+  getTreeMenuItemSelector,
+  getTreeNodeSelector,
+  openContextMenu,
+} from "../utils/shared";
 import { waitForExplorer } from "../utils/waitForExplorer";
 jest.setTimeout(240000);
 
@@ -8,31 +17,43 @@ test("Mongo CRUD", async () => {
   const databaseId = generateDatabaseNameWithTimestamp();
   const containerId = generateUniqueName("container");
 
-  // We can't retrieve AZ CLI credentials from the browser so we get them here.
-  const token = await getAzureCLICredentialsToken();
   page.setDefaultTimeout(50000);
 
-  await page.goto(`https://localhost:1234/testExplorer.html?accountName=portal-mongo32-runner&token=${token}`);
+  const url = await getTestExplorerUrl(AccountType.Mongo32);
+  await page.goto(url);
   const explorer = await waitForExplorer();
 
   // Create new database and collection
   await explorer.click('[data-test="New Collection"]');
+
+  await explorer.waitForSelector(getPanelSelector("New Collection"));
   await explorer.fill('[aria-label="New database id, Type a new database id"]', databaseId);
   await explorer.fill('[aria-label="Collection id, Example Collection1"]', containerId);
   await explorer.fill('[aria-label="Shard key"]', "pk");
   await explorer.click("#sidePanelOkButton");
-  explorer.click(`.nodeItem >> text=${databaseId}`);
+  await explorer.waitForSelector(getPanelSelector("New Collection"), { state: "detached" });
-  explorer.click(`.nodeItem >> text=${containerId}`);
+
+  await explorer.click(getTreeNodeSelector(`DATA/${databaseId}`));
+  await explorer.click(getTreeNodeSelector(`DATA/${databaseId}/${containerId}`));
+
   // Delete database and collection
-  explorer.click(`[data-test="${containerId}"] [aria-label="More options"]`);
+  await openContextMenu(explorer, `DATA/${databaseId}/${containerId}`);
-  explorer.click('button[role="menuitem"]:has-text("Delete Collection")');
+  await explorer.click(getTreeMenuItemSelector(`DATA/${databaseId}/${containerId}`, "Delete Collection"));
+
+  await explorer.waitForSelector(getPanelSelector("Delete Collection"));
   await explorer.fill('text=* Confirm by typing the collection id >> input[type="text"]', containerId);
   await explorer.click('[aria-label="OK"]');
-  await explorer.click(`[data-test="${databaseId}"] [aria-label="More options"]`);
+  await explorer.waitForSelector(getPanelSelector("Delete Collection"), { state: "detached" });
-  await explorer.click('button[role="menuitem"]:has-text("Delete Database")');
+
+  await openContextMenu(explorer, `DATA/${databaseId}`);
+  await explorer.click(getTreeMenuItemSelector(`DATA/${databaseId}`, "Delete Database"));
+
+  await explorer.waitForSelector(getPanelSelector("Delete Database"));
   await explorer.click('text=* Confirm by typing the database id >> input[type="text"]');
   await explorer.fill('text=* Confirm by typing the database id >> input[type="text"]', databaseId);
   await explorer.click("#sidePanelOkButton");
+  await explorer.waitForSelector(getPanelSelector("Delete Database"), { state: "detached" });
+
   await expect(explorer).not.toHaveText(".dataResourceTree", databaseId);
   await expect(explorer).not.toHaveText(".dataResourceTree", containerId);
 });

test/sql/container.spec.ts

@@ -1,6 +1,14 @@
 import { jest } from "@jest/globals";
 import "expect-playwright";
-import { generateUniqueName, getAzureCLICredentialsToken } from "../utils/shared";
+import {
+  AccountType,
+  generateUniqueName,
+  getPanelSelector,
+  getTestExplorerUrl,
+  getTreeMenuItemSelector,
+  getTreeNodeSelector,
+  openContextMenu,
+} from "../utils/shared";
 import { waitForExplorer } from "../utils/waitForExplorer";
 jest.setTimeout(120000);
 
@@ -8,28 +16,40 @@ test("SQL CRUD", async () => {
   const databaseId = generateUniqueName("db");
   const containerId = generateUniqueName("container");
 
-  // We can't retrieve AZ CLI credentials from the browser so we get them here.
-  const token = await getAzureCLICredentialsToken();
   page.setDefaultTimeout(50000);
 
-  await page.goto(`https://localhost:1234/testExplorer.html?accountName=portal-sql-runner-west-us&token=${token}`);
+  const url = await getTestExplorerUrl(AccountType.SQL);
+  await page.goto(url);
   const explorer = await waitForExplorer();
 
   await explorer.click('[data-test="New Container"]');
+
+  await explorer.waitForSelector(getPanelSelector("New Container"));
   await explorer.fill('[aria-label="New database id, Type a new database id"]', databaseId);
   await explorer.fill('[aria-label="Container id, Example Container1"]', containerId);
   await explorer.fill('[aria-label="Partition key"]', "/pk");
   await explorer.click("#sidePanelOkButton");
-  await explorer.click(`.nodeItem >> text=${databaseId}`);
+  await explorer.waitForSelector(getPanelSelector("New Container"), { state: "detached" });
-  await explorer.click(`[data-test="${containerId}"] [aria-label="More options"]`);
+
-  await explorer.click('button[role="menuitem"]:has-text("Delete Container")');
+  await explorer.click(getTreeNodeSelector(`DATA/${databaseId}`));
+  await explorer.hover(getTreeNodeSelector(`DATA/${databaseId}/${containerId}`));
+  await openContextMenu(explorer, `DATA/${databaseId}/${containerId}`);
+  await explorer.click(getTreeMenuItemSelector(`DATA/${databaseId}/${containerId}`, "Delete Container"));
+
+  await explorer.waitForSelector(getPanelSelector("Delete Container"));
   await explorer.fill('text=* Confirm by typing the container id >> input[type="text"]', containerId);
   await explorer.click('[aria-label="OK"]');
-  await explorer.click(`[data-test="${databaseId}"] [aria-label="More options"]`);
+  await explorer.waitForSelector(getPanelSelector("Delete Container"), { state: "detached" });
-  await explorer.click('button[role="menuitem"]:has-text("Delete Database")');
+
+  await openContextMenu(explorer, `DATA/${databaseId}`);
+  await explorer.click(getTreeMenuItemSelector(`DATA/${databaseId}`, "Delete Database"));
+
+  await explorer.waitForSelector(getPanelSelector("Delete Database"));
   await explorer.click('text=* Confirm by typing the database id >> input[type="text"]');
   await explorer.fill('text=* Confirm by typing the database id >> input[type="text"]', databaseId);
   await explorer.click("#sidePanelOkButton");
+  await explorer.waitForSelector(getPanelSelector("Delete Database"), { state: "detached" });
+
   await expect(explorer).not.toHaveText(".dataResourceTree", databaseId);
   await expect(explorer).not.toHaveText(".dataResourceTree", containerId);
 });

test/sql/resourceToken.spec.ts

@@ -2,7 +2,8 @@ import { CosmosDBManagementClient } from "@azure/arm-cosmosdb";
 import { CosmosClient, PermissionMode } from "@azure/cosmos";
 import { jest } from "@jest/globals";
 import "expect-playwright";
-import { generateUniqueName, getAzureCLICredentials } from "../utils/shared";
+import { generateUniqueName, getAzureCLICredentials, getTreeNodeSelector } from "../utils/shared";
+import { waitForExplorer } from "../utils/waitForExplorer";
 jest.setTimeout(120000);
 
 const subscriptionId = process.env["AZURE_SUBSCRIPTION_ID"] ?? "";
@@ -16,7 +17,7 @@ test("Resource token", async () => {
   const dbId = generateUniqueName("db");
   const collectionId = generateUniqueName("col");
   const client = new CosmosClient({
-    endpoint: account.documentEndpoint,
+    endpoint: account.documentEndpoint!,
     key: keys.primaryMasterKey,
   });
   const { database } = await client.databases.createIfNotExists({ id: dbId });
@@ -32,11 +33,10 @@ test("Resource token", async () => {
   await page.goto("https://localhost:1234/hostedExplorer.html");
   await page.waitForSelector("div > p.switchConnectTypeText");
   await page.click("div > p.switchConnectTypeText");
-  await page.type("input[class='inputToken']", resourceTokenConnectionString);
+  await page.fill("input[class='inputToken']", resourceTokenConnectionString);
   await page.click("input[value='Connect']");
-  await page.waitForSelector("iframe");
+  const explorer = await waitForExplorer();
-  const explorer = await page.frame({
+
-    name: "explorer",
+  const collectionNodeLabel = await explorer.textContent(getTreeNodeSelector(`DATA/${collectionId}`));
-  });
+  expect(collectionNodeLabel).toBe(collectionId);
-  await explorer.textContent(`css=.dataResourceTree >> "${collectionId}"`);
 });

test/tables/container.spec.ts

@@ -1,27 +1,40 @@
 import { jest } from "@jest/globals";
 import "expect-playwright";
-import { generateUniqueName, getAzureCLICredentialsToken } from "../utils/shared";
+import {
+  AccountType,
+  generateUniqueName,
+  getPanelSelector,
+  getTestExplorerUrl,
+  getTreeMenuItemSelector,
+  openContextMenu,
+} from "../utils/shared";
 import { waitForExplorer } from "../utils/waitForExplorer";
 
 jest.setTimeout(120000);
 
 test("Tables CRUD", async () => {
   const tableId = generateUniqueName("table");
-  // We can't retrieve AZ CLI credentials from the browser so we get them here.
-  const token = await getAzureCLICredentialsToken();
   page.setDefaultTimeout(50000);
 
-  await page.goto(`https://localhost:1234/testExplorer.html?accountName=portal-tables-runner&token=${token}`);
+  const url = await getTestExplorerUrl(AccountType.Tables);
+  await page.goto(url);
   const explorer = await waitForExplorer();
 
   await page.waitForSelector('text="Querying databases"', { state: "detached" });
   await explorer.click('[data-test="New Table"]');
+
+  await explorer.waitForSelector(getPanelSelector("New Table"));
   await explorer.fill('[aria-label="Table id, Example Table1"]', tableId);
   await explorer.click("#sidePanelOkButton");
-  await explorer.click(`[data-test="TablesDB"]`);
+  await explorer.waitForSelector(getPanelSelector("New Table"), { state: "detached" });
-  await explorer.click(`[data-test="${tableId}"] [aria-label="More options"]`);
+
-  await explorer.click('button[role="menuitem"]:has-text("Delete Table")');
+  await openContextMenu(explorer, `DATA/TablesDB/${tableId}`);
+  await explorer.click(getTreeMenuItemSelector(`DATA/TablesDB/${tableId}`, "Delete Table"));
+
+  await explorer.waitForSelector(getPanelSelector("Delete Table"));
   await explorer.fill('text=* Confirm by typing the table id >> input[type="text"]', tableId);
   await explorer.click('[aria-label="OK"]');
+  await explorer.waitForSelector(getPanelSelector("Delete Table"), { state: "detached" });
+
   await expect(explorer).not.toHaveText(".dataResourceTree", tableId);
 });

test/testExplorer/TestExplorer.ts

@@ -4,9 +4,9 @@ import { DataExplorerInputsFrame } from "../../src/Contracts/ViewModels";
 import { updateUserContext } from "../../src/UserContext";
 import { get, listKeys } from "../../src/Utils/arm/generatedClients/cosmos/databaseAccounts";
 
-const resourceGroup = process.env.RESOURCE_GROUP || "";
-const subscriptionId = process.env.SUBSCRIPTION_ID || "";
 const urlSearchParams = new URLSearchParams(window.location.search);
+const resourceGroup = urlSearchParams.get("resourceGroup") || process.env.RESOURCE_GROUP || "";
+const subscriptionId = urlSearchParams.get("subscriptionId") || process.env.SUBSCRIPTION_ID || "";
 const accountName = urlSearchParams.get("accountName") || "portal-sql-runner-west-us";
 const selfServeType = urlSearchParams.get("selfServeType") || "example";
 const iframeSrc = urlSearchParams.get("iframeSrc") || "explorer.html?platform=Portal&disablePortalInitCache";

test/utils/shared.ts

@@ -1,5 +1,35 @@
 import { AzureCliCredentials } from "@azure/ms-rest-nodeauth";
 import crypto from "crypto";
+import { Frame } from "playwright";
+
+export enum AccountType {
+  Tables = "Tables",
+  Cassandra = "Cassandra",
+  Gremlin = "Gremlin",
+  Mongo = "Mongo",
+  Mongo32 = "Mongo32",
+  SQL = "SQL",
+}
+
+export const defaultAccounts: Record<AccountType, string> = {
+  [AccountType.Tables]: "portal-tables-runner",
+  [AccountType.Cassandra]: "portal-cassandra-runner",
+  [AccountType.Gremlin]: "portal-gremlin-runner",
+  [AccountType.Mongo]: "portal-mongo-runner",
+  [AccountType.Mongo32]: "portal-mongo32-runner",
+  [AccountType.SQL]: "portal-sql-runner-west-us",
+};
+
+const resourceGroup = process.env.DE_TEST_RESOURCE_GROUP ?? "runners";
+const subscriptionId = process.env.DE_TEST_SUBSCRIPTION_ID ?? "69e02f2d-f059-4409-9eac-97e8a276ae2c";
+
+export async function getTestExplorerUrl(accountType: AccountType) {
+  // We can't retrieve AZ CLI credentials from the browser so we get them here.
+  const token = await getAzureCLICredentialsToken();
+  const accountName =
+    process.env[`DE_TEST_ACCOUNT_NAME_${accountType.toLocaleUpperCase()}`] ?? defaultAccounts[accountType];
+  return `https://localhost:1234/testExplorer.html?accountName=${accountName}&resourceGroup=${resourceGroup}&subscriptionId=${subscriptionId}&token=${token}`;
+}
 
 export function generateUniqueName(baseName = "", length = 4): string {
   return `${baseName}${crypto.randomBytes(length).toString("hex")}`;
@@ -18,3 +48,22 @@ export async function getAzureCLICredentialsToken(): Promise<string> {
   const token = (await credentials.getToken()).accessToken;
   return token;
 }
+
+export function getPanelSelector(title: string) {
+  return `[data-test="Panel:${title}"]`;
+}
+
+export function getTreeNodeSelector(id: string) {
+  return `[data-test="TreeNode:${id}"]`;
+}
+
+export function getTreeMenuItemSelector(nodeId: string, itemLabel: string) {
+  return `[data-test="TreeNode/ContextMenu:${nodeId}"] [data-test="TreeNode/ContextMenuItem:${itemLabel}"]`;
+}
+
+export async function openContextMenu(explorer: Frame, nodeIdentifier: string) {
+  const nodeSelector = getTreeNodeSelector(nodeIdentifier);
+  await explorer.hover(nodeSelector);
+  await explorer.click(`${nodeSelector} [data-test="TreeNode/ContextMenuTrigger"]`);
+  await explorer.waitForSelector(`[data-test="TreeNode/ContextMenu:${nodeIdentifier}"]`);
+}

utils/generateResourceToken.js

@@ -1,24 +1,79 @@
 const { CosmosClient } = require("@azure/cosmos");
 
-async function main() {
+let endpoint = process.env.ENDPOINT;
-  const endpoint = process.env.ENDPOINT;
+let key = process.env.KEY;
-  if (!endpoint) {
+let databaseId = process.env.DATABASE;
-    throw new Error("Expected env var ENDPOINT");
+let containerId = process.env.CONTAINER;
-  }
+let partitionKey = process.env.partitionKey;
-  const key = process.env.KEY;
-  if (!key) {
-    throw new Error("Expected env var KEY");
-  }
-  const databaseId = process.env.DATABASE;
-  if (!databaseId) {
-    throw new Error("Expected env var DATABASE");
-  }
-  const containerId = process.env.CONTAINER;
-  if (!containerId) {
-    throw new Error("Expected env var CONTAINER");
-  }
-  const partitionKey = process.env.partitionKey;
 
+// Super rudimentary, but dependency-free arg parsing
+const args = process.argv.slice(2);
+for (let i = 0; i < args.length; i++) {
+  switch (args[i]) {
+    case "--help":
+      console.log("Usage: node generateResourceToken.js --endpoint <endpoint> --key <key> --database <database> --container <container> --partitionKey <partitionKey>");
+      console.log("Options");
+      console.log("  --endpoint: Cosmos DB endpoint (can also be set via ENDPOINT env var)");
+      console.log("  --key: Cosmos DB key (can also be set via KEY env var)");
+      console.log("  --database: Cosmos DB database ID (can also be set via DATABASE env var)");
+      console.log("  --container: Cosmos DB container ID (can also be set via CONTAINER env var)");
+      console.log("  --partitionKey: Cosmos DB container partition key (can also be set via partitionKey env var; Optional)");
+      process.exit(0);
+      break;
+    case "--endpoint":
+      i++;
+      if (i >= args.length) {
+        throw new Error("--endpoint requires an argument");
+      }
+      endpoint = args[i];
+      break;
+    case "--key":
+      i++;
+      if (i >= args.length) {
+        throw new Error("--key requires an argument");
+      }
+      key = args[i];
+      break;
+    case "--database":
+      i++;
+      if (i >= args.length) {
+        throw new Error("--database requires an argument");
+      }
+      databaseId = args[i];
+      break;
+    case "--container":
+      i++;
+      if (i >= args.length) {
+        throw new Error("--container requires an argument");
+      }
+      containerId = args[i];
+      break;
+    case "--partitionKey":
+      i++;
+      if (i >= args.length) {
+        throw new Error("--partitionKey requires an argument");
+      }
+      partitionKey = args[i];
+      break;
+    default:
+      throw new Error(`Unknown argument: ${args[i]}`);
+  }
+}
+
+if (!endpoint) {
+  throw new Error("Endpoint is required, either via --endpoint or ENDPOINT env var");
+}
+if (!key) {
+  throw new Error("Key is required, either via --key or KEY env var");
+}
+if (!databaseId) {
+  throw new Error("Database is required, either via --database or DATABASE env var");
+}
+if (!containerId) {
+  throw new Error("Container is required, either via --container or CONTAINER env var");
+}
+
+async function main() {
   const client = new CosmosClient({
     endpoint,
     key
@@ -53,4 +108,4 @@ async function main() {
 main().catch(error => {
   console.log("Error!");
   throw error;
-});
+}).then(() => process.exit(0));