[TreeView] correct visibleNodes on re-render

mui · Mar 17, 2020 · d2eb0d2 · d2eb0d2
1 parent 701e3ad
commit d2eb0d2
Show file tree

Hide file tree

Showing 2 changed files with 35 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
@@ -436,6 +436,40 @@ describe('<TreeItem />', () => {
           expect(getByTestId('two')).to.have.focus;
         });
 
+        it('moves focus to a sibling node after conditional rendering', () => {
+          function TestComponent() {
+            const [hide, setState] = React.useState(false);
+
+            return (
+              <React.Fragment>
+                <button type="button" onClick={() => setState(value => !value)}>
+                  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 { getByText, queryByText, getByTestId } = render(<TestComponent />);
+
+          expect(getByText('one')).to.not.be.null;
+          fireEvent.click(getByText('Hide'));
+          expect(queryByText('one')).to.be.null;
+          fireEvent.click(getByText('Hide'));
+          expect(getByText('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 child node', () => {
           const { getByTestId } = render(
             <TreeView defaultExpanded={['one']}>

diff --git a/packages/material-ui-lab/src/TreeView/TreeView.js b/packages/material-ui-lab/src/TreeView/TreeView.js
@@ -455,7 +455,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;
@@ -480,7 +479,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]);
 
   return (
     <TreeViewContext.Provider