From 598c6fe32e99345cd487de101f530ed8b94ccdd5 Mon Sep 17 00:00:00 2001
From: Flavien DELANGLE <flaviendelangle@gmail.com>
Date: Mon, 25 Sep 2023 11:55:40 +0200
Subject: [PATCH] [TreeView] Do not try to focus a collapsed node when
 re-focusing the TreeView (#10422)

---
 .../src/TreeView/TreeView.test.tsx            | 93 +++++++++++++++++++
 .../useTreeViewFocus/useTreeViewFocus.ts      | 21 ++++-
 .../useTreeViewFocus.types.ts                 |  3 +-
 3 files changed, 112 insertions(+), 5 deletions(-)

diff --git a/packages/x-tree-view/src/TreeView/TreeView.test.tsx b/packages/x-tree-view/src/TreeView/TreeView.test.tsx
index 01d9aa84b2a8..708efa0de146 100644
--- a/packages/x-tree-view/src/TreeView/TreeView.test.tsx
+++ b/packages/x-tree-view/src/TreeView/TreeView.test.tsx
@@ -442,6 +442,99 @@ describe('<TreeView />', () => {
     });
   });
 
+  describe('useTreeViewFocus', () => {
+    it('should focus the selected item when the tree is focused', () => {
+      const onNodeFocus = spy();
+
+      const { getByRole } = render(
+        <TreeView selected={'2'} onNodeFocus={onNodeFocus}>
+          <TreeItem nodeId="1" label="1" />
+          <TreeItem nodeId="2" label="2" />
+        </TreeView>,
+      );
+
+      act(() => {
+        getByRole('tree').focus();
+      });
+
+      expect(onNodeFocus.lastCall.lastArg).to.equal('2');
+    });
+
+    it('should focus the selected item when the tree is focused (multi select)', () => {
+      const onNodeFocus = spy();
+
+      const { getByRole } = render(
+        <TreeView multiSelect selected={['2']} onNodeFocus={onNodeFocus}>
+          <TreeItem nodeId="1" label="1" />
+          <TreeItem nodeId="2" label="2" />
+        </TreeView>,
+      );
+
+      act(() => {
+        getByRole('tree').focus();
+      });
+
+      expect(onNodeFocus.lastCall.lastArg).to.equal('2');
+    });
+
+    it('should focus the first visible selected item when the tree is focused (multi select)', () => {
+      const onNodeFocus = spy();
+
+      const { getByRole } = render(
+        <TreeView multiSelect selected={['1.1', '2']} onNodeFocus={onNodeFocus}>
+          <TreeItem nodeId="1" label="1">
+            <TreeItem nodeId="1.1" label="1.1" />
+          </TreeItem>
+          <TreeItem nodeId="2" label="2" />
+        </TreeView>,
+      );
+
+      act(() => {
+        getByRole('tree').focus();
+      });
+
+      expect(onNodeFocus.lastCall.lastArg).to.equal('2');
+    });
+
+    it('should focus the first item if the selected item is not visible', () => {
+      const onNodeFocus = spy();
+
+      const { getByRole } = render(
+        <TreeView selected="1.1" onNodeFocus={onNodeFocus}>
+          <TreeItem nodeId="1" label="1">
+            <TreeItem nodeId="1.1" label="1.1" />
+          </TreeItem>
+          <TreeItem nodeId="2" label="2" />
+        </TreeView>,
+      );
+
+      act(() => {
+        getByRole('tree').focus();
+      });
+
+      expect(onNodeFocus.lastCall.lastArg).to.equal('1');
+    });
+
+    it('should focus the first item if no selected item is visible (multi select)', () => {
+      const onNodeFocus = spy();
+
+      const { getByRole } = render(
+        <TreeView multiSelect selected={['1.1']} onNodeFocus={onNodeFocus}>
+          <TreeItem nodeId="1" label="1">
+            <TreeItem nodeId="1.1" label="1.1" />
+          </TreeItem>
+          <TreeItem nodeId="2" label="2" />
+        </TreeView>,
+      );
+
+      act(() => {
+        getByRole('tree').focus();
+      });
+
+      expect(onNodeFocus.lastCall.lastArg).to.equal('1');
+    });
+  });
+
   describe('Accessibility', () => {
     it('(TreeView) should have the role `tree`', () => {
       const { getByRole } = render(<TreeView />);
diff --git a/packages/x-tree-view/src/internals/plugins/useTreeViewFocus/useTreeViewFocus.ts b/packages/x-tree-view/src/internals/plugins/useTreeViewFocus/useTreeViewFocus.ts
index 7072a35c0dfa..6f12f844cfd3 100644
--- a/packages/x-tree-view/src/internals/plugins/useTreeViewFocus/useTreeViewFocus.ts
+++ b/packages/x-tree-view/src/internals/plugins/useTreeViewFocus/useTreeViewFocus.ts
@@ -58,10 +58,23 @@ export const useTreeViewFocus: TreeViewPlugin<UseTreeViewFocusSignature> = ({
 
       // if the event bubbled (which is React specific) we don't want to steal focus
       if (event.target === event.currentTarget) {
-        const firstSelected = Array.isArray(models.selected.value)
-          ? models.selected.value[0]
-          : models.selected.value;
-        instance.focusNode(event, firstSelected || instance.getNavigableChildrenIds(null)[0]);
+        const isNodeVisible = (nodeId: string) => {
+          const node = instance.getNode(nodeId);
+          return node && (node.parentId == null || instance.isNodeExpanded(node.parentId));
+        };
+
+        let nodeToFocusId: string | null | undefined;
+        if (Array.isArray(models.selected.value)) {
+          nodeToFocusId = models.selected.value.find(isNodeVisible);
+        } else if (models.selected.value != null && isNodeVisible(models.selected.value)) {
+          nodeToFocusId = models.selected.value;
+        }
+
+        if (nodeToFocusId == null) {
+          nodeToFocusId = instance.getNavigableChildrenIds(null)[0];
+        }
+
+        instance.focusNode(event, nodeToFocusId);
       }
     };
 
diff --git a/packages/x-tree-view/src/internals/plugins/useTreeViewFocus/useTreeViewFocus.types.ts b/packages/x-tree-view/src/internals/plugins/useTreeViewFocus/useTreeViewFocus.types.ts
index 629e28f0933f..b88ee531e2ef 100644
--- a/packages/x-tree-view/src/internals/plugins/useTreeViewFocus/useTreeViewFocus.types.ts
+++ b/packages/x-tree-view/src/internals/plugins/useTreeViewFocus/useTreeViewFocus.types.ts
@@ -2,6 +2,7 @@ import * as React from 'react';
 import { TreeViewPluginSignature } from '../../models';
 import type { UseTreeViewNodesSignature } from '../useTreeViewNodes';
 import type { UseTreeViewSelectionSignature } from '../useTreeViewSelection';
+import { UseTreeViewExpansionSignature } from '../useTreeViewExpansion';
 
 export interface UseTreeViewFocusInstance {
   isNodeFocused: (nodeId: string) => boolean;
@@ -31,5 +32,5 @@ export type UseTreeViewFocusSignature = TreeViewPluginSignature<
   {},
   UseTreeViewFocusState,
   never,
-  [UseTreeViewNodesSignature, UseTreeViewSelectionSignature<any>]
+  [UseTreeViewNodesSignature, UseTreeViewSelectionSignature<any>, UseTreeViewExpansionSignature]
 >;