[TreeView] Correct visibleNodes on re-render (#20157)

mui · Mar 29, 2020 · 1edf5a7 · 1edf5a7
1 parent 0b9329f
commit 1edf5a7
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 2 deletions.
diff --git a/packages/material-ui-lab/src/TreeItem/TreeItem.test.js b/packages/material-ui-lab/src/TreeItem/TreeItem.test.js
@@ -484,6 +484,45 @@ describe('<TreeItem />', () => {
           expect(getByTestId('two')).to.have.focus;
         });
 
+        it('moves focus to a child node works with a dynamic tree', () => {
+          function TestComponent() {
+            const [hide, setState] = React.useState(false);
+
+            return (
+              <React.Fragment>
+                <button
+                  data-testid="button"
+                  type="button"
+                  onClick={() => setState(value => !value)}
+                >
+                  Toggle Hide
+                </button>
+                <TreeView defaultExpanded={['one']}>
+                  {!hide && (
+                    <TreeItem nodeId="one" label="one" data-testid="one">
+                      <TreeItem nodeId="two" label="two" data-testid="two" />
+                    </TreeItem>
+                  )}
+                  <TreeItem nodeId="three" label="three" />
+                </TreeView>
+              </React.Fragment>
+            );
+          }
+
+          const { queryByTestId, getByTestId } = render(<TestComponent />);
+
+          expect(getByTestId('one')).to.not.be.null;
+          fireEvent.click(getByTestId('button'));
+          expect(queryByTestId('one')).to.be.null;
+          fireEvent.click(getByTestId('button'));
+          expect(getByTestId('one')).to.not.be.null;
+
+          getByTestId('one').focus();
+          fireEvent.keyDown(document.activeElement, { key: 'ArrowDown' });
+
+          expect(getByTestId('two')).to.have.focus;
+        });
+
         it("moves focus to a parent's sibling", () => {
           const { getByTestId, getByText } = render(
             <TreeView defaultExpanded={['one']}>

diff --git a/packages/material-ui-lab/src/TreeView/TreeView.js b/packages/material-ui-lab/src/TreeView/TreeView.js
@@ -468,7 +468,6 @@ const TreeView = React.forwardRef(function TreeView(props, ref) {
         if (index === 0) {
           setTabbable(id);
         }
-        nodeMap.current[id] = { parent: null };
       });
       visibleNodes.current = nodeMap.current[-1].children;
       prevChildIds.current = childIds;
@@ -493,7 +492,7 @@ const TreeView = React.forwardRef(function TreeView(props, ref) {
     if (childrenCalculated) {
       visibleNodes.current = buildVisible(nodeMap.current[-1].children);
     }
-  }, [expanded, childrenCalculated, isExpanded]);
+  }, [expanded, childrenCalculated, isExpanded, children]);
 
   const noopSelection = () => {
     return false;