TreeView: Refactor loading focus management (#2868)

* Refactor loading focus management

Co-authored-by: Josh Black <joshblack@users.noreply.github.com>

* Update spacing

* Revert loading item id

* Update src/TreeView/TreeView.tsx

* Update src/TreeView/TreeView.tsx

* Remove children from dep array

* Create .changeset/chilly-rabbits-invent.md

* Update TreeView tests

---------

Co-authored-by: Josh Black <joshblack@users.noreply.github.com>

Author: colebemis

.changeset/chilly-rabbits-invent.md

@@ -0,0 +1,5 @@
+---
+"@primer/react": patch
+---
+
+TreeView: Reliably move focus from loading item to first asynchronously loaded item

src/TreeView/TreeView.test.tsx

@@ -1230,6 +1230,11 @@ describe('Asyncronous loading', () => {
   })
 
   it('moves focus from loading item to first child', async () => {
+    // We get a focus zone warning in this test that doesn't
+    // happen in the browser. We're not sure why, so we're
+    // suppressing it for now.
+    jest.spyOn(console, 'warn').mockImplementation()
+
     function TestTree() {
       const [state, setState] = React.useState<SubTreeState>('loading')
 
@@ -1259,21 +1264,21 @@ describe('Asyncronous loading', () => {
     act(() => {
       // Focus first item
       parentItem.focus()
-
-      // Press ↓ to move focus to loading item
-      fireEvent.keyDown(document.activeElement || document.body, {key: 'ArrowDown'})
     })
 
+    // Press ↓ to move focus to loading item
+    fireEvent.keyDown(document.activeElement || document.body, {key: 'ArrowDown'})
+
     // Loading item should be focused
     expect(loadingItem).toHaveFocus()
 
-    // Wait for async loading to complete
-    const firstChild = await findByRole('treeitem', {name: 'Child 1'})
-
     act(() => {
       jest.runAllTimers()
     })
 
+    // Wait for async loading to complete
+    const firstChild = await findByRole('treeitem', {name: 'Child 1'})
+
     // First child should be focused
     expect(firstChild).toHaveFocus()
   })

src/TreeView/TreeView.tsx

@@ -10,7 +10,6 @@ import styled, {keyframes} from 'styled-components'
 import {get} from '../constants'
 import {ConfirmationDialog} from '../Dialog/ConfirmationDialog'
 import {useControllableState} from '../hooks/useControllableState'
-import useSafeTimeout from '../hooks/useSafeTimeout'
 import {useId} from '../hooks/useId'
 import Spinner from '../Spinner'
 import sx, {SxProp} from '../sx'
@@ -510,17 +509,10 @@ export type TreeViewSubTreeProps = {
 const SubTree: React.FC<TreeViewSubTreeProps> = ({count, state, children}) => {
   const {announceUpdate} = React.useContext(RootContext)
   const {itemId, isExpanded, isSubTreeEmpty, setIsSubTreeEmpty} = React.useContext(ItemContext)
-  const [isLoadingItemVisible, setIsLoadingItemVisible] = React.useState(false)
-  const {safeSetTimeout} = useSafeTimeout()
   const loadingItemRef = React.useRef<HTMLElement>(null)
   const ref = React.useRef<HTMLElement>(null)
-  const [isPending, setPending] = React.useState(state === 'loading')
-
-  React.useEffect(() => {
-    if (state === 'loading') {
-      setPending(true)
-    }
-  }, [state])
+  const [loadingFocused, setLoadingFocused] = React.useState(false)
+  const previousState = usePreviousValue(state)
 
   React.useEffect(() => {
     // If `state` is undefined, we're working in a synchronous context and need
@@ -536,61 +528,66 @@ const SubTree: React.FC<TreeViewSubTreeProps> = ({count, state, children}) => {
     }
   }, [state, isSubTreeEmpty, setIsSubTreeEmpty, children])
 
-  // If a consumer sets state="done" without having a previous state (like `loading`),
-  // then it would announce on the first render. Using isPending is to only
-  // announce being "loaded" when the state has changed from `loading` --> `done`.
+  // Handle transition from loading to done state
   React.useEffect(() => {
-    if (isPending && state === 'done') {
-      const parentItem = document.getElementById(itemId)
+    if (previousState === 'loading' && state === 'done') {
+      const parentElement = document.getElementById(itemId)
+      if (!parentElement) return
+
+      // Announce update to screen readers
+      const parentName = getAccessibleName(parentElement)
 
-      if (!parentItem) return
+      if (ref.current?.childElementCount) {
+        announceUpdate(`${parentName} content loaded`)
+      } else {
+        announceUpdate(`${parentName} is empty`)
+      }
 
-      const {current: node} = ref
-      const parentName = getAccessibleName(parentItem)
+      // Move focus to the first child if the loading indicator
+      // was focused when the async items finished loading
+      if (loadingFocused) {
+        const firstChild = getFirstChildElement(parentElement)
 
-      safeSetTimeout(() => {
-        if (node && node.childElementCount > 0) {
-          announceUpdate(`${parentName} content loaded`)
+        if (firstChild) {
+          firstChild.focus()
         } else {
-          announceUpdate(`${parentName} is empty`)
+          parentElement.focus()
         }
-      })
 
-      setPending(false)
+        setLoadingFocused(false)
+      }
     }
-  }, [state, itemId, announceUpdate, safeSetTimeout, isPending])
+  }, [loadingFocused, previousState, state, itemId, announceUpdate, ref])
 
-  // Manage loading indicator state
+  // Track focus on the loading indicator
   React.useEffect(() => {
-    // If we're in the loading state, but not showing the loading indicator yet,
-    // show the loading indicator
-    if (state === 'loading' && !isLoadingItemVisible) {
-      setIsLoadingItemVisible(true)
+    function handleFocus() {
+      setLoadingFocused(true)
     }
 
-    // If we're not in the loading state, but we're still showing a loading indicator,
-    // hide the loading indicator and move focus if necessary
-    if (state !== 'loading' && isLoadingItemVisible) {
-      const isLoadingItemFocused = document.activeElement === loadingItemRef.current
+    function handleBlur(event: FocusEvent) {
+      // Skip blur events that are caused by the element being removed from the DOM.
+      // This can happen when the loading indicator is focused when async items are
+      // done loading and the loading indicator is removed from the DOM.
+      // If `loadingFocused` is `true` when `state` is `"done"` then the loading indicator
+      // was focused when the async items finished loading and we need to move focus to the
+      // first child.
+      if (!event.relatedTarget) return
 
-      setIsLoadingItemVisible(false)
+      setLoadingFocused(false)
+    }
 
-      if (isLoadingItemFocused) {
-        safeSetTimeout(() => {
-          const parentElement = document.getElementById(itemId)
-          if (!parentElement) return
+    const loadingElement = loadingItemRef.current
+    if (!loadingElement) return
 
-          const firstChild = getFirstChildElement(parentElement)
+    loadingElement.addEventListener('focus', handleFocus)
+    loadingElement.addEventListener('blur', handleBlur)
 
-          if (firstChild) {
-            firstChild.focus()
-          } else {
-            parentElement.focus()
-          }
-        })
-      }
+    return () => {
+      loadingElement.removeEventListener('focus', handleFocus)
+      loadingElement.removeEventListener('blur', handleBlur)
     }
-  }, [state, safeSetTimeout, isLoadingItemVisible, itemId])
+  }, [loadingItemRef, state])
 
   if (!isExpanded) {
     return null
@@ -607,13 +604,23 @@ const SubTree: React.FC<TreeViewSubTreeProps> = ({count, state, children}) => {
       // @ts-ignore Box doesn't have type support for `ref` used in combination with `as`
       ref={ref}
     >
-      {isLoadingItemVisible ? <LoadingItem ref={loadingItemRef} count={count} /> : children}
+      {state === 'loading' ? <LoadingItem ref={loadingItemRef} count={count} /> : children}
     </ul>
   )
 }
 
 SubTree.displayName = 'TreeView.SubTree'
 
+function usePreviousValue<T>(value: T): T {
+  const ref = React.useRef(value)
+
+  React.useEffect(() => {
+    ref.current = value
+  }, [value])
+
+  return ref.current
+}
+
 const shimmer = keyframes`
   from { mask-position: 200%; }
   to { mask-position: 0%; }