diff --git a/packages/material-ui-lab/src/TreeItem/TreeItem.test.js b/packages/material-ui-lab/src/TreeItem/TreeItem.test.js
index 6409de5ea6cca0..3d13c4c6d8e1dd 100644
--- a/packages/material-ui-lab/src/TreeItem/TreeItem.test.js
+++ b/packages/material-ui-lab/src/TreeItem/TreeItem.test.js
@@ -436,14 +436,33 @@ describe('<TreeItem />', () => {
           expect(getByTestId('two')).to.have.focus;
         });
 
-        it('moves focus to a sibling node after conditional rendering', () => {
+        it('moves focus to a child node', () => {
+          const { getByTestId } = render(
+            <TreeView defaultExpanded={['one']}>
+              <TreeItem nodeId="one" label="one" data-testid="one">
+                <TreeItem nodeId="two" label="two" data-testid="two" />
+              </TreeItem>
+            </TreeView>,
+          );
+
+          expect(getByTestId('one')).to.have.attribute('aria-expanded', 'true');
+          getByTestId('one').focus();
+          fireEvent.keyDown(document.activeElement, { key: 'ArrowDown' });
+          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 type="button" onClick={() => setState(value => !value)}>
-                  Hide
+                <button
+                  data-testid="button"
+                  type="button"
+                  onClick={() => setState(value => !value)}
+                >
+                  Toggle Hide
                 </button>
                 <TreeView defaultExpanded={['one']}>
                   {!hide && (
@@ -457,31 +476,17 @@ describe('<TreeItem />', () => {
             );
           }
 
-          const { getByText, queryByText, getByTestId } = render(<TestComponent />);
+          const { queryByTestId, 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;
-        });
+          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;
 
-        it('moves focus to a child node', () => {
-          const { getByTestId } = render(
-            <TreeView defaultExpanded={['one']}>
-              <TreeItem nodeId="one" label="one" data-testid="one">
-                <TreeItem nodeId="two" label="two" data-testid="two" />
-              </TreeItem>
-            </TreeView>,
-          );
-
-          expect(getByTestId('one')).to.have.attribute('aria-expanded', 'true');
           getByTestId('one').focus();
           fireEvent.keyDown(document.activeElement, { key: 'ArrowDown' });
+
           expect(getByTestId('two')).to.have.focus;
         });