[TreeView] Correctly select items in deeply nested trees (#26413)

mui · May 22, 2021 · 5371b14 · 5371b14
1 parent 606c288
commit 5371b14
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 12 deletions.
diff --git a/packages/material-ui-lab/src/TreeView/TreeView.js b/packages/material-ui-lab/src/TreeView/TreeView.js
@@ -198,22 +198,20 @@ const TreeView = React.forwardRef(function TreeView(inProps, ref) {
       return getNavigableChildrenIds(id)[0];
     }
 
-    // Try to get next sibling
-    const node = nodeMap.current[id];
-    const siblings = getNavigableChildrenIds(node.parentId);
+    let node = nodeMap.current[id];
+    while (node != null) {
+      // Try to get next sibling
+      const siblings = getNavigableChildrenIds(node.parentId);
+      const nextSibling = siblings[siblings.indexOf(node.id) + 1];
 
-    const nextSibling = siblings[siblings.indexOf(id) + 1];
+      if (nextSibling) {
+        return nextSibling;
+      }
 
-    if (nextSibling) {
-      return nextSibling;
+      // If the sibling does not exist, go up a level to the parent and try again.
+      node = nodeMap.current[node.parentId];
     }
 
-    // try to get parent's next sibling
-    const parent = nodeMap.current[node.parentId];
-    if (parent) {
-      const parentSiblings = getNavigableChildrenIds(parent.parentId);
-      return parentSiblings[parentSiblings.indexOf(parent.id) + 1];
-    }
     return null;
   };
 

diff --git a/packages/material-ui-lab/src/TreeView/TreeView.test.js b/packages/material-ui-lab/src/TreeView/TreeView.test.js
@@ -68,6 +68,24 @@ describe('<TreeView />', () => {
       fireEvent.click(screen.getByText('one'), { shiftKey: true });
     });
 
+    it('should not crash when selecting multiple items in a deeply nested tree', () => {
+      render(
+        <TreeView multiSelect defaultExpanded={['1', '1.1', '2']}>
+          <TreeItem nodeId="1" label="Item 1">
+            <TreeItem nodeId="1.1" label="Item 1.1">
+              <TreeItem nodeId="1.1.1" data-testid="item-1.1.1" label="Item 1.1.1" />
+            </TreeItem>
+          </TreeItem>
+          <TreeItem nodeId="2" data-testid="item-2" label="Item 2" />
+        </TreeView>,
+      );
+      fireEvent.click(screen.getByText('Item 1.1.1'));
+      fireEvent.click(screen.getByText('Item 2'), { shiftKey: true });
+
+      expect(screen.getByTestId('item-1.1.1')).to.have.attribute('aria-selected', 'true');
+      expect(screen.getByTestId('item-2')).to.have.attribute('aria-selected', 'true');
+    });
+
     it('should not crash on keydown on an empty tree', () => {
       render(<TreeView />);