[TreeView] Allow TreeItem to have conditional child (#20238)

mui · Mar 29, 2020 · 8b795b2 · 8b795b2
1 parent 1edf5a7
commit 8b795b2
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 1 deletion.
diff --git a/packages/material-ui-lab/src/TreeItem/TreeItem.js b/packages/material-ui-lab/src/TreeItem/TreeItem.js
@@ -325,8 +325,13 @@ const TreeItem = React.forwardRef(function TreeItem(props, ref) {
   };
 
   React.useEffect(() => {
-    const childIds = React.Children.map(children, (child) => child.props.nodeId) || [];
     if (addNodeToNodeMap) {
+      const childIds = [];
+      React.Children.forEach(children, (child) => {
+        if (React.isValidElement(child) && child.props.nodeId) {
+          childIds.push(child.props.nodeId);
+        }
+      });
       addNodeToNodeMap(nodeId, childIds);
     }
   }, [children, nodeId, addNodeToNodeMap]);

diff --git a/packages/material-ui-lab/src/TreeItem/TreeItem.test.js b/packages/material-ui-lab/src/TreeItem/TreeItem.test.js
@@ -77,6 +77,32 @@ describe('<TreeItem />', () => {
     expect(getIcon('6')).attribute('data-test').to.equal('endIcon');
   });
 
+  it('should allow conditional child', () => {
+    function TestComponent() {
+      const [hide, setState] = React.useState(false);
+
+      return (
+        <React.Fragment>
+          <button data-testid="button" type="button" onClick={() => setState(true)}>
+            Hide
+          </button>
+          <TreeView defaultExpanded={['1']}>
+            <TreeItem nodeId="1" data-testid="1">
+              {!hide && <TreeItem nodeId="2" data-testid="2" />}
+            </TreeItem>
+          </TreeView>
+        </React.Fragment>
+      );
+    }
+    const { getByTestId, queryByTestId } = render(<TestComponent />);
+
+    expect(getByTestId('1')).to.have.attribute('aria-expanded', 'true');
+    expect(getByTestId('2')).to.not.be.null;
+    fireEvent.click(getByTestId('button'));
+    expect(getByTestId('1')).to.not.have.attribute('aria-expanded');
+    expect(queryByTestId('2')).to.be.null;
+  });
+
   it('should treat an empty array equally to no children', () => {
     const { getByTestId } = render(
       <TreeView defaultExpanded={['1']}>