Fix focus timing issues (#2911)

Author: colebemis

src/TreeView/TreeView.test.tsx

@@ -1230,11 +1230,6 @@ 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')
 
@@ -1272,13 +1267,13 @@ describe('Asyncronous loading', () => {
     // 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

@@ -11,6 +11,7 @@ import {get} from '../constants'
 import {ConfirmationDialog} from '../Dialog/ConfirmationDialog'
 import {useControllableState} from '../hooks/useControllableState'
 import {useId} from '../hooks/useId'
+import useSafeTimeout from '../hooks/useSafeTimeout'
 import Spinner from '../Spinner'
 import sx, {SxProp} from '../sx'
 import Text from '../Text'
@@ -513,6 +514,7 @@ const SubTree: React.FC<TreeViewSubTreeProps> = ({count, state, children}) => {
   const ref = React.useRef<HTMLElement>(null)
   const [loadingFocused, setLoadingFocused] = React.useState(false)
   const previousState = usePreviousValue(state)
+  const {safeSetTimeout} = useSafeTimeout()
 
   React.useEffect(() => {
     // If `state` is undefined, we're working in a synchronous context and need
@@ -549,15 +551,19 @@ const SubTree: React.FC<TreeViewSubTreeProps> = ({count, state, children}) => {
         const firstChild = getFirstChildElement(parentElement)
 
         if (firstChild) {
-          firstChild.focus()
+          safeSetTimeout(() => {
+            firstChild.focus()
+          })
         } else {
-          parentElement.focus()
+          safeSetTimeout(() => {
+            parentElement.focus()
+          })
         }
 
         setLoadingFocused(false)
       }
     }
-  }, [loadingFocused, previousState, state, itemId, announceUpdate, ref])
+  }, [loadingFocused, previousState, state, itemId, announceUpdate, ref, safeSetTimeout])
 
   // Track focus on the loading indicator
   React.useEffect(() => {