From 55bc02430fe898e00f6fcc2cb73b9c917f71650a Mon Sep 17 00:00:00 2001
From: Flavien DELANGLE <flaviendelangle@gmail.com>
Date: Thu, 25 Apr 2024 16:12:14 +0200
Subject: [PATCH] [core] Use `describeTreeView` for type-ahead tests (#12811)

---
 .../src/TreeItem/TreeItem.test.tsx            | 155 -----------
 .../useTreeViewKeyboardNavigation.test.tsx    | 251 ++++++++++++++----
 .../useTreeViewKeyboardNavigation.ts          |   6 +-
 .../describeTreeView/describeTreeView.tsx     |  12 +-
 .../describeTreeView.types.ts                 |   2 +-
 5 files changed, 219 insertions(+), 207 deletions(-)

diff --git a/packages/x-tree-view/src/TreeItem/TreeItem.test.tsx b/packages/x-tree-view/src/TreeItem/TreeItem.test.tsx
index beebfefe762d1..857b8866ae147 100644
--- a/packages/x-tree-view/src/TreeItem/TreeItem.test.tsx
+++ b/packages/x-tree-view/src/TreeItem/TreeItem.test.tsx
@@ -683,129 +683,6 @@ describe('<TreeItem />', () => {
         });
       });
 
-      describe('type-ahead functionality', () => {
-        it('moves focus to the next item with a name that starts with the typed character', () => {
-          const { getByTestId } = render(
-            <SimpleTreeView>
-              <TreeItem itemId="one" label="one" data-testid="one" />
-              <TreeItem itemId="two" label={<span>two</span>} data-testid="two" />
-              <TreeItem itemId="three" label="three" data-testid="three" />
-              <TreeItem itemId="four" label="four" data-testid="four" />
-            </SimpleTreeView>,
-          );
-
-          act(() => {
-            getByTestId('one').focus();
-          });
-
-          expect(getByTestId('one')).toHaveFocus();
-
-          fireEvent.keyDown(getByTestId('one'), { key: 't' });
-
-          expect(getByTestId('two')).toHaveFocus();
-
-          fireEvent.keyDown(getByTestId('two'), { key: 'f' });
-
-          expect(getByTestId('four')).toHaveFocus();
-
-          fireEvent.keyDown(getByTestId('four'), { key: 'o' });
-
-          expect(getByTestId('one')).toHaveFocus();
-        });
-
-        it('moves focus to the next item with the same starting character', () => {
-          const { getByTestId } = render(
-            <SimpleTreeView>
-              <TreeItem itemId="one" label="one" data-testid="one" />
-              <TreeItem itemId="two" label="two" data-testid="two" />
-              <TreeItem itemId="three" label="three" data-testid="three" />
-              <TreeItem itemId="four" label="four" data-testid="four" />
-            </SimpleTreeView>,
-          );
-
-          act(() => {
-            getByTestId('one').focus();
-          });
-
-          expect(getByTestId('one')).toHaveFocus();
-
-          fireEvent.keyDown(getByTestId('one'), { key: 't' });
-
-          expect(getByTestId('two')).toHaveFocus();
-
-          fireEvent.keyDown(getByTestId('two'), { key: 't' });
-
-          expect(getByTestId('three')).toHaveFocus();
-
-          fireEvent.keyDown(getByTestId('three'), { key: 't' });
-
-          expect(getByTestId('two')).toHaveFocus();
-        });
-
-        it('should not move focus when pressing a modifier key + letter', () => {
-          const { getByTestId } = render(
-            <SimpleTreeView>
-              <TreeItem itemId="one" data-testid="one" />
-              <TreeItem itemId="two" data-testid="two" />
-              <TreeItem itemId="three" data-testid="three" />
-              <TreeItem itemId="four" data-testid="four" />
-            </SimpleTreeView>,
-          );
-
-          act(() => {
-            getByTestId('one').focus();
-          });
-
-          expect(getByTestId('one')).toHaveFocus();
-
-          fireEvent.keyDown(getByTestId('one'), { key: 'f', ctrlKey: true });
-
-          expect(getByTestId('one')).toHaveFocus();
-
-          fireEvent.keyDown(getByTestId('one'), { key: 'f', metaKey: true });
-
-          expect(getByTestId('one')).toHaveFocus();
-
-          fireEvent.keyDown(getByTestId('one'), { key: 'f', shiftKey: true });
-
-          expect(getByTestId('one')).toHaveFocus();
-        });
-
-        it('should not throw when an item is removed', () => {
-          function TestComponent() {
-            const [hide, setState] = React.useState(false);
-            return (
-              <React.Fragment>
-                <button type="button" onClick={() => setState(true)}>
-                  Hide
-                </button>
-                <SimpleTreeView>
-                  {!hide && <TreeItem itemId="hide" label="ab" />}
-                  <TreeItem itemId="keyDown" label="keyDown" data-testid="keyDown" />
-                  <TreeItem itemId="navTo" label="ac" data-testid="navTo" />
-                </SimpleTreeView>
-              </React.Fragment>
-            );
-          }
-
-          const { getByText, getByTestId } = render(<TestComponent />);
-          fireEvent.click(getByText('Hide'));
-          expect(getByTestId('navTo')).not.toHaveFocus();
-
-          expect(() => {
-            act(() => {
-              getByTestId('keyDown').focus();
-            });
-
-            expect(getByTestId('keyDown')).toHaveFocus();
-
-            fireEvent.keyDown(getByTestId('keyDown'), { key: 'a' });
-          }).not.to.throw();
-
-          expect(getByTestId('navTo')).toHaveFocus();
-        });
-      });
-
       describe('asterisk key interaction', () => {
         it('expands all siblings that are at the same level as the current item', () => {
           const onExpandedItemsChange = spy();
@@ -1545,22 +1422,6 @@ describe('<TreeItem />', () => {
 
     describe('focus', () => {
       describe('`disabledItemsFocusable={true}`', () => {
-        it('should not prevent focus by type-ahead', () => {
-          const { getByTestId } = render(
-            <SimpleTreeView disabledItemsFocusable>
-              <TreeItem itemId="one" label="one" data-testid="one" />
-              <TreeItem itemId="two" label="two" disabled data-testid="two" />
-            </SimpleTreeView>,
-          );
-
-          act(() => {
-            getByTestId('one').focus();
-          });
-          expect(getByTestId('one')).toHaveFocus();
-          fireEvent.keyDown(getByTestId('one'), { key: 't' });
-          expect(getByTestId('two')).toHaveFocus();
-        });
-
         it('should not prevent focus by arrow keys', () => {
           const { getByTestId } = render(
             <SimpleTreeView disabledItemsFocusable>
@@ -1581,22 +1442,6 @@ describe('<TreeItem />', () => {
       });
 
       describe('`disabledItemsFocusable=false`', () => {
-        it('should prevent focus by type-ahead', () => {
-          const { getByTestId } = render(
-            <SimpleTreeView>
-              <TreeItem itemId="one" label="one" data-testid="one" />
-              <TreeItem itemId="two" label="two" disabled data-testid="two" />
-            </SimpleTreeView>,
-          );
-
-          act(() => {
-            getByTestId('one').focus();
-          });
-          expect(getByTestId('one')).toHaveFocus();
-          fireEvent.keyDown(getByTestId('one'), { key: 't' });
-          expect(getByTestId('one')).toHaveFocus();
-        });
-
         it('should be skipped on navigation with arrow keys', () => {
           const { getByTestId } = render(
             <SimpleTreeView>
diff --git a/packages/x-tree-view/src/internals/plugins/useTreeViewKeyboardNavigation/useTreeViewKeyboardNavigation.test.tsx b/packages/x-tree-view/src/internals/plugins/useTreeViewKeyboardNavigation/useTreeViewKeyboardNavigation.test.tsx
index 970bed10f7210..5b5c7a4990f52 100644
--- a/packages/x-tree-view/src/internals/plugins/useTreeViewKeyboardNavigation/useTreeViewKeyboardNavigation.test.tsx
+++ b/packages/x-tree-view/src/internals/plugins/useTreeViewKeyboardNavigation/useTreeViewKeyboardNavigation.test.tsx
@@ -1,53 +1,208 @@
 import * as React from 'react';
 import { expect } from 'chai';
-import { act, createRenderer, fireEvent } from '@mui-internal/test-utils';
-import { RichTreeView } from '@mui/x-tree-view/RichTreeView';
-
-describe('useTreeViewKeyboardNavigation', () => {
-  const { render } = createRenderer();
-
-  it('should work after adding / removing items', () => {
-    const { getByRole, setProps } = render(
-      <RichTreeView
-        items={[
-          { id: 'one', label: 'one' },
-          { id: 'two', label: 'two' },
-          { id: 'three', label: 'three' },
-          { id: 'four', label: 'four' },
-        ]}
-      />,
-    );
-
-    act(() => {
-      getByRole('treeitem', { name: 'one' }).focus();
-    });
+import { act, fireEvent } from '@mui-internal/test-utils';
+import { describeTreeView } from 'test/utils/tree-view/describeTreeView';
+import {
+  UseTreeViewItemsSignature,
+  UseTreeViewKeyboardNavigationSignature,
+} from '@mui/x-tree-view/internals';
 
-    fireEvent.keyDown(getByRole('treeitem', { name: 'one' }), { key: 'f' });
-    expect(getByRole('treeitem', { name: 'four' })).toHaveFocus();
+describeTreeView<[UseTreeViewKeyboardNavigationSignature, UseTreeViewItemsSignature]>(
+  'useTreeViewKeyboardNavigation',
+  ({ render, treeViewComponent }) => {
+    describe('Type-ahead', () => {
+      it('should move the focus to the next item with a name that starts with the typed character', () => {
+        const { getFocusedItemId, getItemRoot } = render({
+          items: [
+            { id: '1', label: 'one' },
+            { id: '2', label: 'two' },
+            { id: '3', label: 'three' },
+            { id: '4', label: 'four' },
+          ],
+        });
 
-    setProps({
-      items: [
-        { id: 'one', label: 'one' },
-        { id: 'two', label: 'two' },
-        { id: 'three', label: 'three' },
-      ],
-    });
-    expect(getByRole('treeitem', { name: 'one' })).toHaveFocus();
-
-    fireEvent.keyDown(getByRole('treeitem', { name: 'one' }), { key: 't' });
-    expect(getByRole('treeitem', { name: 'two' })).toHaveFocus();
-
-    setProps({
-      items: [
-        { id: 'one', label: 'one' },
-        { id: 'two', label: 'two' },
-        { id: 'three', label: 'three' },
-        { id: 'four', label: 'four' },
-      ],
-    });
-    expect(getByRole('treeitem', { name: 'two' })).toHaveFocus();
+        act(() => {
+          getItemRoot('1').focus();
+        });
+        expect(getFocusedItemId()).to.equal('1');
+
+        fireEvent.keyDown(getItemRoot('1'), { key: 't' });
+        expect(getFocusedItemId()).to.equal('2');
+
+        fireEvent.keyDown(getItemRoot('2'), { key: 'f' });
+        expect(getFocusedItemId()).to.equal('4');
+
+        fireEvent.keyDown(getItemRoot('4'), { key: 'o' });
+        expect(getFocusedItemId()).to.equal('1');
+      });
+
+      it('should move to the next item in the displayed order when typing the same starting character', () => {
+        const { getFocusedItemId, getItemRoot } = render({
+          items: [{ id: 'A1' }, { id: 'B1' }, { id: 'A2' }, { id: 'B3' }, { id: 'B2' }],
+        });
+
+        act(() => {
+          getItemRoot('A1').focus();
+        });
+        expect(getFocusedItemId()).to.equal('A1');
+
+        fireEvent.keyDown(getItemRoot('A1'), { key: 'b' });
+        expect(getFocusedItemId()).to.equal('B1');
+
+        fireEvent.keyDown(getItemRoot('B1'), { key: 'b' });
+        expect(getFocusedItemId()).to.equal('B3');
+
+        fireEvent.keyDown(getItemRoot('B3'), { key: 'b' });
+        expect(getFocusedItemId()).to.equal('B2');
+
+        fireEvent.keyDown(getItemRoot('B2'), { key: 'b' });
+        expect(getFocusedItemId()).to.equal('B1');
+      });
+
+      it('should work with capitalized label', () => {
+        const { getFocusedItemId, getItemRoot } = render({
+          items: [
+            { id: '1', label: 'One' },
+            { id: '2', label: 'Two' },
+            { id: '3', label: 'Three' },
+            { id: '4', label: 'Four' },
+          ],
+        });
+
+        act(() => {
+          getItemRoot('1').focus();
+        });
+        expect(getFocusedItemId()).to.equal('1');
+
+        fireEvent.keyDown(getItemRoot('1'), { key: 't' });
+        expect(getFocusedItemId()).to.equal('2');
+
+        fireEvent.keyDown(getItemRoot('2'), { key: 'f' });
+        expect(getFocusedItemId()).to.equal('4');
+
+        fireEvent.keyDown(getItemRoot('4'), { key: 'o' });
+        expect(getFocusedItemId()).to.equal('1');
+      });
+
+      it('should work with ReactElement label', function test() {
+        // Only the SimpleTreeView can have React Element labels.
+        if (treeViewComponent !== 'SimpleTreeView') {
+          this.skip();
+        }
+
+        const { getFocusedItemId, getItemRoot } = render({
+          items: [
+            { id: '1', label: <span>one</span> },
+            { id: '2', label: <span>two</span> },
+            { id: '3', label: <span>three</span> },
+            { id: '4', label: <span>four</span> },
+          ],
+        });
+
+        act(() => {
+          getItemRoot('1').focus();
+        });
+        expect(getFocusedItemId()).to.equal('1');
+
+        fireEvent.keyDown(getItemRoot('1'), { key: 't' });
+        expect(getFocusedItemId()).to.equal('2');
 
-    fireEvent.keyDown(getByRole('treeitem', { name: 'two' }), { key: 'f' });
-    expect(getByRole('treeitem', { name: 'four' })).toHaveFocus();
-  });
-});
+        fireEvent.keyDown(getItemRoot('2'), { key: 'f' });
+        expect(getFocusedItemId()).to.equal('4');
+
+        fireEvent.keyDown(getItemRoot('4'), { key: 'o' });
+        expect(getFocusedItemId()).to.equal('1');
+      });
+
+      it('should work after adding / removing items', () => {
+        const { getFocusedItemId, getItemRoot, setItems } = render({
+          items: [{ id: '1' }, { id: '2' }, { id: '3' }, { id: '4' }],
+        });
+
+        act(() => {
+          getItemRoot('1').focus();
+        });
+
+        fireEvent.keyDown(getItemRoot('1'), { key: '4' });
+        expect(getFocusedItemId()).to.equal('4');
+
+        setItems([{ id: '1' }, { id: '2' }, { id: '3' }]);
+        expect(getFocusedItemId()).to.equal('1');
+
+        fireEvent.keyDown(getItemRoot('1'), { key: '2' });
+        expect(getFocusedItemId()).to.equal('2');
+
+        setItems([{ id: '1' }, { id: '2' }, { id: '3' }, { id: '4' }]);
+        expect(getFocusedItemId()).to.equal('2');
+
+        fireEvent.keyDown(getItemRoot('2'), { key: '4' });
+        expect(getFocusedItemId()).to.equal('4');
+      });
+
+      it('should not move focus when pressing a modifier key and a letter', () => {
+        const { getFocusedItemId, getItemRoot } = render({
+          items: [
+            { id: '1', label: 'one' },
+            { id: '2', label: 'two' },
+            { id: '3', label: 'three' },
+            { id: '4', label: 'four' },
+          ],
+        });
+
+        act(() => {
+          getItemRoot('1').focus();
+        });
+        expect(getFocusedItemId()).to.equal('1');
+
+        fireEvent.keyDown(getItemRoot('1'), { key: 't', ctrlKey: true });
+        expect(getFocusedItemId()).to.equal('1');
+
+        fireEvent.keyDown(getItemRoot('1'), { key: 't', metaKey: true });
+        expect(getFocusedItemId()).to.equal('1');
+
+        fireEvent.keyDown(getItemRoot('1'), { key: 't', shiftKey: true });
+        expect(getFocusedItemId()).to.equal('1');
+      });
+
+      it('should work on disabled item when disabledItemsFocusable={true}', () => {
+        const { getFocusedItemId, getItemRoot } = render({
+          items: [
+            { id: '1', label: 'one', disabled: true },
+            { id: '2', label: 'two', disabled: true },
+            { id: '3', label: 'three', disabled: true },
+            { id: '4', label: 'four', disabled: true },
+          ],
+          disabledItemsFocusable: true,
+        });
+
+        act(() => {
+          getItemRoot('1').focus();
+        });
+        expect(getFocusedItemId()).to.equal('1');
+
+        fireEvent.keyDown(getItemRoot('1'), { key: 't' });
+        expect(getFocusedItemId()).to.equal('2');
+      });
+
+      it('should not move focus on disabled item when disabledItemsFocusable={false}', () => {
+        const { getFocusedItemId, getItemRoot } = render({
+          items: [
+            { id: '1', label: 'one', disabled: true },
+            { id: '2', label: 'two', disabled: true },
+            { id: '3', label: 'three', disabled: true },
+            { id: '4', label: 'four', disabled: true },
+          ],
+          disabledItemsFocusable: false,
+        });
+
+        act(() => {
+          getItemRoot('1').focus();
+        });
+        expect(getFocusedItemId()).to.equal('1');
+
+        fireEvent.keyDown(getItemRoot('1'), { key: 't' });
+        expect(getFocusedItemId()).to.equal('1');
+      });
+    });
+  },
+);
diff --git a/packages/x-tree-view/src/internals/plugins/useTreeViewKeyboardNavigation/useTreeViewKeyboardNavigation.ts b/packages/x-tree-view/src/internals/plugins/useTreeViewKeyboardNavigation/useTreeViewKeyboardNavigation.ts
index df6d4b7a2a335..75a7538bc1d1a 100644
--- a/packages/x-tree-view/src/internals/plugins/useTreeViewKeyboardNavigation/useTreeViewKeyboardNavigation.ts
+++ b/packages/x-tree-view/src/internals/plugins/useTreeViewKeyboardNavigation/useTreeViewKeyboardNavigation.ts
@@ -61,11 +61,13 @@ export const useTreeViewKeyboardNavigation: TreeViewPlugin<
 
     let matchingItemId: string | null = null;
     let currentItemId: string = getNextItem(itemId);
-    // The "currentItemId !== itemId" condition is to avoid an infinite loop when there is no matching item.
-    while (matchingItemId == null && currentItemId !== itemId) {
+    const checkedItems: Record<string, true> = {};
+    // The "!checkedItems[currentItemId]" condition avoids an infinite loop when there is no matching item.
+    while (matchingItemId == null && !checkedItems[currentItemId]) {
       if (firstCharMap.current[currentItemId] === cleanQuery) {
         matchingItemId = currentItemId;
       } else {
+        checkedItems[currentItemId] = true;
         currentItemId = getNextItem(currentItemId);
       }
     }
diff --git a/test/utils/tree-view/describeTreeView/describeTreeView.tsx b/test/utils/tree-view/describeTreeView/describeTreeView.tsx
index 6cfccabee847e..3eb291137dbeb 100644
--- a/test/utils/tree-view/describeTreeView/describeTreeView.tsx
+++ b/test/utils/tree-view/describeTreeView/describeTreeView.tsx
@@ -91,7 +91,17 @@ const innerDescribeTreeView = <TPlugins extends TreeViewAnyPluginSignature[]>(
                 'data-testid': ownerState.itemId,
               }) as any,
           }}
-          getItemLabel={(item) => item.label ?? item.id}
+          getItemLabel={(item) => {
+            if (item.label) {
+              if (typeof item.label !== 'string') {
+                throw new Error('Only use string labels when testing RichTreeView(Pro)');
+              }
+
+              return item.label;
+            }
+
+            return item.id;
+          }}
           isItemDisabled={(item) => !!item.disabled}
           {...other}
         />
diff --git a/test/utils/tree-view/describeTreeView/describeTreeView.types.ts b/test/utils/tree-view/describeTreeView/describeTreeView.types.ts
index ba15ccee9e59e..082f2fb3bd4d9 100644
--- a/test/utils/tree-view/describeTreeView/describeTreeView.types.ts
+++ b/test/utils/tree-view/describeTreeView/describeTreeView.types.ts
@@ -115,7 +115,7 @@ interface DescribeTreeViewTestRunnerParams<TPlugins extends TreeViewAnyPluginSig
 
 export interface DescribeTreeViewItem {
   id: string;
-  label?: string;
+  label?: React.ReactNode;
   disabled?: boolean;
   children?: readonly DescribeTreeViewItem[];
 }