From 1c847bdf75e28d62015c5b8e9d77c26e5b0995dc Mon Sep 17 00:00:00 2001
From: Stephanie Hong <41085564+JelloBagel@users.noreply.github.com>
Date: Wed, 24 Apr 2024 07:55:28 -0700
Subject: [PATCH] [TreeView] Check if the active element is a treeitem (#4484)

* Check if the active element is a treeitem

* changeset

* Update check

* Add focus story and return focus test

* Update rare-rings-argue.md

* patch, not minor

---------

Co-authored-by: Siddharth Kshetrapal <siddharthkp@github.com>
---
 .changeset/rare-rings-argue.md                |  5 ++
 .../TreeView/TreeView.features.stories.tsx    | 62 ++++++++++++++++++-
 packages/react/src/TreeView/TreeView.test.tsx | 34 ++++++++++
 .../react/src/TreeView/useRovingTabIndex.ts   |  6 +-
 4 files changed, 104 insertions(+), 3 deletions(-)
 create mode 100644 .changeset/rare-rings-argue.md

diff --git a/.changeset/rare-rings-argue.md b/.changeset/rare-rings-argue.md
new file mode 100644
index 00000000000..7f305e1c784
--- /dev/null
+++ b/.changeset/rare-rings-argue.md
@@ -0,0 +1,5 @@
+---
+'@primer/react': patch
+---
+
+TreeView: Fix returning focus when active element is not a treeitem
diff --git a/packages/react/src/TreeView/TreeView.features.stories.tsx b/packages/react/src/TreeView/TreeView.features.stories.tsx
index 4f73e1dcb7d..9470ad91730 100644
--- a/packages/react/src/TreeView/TreeView.features.stories.tsx
+++ b/packages/react/src/TreeView/TreeView.features.stories.tsx
@@ -1,8 +1,15 @@
-import {DiffAddedIcon, DiffModifiedIcon, DiffRemovedIcon, DiffRenamedIcon, FileIcon} from '@primer/octicons-react'
+import {
+  DiffAddedIcon,
+  DiffModifiedIcon,
+  DiffRemovedIcon,
+  DiffRenamedIcon,
+  FileIcon,
+  KebabHorizontalIcon,
+} from '@primer/octicons-react'
 import type {Meta, Story} from '@storybook/react'
 import React from 'react'
 import Box from '../Box'
-import {Button} from '../Button'
+import {Button, IconButton} from '../Button'
 import Octicon from '../Octicon'
 import type {SubTreeState} from './TreeView'
 import {TreeView} from './TreeView'
@@ -891,6 +898,57 @@ export const InitialFocus: Story = () => (
   </div>
 )
 
+export const FocusManagement: Story = () => (
+  <div>
+    <Button>Focusable element before TreeView</Button>
+    <TreeView aria-label="Test tree">
+      <TreeView.Item id="src" defaultExpanded>
+        <TreeView.LeadingVisual>
+          <TreeView.DirectoryIcon />
+        </TreeView.LeadingVisual>
+        src
+        <TreeView.SubTree>
+          <TreeView.Item id="src/Avatar.tsx">
+            <TreeView.LeadingVisual>
+              <FileIcon />
+            </TreeView.LeadingVisual>
+            Avatar.tsx
+            <IconButton
+              variant="invisible"
+              icon={KebabHorizontalIcon}
+              aria-label="Secondary actions"
+              aria-hidden
+              tabIndex={-1}
+            ></IconButton>
+          </TreeView.Item>
+          <TreeView.Item id="src/Button" defaultExpanded>
+            <TreeView.LeadingVisual>
+              <TreeView.DirectoryIcon />
+            </TreeView.LeadingVisual>
+            Button
+            <TreeView.SubTree>
+              <TreeView.Item id="src/Button/Button.tsx">
+                <TreeView.LeadingVisual>
+                  <FileIcon />
+                </TreeView.LeadingVisual>
+                Button.tsx
+                <IconButton
+                  variant="invisible"
+                  icon={KebabHorizontalIcon}
+                  aria-label="Secondary actions"
+                  aria-hidden
+                  tabIndex={-1}
+                ></IconButton>
+              </TreeView.Item>
+            </TreeView.SubTree>
+          </TreeView.Item>
+        </TreeView.SubTree>
+      </TreeView.Item>
+    </TreeView>
+    <Button>Focusable element after TreeView</Button>
+  </div>
+)
+
 export const WithoutIndentation: Story = () => (
   <nav aria-label="Files changed">
     <TreeView aria-label="Files changed" flat>
diff --git a/packages/react/src/TreeView/TreeView.test.tsx b/packages/react/src/TreeView/TreeView.test.tsx
index 40afaf1edd3..00f076da368 100644
--- a/packages/react/src/TreeView/TreeView.test.tsx
+++ b/packages/react/src/TreeView/TreeView.test.tsx
@@ -285,6 +285,40 @@ describe('Markup', () => {
     const subItem1 = getByRole('treeitem', {name: /SubItem 1/})
     expect(subItem1).toBeInTheDocument()
   })
+
+  it("should move focus to first treeitem when focusing back in after clicking on a treeitem's secondary action", async () => {
+    const user = userEvent.setup({delay: null})
+    const {getByRole, getByText} = renderWithTheme(
+      <div>
+        <TreeView aria-label="Test tree">
+          <TreeView.Item id="item-1">Item 1</TreeView.Item>
+          <TreeView.Item id="item-2">
+            Item 2
+            <button id="item-2-button" tabIndex={-1} aria-hidden>
+              Link in Item 2
+            </button>
+          </TreeView.Item>
+          <TreeView.Item id="item-3">Item 3</TreeView.Item>
+        </TreeView>
+        <button>Focusable element</button>
+      </div>,
+    )
+
+    // Click on treeitem's secondary action
+    const item2Button = getByText(/Link in Item 2/i)
+    await user.click(item2Button)
+    expect(item2Button).toHaveFocus()
+
+    // Move focus to button outside of TreeView
+    await user.tab()
+    const outerButton = getByRole('button', {name: /Focusable element/})
+    expect(outerButton).toHaveFocus()
+
+    // Move focus into TreeView. Focus should be on first treeitem
+    await user.tab({shift: true})
+    const item1 = getByRole('treeitem', {name: /Item 1/})
+    expect(item1).toHaveFocus()
+  })
 })
 
 describe('Keyboard interactions', () => {
diff --git a/packages/react/src/TreeView/useRovingTabIndex.ts b/packages/react/src/TreeView/useRovingTabIndex.ts
index 74e569ba375..0c31272b2af 100644
--- a/packages/react/src/TreeView/useRovingTabIndex.ts
+++ b/packages/react/src/TreeView/useRovingTabIndex.ts
@@ -52,7 +52,11 @@ export function useRovingTabIndex({
       }
 
       // Otherwise, focus the activeElement if it's a treeitem
-      if (document.activeElement instanceof HTMLElement && containerRef.current?.contains(document.activeElement)) {
+      if (
+        document.activeElement instanceof HTMLElement &&
+        containerRef.current?.contains(document.activeElement) &&
+        document.activeElement.getAttribute('role') === 'treeitem'
+      ) {
         return document.activeElement
       }