fix: tree component selection

- added selectedItems prop to useItemSelected hook to make it
controllable input like component.
- added active-node class to TreeNodeBase

src/components/Tree/Tree.stories.tsx

@@ -1,4 +1,4 @@
-import React, { useMemo, useRef, useState } from 'react';
+import React, { useCallback, useMemo, useRef, useState } from 'react';
 import Text from '../Text';
 import { Template } from '../../storybook/helper.stories.templates';
 import { DocumentationPage } from '../../storybook/helper.stories.docs';
@@ -17,6 +17,7 @@ import ButtonCircle from '../ButtonCircle';
 import { TreeNodeRecord, TreeRefObject } from './Tree.types';
 import ButtonPill from '../ButtonPill';
 import Flex from '../Flex';
+import { v4 as uuidV4 } from 'uuid';
 
 // prettier-ignore
 const exampleTree =
@@ -190,58 +191,119 @@ const DynamicTree = Template(() => {
   );
 }).bind({});
 
-const SelectionTemplate = (props: Partial<TreeProps>) =>
-  Template(() => {
-    const ref = useRef<TreeRefObject>();
-    const [selected, setSelected] = useState('');
-
-    const onSelectHandler = (ids: Array<string>) => setSelected(ids.join(', '));
-
-    return (
-      <>
-        <Flex alignItems={'center'} xgap={'1rem'}>
-          <ButtonPill onPress={() => ref.current.clearSelection()}>Clear Selection</ButtonPill>
-          <Text>Selected nodes: [{selected}]</Text>
-        </Flex>
-        <hr />
-        <Tree
-          ref={ref as any}
-          treeStructure={exampleTree}
-          isRenderedFlat={true}
-          shouldNodeFocusBeInset={true}
-          onSelectionChange={onSelectHandler}
-          {...props}
-        >
-          {mapTree(exampleTreeMap, (node) => (
-            <ExampleTreeNode key={node.id.toString()} node={node} />
-          ))}
-        </Tree>
-      </>
-    );
-  }).bind({});
-
-const SingleSelectLeafNodesOnly = SelectionTemplate({
-  selectionMode: 'single',
-  selectableNodes: 'leafOnly',
-});
-
-const SingleSelectAnyNodesRequired = SelectionTemplate({
-  selectionMode: 'single',
-  selectableNodes: 'any',
-  isRequired: true,
-});
-
-const MultiSelectLeafNodesOnly = SelectionTemplate({
-  selectionMode: 'multiple',
-  selectableNodes: 'leafOnly',
-});
-
-export {
-  Example,
-  WithRoot,
-  TreeWithScroll,
-  DynamicTree,
-  SingleSelectLeafNodesOnly,
-  SingleSelectAnyNodesRequired,
-  MultiSelectLeafNodesOnly,
+const SelectableTree = ({
+  templateTitle,
+  showSelectedItems,
+  ...props
+}: Partial<TreeProps> & { templateTitle: string; showSelectedItems?: boolean }) => {
+  const ref = useRef<TreeRefObject>();
+  const [selected, setSelected] = useState('');
+  const [id] = useState(uuidV4);
+
+  const onSelectHandler = (ids: Array<string>) => setSelected(ids.join(', '));
+  const clearAll = useCallback(() => ref.current.clearSelection(), [ref]);
+  const selectAll = useCallback(
+    () =>
+      ref.current.updateSelection(
+        Array.from(exampleTreeMap.values())
+          .filter((node) => (props.selectableNodes === 'leafOnly' ? node.isLeaf : true))
+          .map((node) => node.id)
+      ),
+    [ref, exampleTree]
+  );
+
+  return (
+    <div>
+      <h2 id={id}>{templateTitle}</h2>
+      <Flex alignItems={'center'} xgap={'0.5rem'}>
+        <ButtonPill size={28} ghost outline onPress={clearAll}>
+          Clear Selection
+        </ButtonPill>
+        {props.selectionMode === 'multiple' && (
+          <ButtonPill size={28} ghost outline onPress={selectAll}>
+            Select All
+          </ButtonPill>
+        )}
+      </Flex>
+      {showSelectedItems && <Text>Selected nodes: [{selected}]</Text>}
+      <hr />
+      <Tree
+        ref={ref as any}
+        aria-labelledby={id}
+        treeStructure={exampleTree}
+        isRenderedFlat={true}
+        shouldNodeFocusBeInset={true}
+        onSelectionChange={onSelectHandler}
+        {...props}
+      >
+        {mapTree(exampleTreeMap, (node) => (
+          <ExampleTreeNode key={node.id.toString()} node={node} />
+        ))}
+      </Tree>
+    </div>
+  );
 };
+
+const NodeSelection = Template(() => {
+  return (
+    <Flex justifyContent="space-around">
+      <SelectableTree
+        templateTitle="Single Selection, Leaf nodes only"
+        showSelectedItems={true}
+        selectionMode="single"
+        selectableNodes="leafOnly"
+      />
+
+      <SelectableTree
+        templateTitle="Single Selection, Any nodes"
+        showSelectedItems={true}
+        selectionMode="single"
+        selectableNodes="any"
+      />
+
+      <SelectableTree
+        templateTitle="Multi Selection, Leaf nodes only"
+        showSelectedItems={true}
+        selectionMode="multiple"
+        selectableNodes="leafOnly"
+      />
+    </Flex>
+  );
+}).bind({});
+
+const ControlledSelection = Template(() => {
+  const [selected, setSelected] = useState<Array<string>>([]);
+
+  const onSelect = useCallback(
+    (ids: Array<string>) => {
+      setSelected(ids);
+    },
+    [setSelected]
+  );
+
+  return (
+    <>
+      <h2>Sync selection between trees</h2>
+      <Text>Selected nodes: [{selected.join(', ')}]</Text>
+      <Flex justifyContent="space-around">
+        <SelectableTree
+          templateTitle="Tree A"
+          selectionMode="multiple"
+          selectableNodes="leafOnly"
+          selectedItems={selected}
+          onSelectionChange={onSelect}
+        />
+
+        <SelectableTree
+          templateTitle="Tree B"
+          selectionMode="multiple"
+          selectableNodes="leafOnly"
+          selectedItems={selected}
+          onSelectionChange={onSelect}
+        />
+      </Flex>
+    </>
+  );
+}).bind({});
+
+export { Example, WithRoot, TreeWithScroll, DynamicTree, NodeSelection, ControlledSelection };

src/components/Tree/Tree.test.tsx

@@ -3,6 +3,7 @@ import { mount } from 'enzyme';
 import '@testing-library/jest-dom';
 import { render, waitFor, waitForElementToBeRemoved } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
+import { act } from 'react-test-renderer';
 
 import Tree, { TREE_CONSTANTS, TreeProps } from './index';
 import { createTreeNode as tNode } from './test.utils';
@@ -698,19 +699,9 @@ describe('<Tree />', () => {
   });
 
   describe('selection', () => {
-    const getTreeComponent = (
-      tree,
-      { isRequired, selectionMode, selectableNodes, onSelectionChange }: Partial<TreeProps> = {}
-    ) => {
+    const getTreeComponent = (tree, { ref, ...props }: Partial<TreeProps> = {}) => {
       return (
-        <Tree
-          treeStructure={tree}
-          excludeTreeRoot={false}
-          selectableNodes={selectableNodes}
-          onSelectionChange={onSelectionChange}
-          selectionMode={selectionMode}
-          isRequired={isRequired}
-        >
+        <Tree treeStructure={tree} excludeTreeRoot={false} {...props} ref={ref as any}>
           {mapTree(
             convertNestedTree2MappedTree(tree),
             (node) => (
@@ -809,5 +800,42 @@ describe('<Tree />', () => {
         }
       );
     });
+
+    it('should update selected items based on the selectionMode', async () => {
+      const tree = getSampleTree();
+      const { rerender, getByTestId } = render(
+        getTreeComponent(tree, { selectionMode: 'multiple', selectedItems: ['1', '2.2'] })
+      );
+
+      expect(getByTestId('1')).toHaveAttribute('aria-selected', 'true');
+      expect(getByTestId('2.2')).toHaveAttribute('aria-selected', 'true');
+
+      rerender(getTreeComponent(tree, { selectionMode: 'none', selectedItems: ['2.1', '4'] }));
+
+      expect(getByTestId('1')).not.toHaveAttribute('aria-selected', 'true');
+      expect(getByTestId('2.2')).not.toHaveAttribute('aria-selected', 'true');
+    });
+
+    describe('controlled tree selection', () => {
+      it('should change selection based on the selectedItems prop', async () => {
+        const tree = getSampleTree();
+        const { rerender, getByTestId } = render(
+          getTreeComponent(tree, { selectionMode: 'multiple', selectedItems: ['1', '2.2'] })
+        );
+
+        expect(getByTestId('1')).toHaveAttribute('aria-selected', 'true');
+        expect(getByTestId('2.2')).toHaveAttribute('aria-selected', 'true');
+
+        rerender(
+          getTreeComponent(tree, { selectionMode: 'multiple', selectedItems: ['2.1', '4'] })
+        );
+
+        expect(getByTestId('1')).not.toHaveAttribute('aria-selected', 'true');
+        expect(getByTestId('2.2')).not.toHaveAttribute('aria-selected', 'true');
+
+        expect(getByTestId('2.1')).toHaveAttribute('aria-selected', 'true');
+        expect(getByTestId('4')).toHaveAttribute('aria-selected', 'true');
+      });
+    });
   });
 });

src/components/Tree/Tree.tsx

@@ -42,6 +42,7 @@ const Tree = forwardRef((props: Props, ref: ForwardedRef<TreeRefObject>) => {
     selectableNodes = DEFAULTS.SELECTABLE_NODES,
     selectedByDefault,
     onSelectionChange,
+    selectedItems,
     isRequired = DEFAULTS.IS_REQUIRED,
     ...rest
   } = props;
@@ -51,6 +52,7 @@ const Tree = forwardRef((props: Props, ref: ForwardedRef<TreeRefObject>) => {
   const itemSelection = useItemSelected<TreeNodeId>({
     selectionMode,
     selectedByDefault,
+    selectedItems,
     onSelectionChange,
     isRequired,
   });

src/components/TreeNodeBase/TreeNodeBase.style.scss

@@ -17,7 +17,7 @@
   }
 
   &:active,
-  &.active {
+  &.selected {
     color: var(--mds-color-theme-text-primary-normal);
     background-color: var(--mds-color-theme-background-primary-active);
   }

src/components/TreeNodeBase/TreeNodeBase.test.tsx

@@ -465,7 +465,7 @@ describe('TreeNodeBase', () => {
       expect(element.getAttribute('data-shape')).toBe(shape);
     });
 
-    it('should have provided active class when isSelected is provided', () => {
+    it('should have provided selected class when isSelected is provided', () => {
       expect.assertions(2);
 
       treeContextMock.itemSelection.selectionMode = 'single';
@@ -475,9 +475,23 @@ describe('TreeNodeBase', () => {
 
       const element = container.find(TreeNodeBase).getDOMNode();
 
-      expect(element.classList.contains('active')).toBe(true);
+      expect(element.classList.contains('selected')).toBe(true);
       expect(element.getAttribute('aria-selected')).toBe('true');
     });
+
+    it('should have provided active-node class the tree node is active in the tree', () => {
+      expect.assertions(1);
+
+      container = mount(
+        <TreeContext.Provider value={{ activeNodeId: 42 } as any}>
+          <TreeNodeBase nodeId="42">{() => 'Test'}</TreeNodeBase>
+        </TreeContext.Provider>
+      );
+
+      const element = container.find(TreeNodeBase).getDOMNode();
+
+      expect(element.classList.contains('active-node')).toBe(true);
+    });
   });
 
   describe('actions', () => {