Address focus trap review feedback

Copilot-Session: b43bee82-7a2d-40d0-95c2-9b053bdc6075

Author: danielguillan

packages/react/src/hooks/useFocusTrap.test.tsx

@@ -33,8 +33,8 @@ describe('useFocusTrap', () => {
   })
 
   it('preserves provided refs', () => {
-    const containerRef = React.createRef<HTMLElement>()
-    const initialFocusRef = React.createRef<HTMLElement>()
+    const containerRef = React.createRef<HTMLDivElement>()
+    const initialFocusRef = React.createRef<HTMLButtonElement>()
     const {result} = renderHook(() => useFocusTrap({containerRef, initialFocusRef, disabled: true}))
 
     expect(result.current.containerRef).toBe(containerRef)
@@ -43,11 +43,11 @@ describe('useFocusTrap', () => {
 
   it('passes generated ref elements to the focus trap', () => {
     const TestComponent = () => {
-      const {containerRef, initialFocusRef} = useFocusTrap()
+      const {containerRef, initialFocusRef} = useFocusTrap<HTMLDivElement, HTMLButtonElement>()
 
       return (
-        <div ref={containerRef as React.RefObject<HTMLDivElement | null>}>
-          <button ref={initialFocusRef as React.RefObject<HTMLButtonElement | null>} />
+        <div ref={containerRef}>
+          <button ref={initialFocusRef} />
         </div>
       )
     }
@@ -56,4 +56,57 @@ describe('useFocusTrap', () => {
 
     expect(mockFocusTrap).toHaveBeenCalledWith(container.querySelector('div'), container.querySelector('button'))
   })
+
+  it('starts and aborts the focus trap when disabled changes', () => {
+    const abortController = new AbortController()
+    const abortSpy = jest.spyOn(abortController, 'abort')
+    mockFocusTrap.mockReturnValue(abortController)
+
+    const TestComponent = ({disabled}: {disabled: boolean}) => {
+      const {containerRef} = useFocusTrap<HTMLDivElement>({disabled})
+      return <div ref={containerRef} />
+    }
+
+    const {rerender} = render(<TestComponent disabled />)
+
+    expect(mockFocusTrap).not.toHaveBeenCalled()
+
+    rerender(<TestComponent disabled={false} />)
+
+    expect(mockFocusTrap).toHaveBeenCalledTimes(1)
+
+    rerender(<TestComponent disabled />)
+
+    expect(abortSpy).toHaveBeenCalled()
+  })
+
+  it('captures fresh focus after cleaning up a non-HTMLElement active element', () => {
+    const svg = document.createElementNS('http://www.w3.org/2000/svg', 'svg')
+    svg.setAttribute('tabindex', '0')
+    document.body.append(svg)
+    svg.focus()
+
+    const nextFocusedElement = document.createElement('button')
+    document.body.append(nextFocusedElement)
+
+    const TestComponent = ({disabled}: {disabled: boolean}) => {
+      const {containerRef} = useFocusTrap<HTMLDivElement>({disabled, restoreFocusOnCleanUp: true})
+      return <div ref={containerRef} />
+    }
+
+    const {rerender} = render(<TestComponent disabled={false} />)
+
+    rerender(<TestComponent disabled />)
+    nextFocusedElement.focus()
+
+    const focusSpy = jest.spyOn(nextFocusedElement, 'focus')
+
+    rerender(<TestComponent disabled={false} />)
+    rerender(<TestComponent disabled />)
+
+    expect(focusSpy).toHaveBeenCalled()
+
+    svg.remove()
+    nextFocusedElement.remove()
+  })
 })

packages/react/src/hooks/useFocusTrap.ts

@@ -2,18 +2,21 @@ import React from 'react'
 import {focusTrap} from '@primer/behaviors'
 import {useProvidedRefOrCreate} from './useRef'
 
-export interface FocusTrapHookSettings {
+export interface FocusTrapHookSettings<
+  ContainerElement extends HTMLElement = HTMLElement,
+  InitialFocusElement extends HTMLElement = HTMLElement,
+> {
   /**
    * Ref object that will be used for the trapping container. If not provided, one will
    * be created by this hook and returned.
    */
-  containerRef?: React.RefObject<HTMLElement | null>
+  containerRef?: React.RefObject<ContainerElement | null>
 
   /**
    * Ref object for the element that should receive focus when the focus trap is first enabled.
    * If not provided, one will be created by this hook and returned. Its use is optional.
    */
-  initialFocusRef?: React.RefObject<HTMLElement | null>
+  initialFocusRef?: React.RefObject<InitialFocusElement | null>
 
   /**
    * Set to true to disable the focus trap and clean up listeners. Can be re-enabled at any time.
@@ -31,15 +34,18 @@ export interface FocusTrapHookSettings {
  * that should trap focus.
  * @param settings {FocusTrapHookSettings}
  */
-export function useFocusTrap(
-  settings?: FocusTrapHookSettings,
+export function useFocusTrap<
+  ContainerElement extends HTMLElement = HTMLElement,
+  InitialFocusElement extends HTMLElement = HTMLElement,
+>(
+  settings?: FocusTrapHookSettings<ContainerElement, InitialFocusElement>,
   dependencies: React.DependencyList = [],
 ): {
-  containerRef: React.RefObject<HTMLElement | null>
-  initialFocusRef: React.RefObject<HTMLElement | null>
+  containerRef: React.RefObject<ContainerElement | null>
+  initialFocusRef: React.RefObject<InitialFocusElement | null>
 } {
-  const containerRef = useProvidedRefOrCreate<HTMLElement | null>(settings?.containerRef)
-  const initialFocusRef = useProvidedRefOrCreate<HTMLElement | null>(settings?.initialFocusRef)
+  const containerRef = useProvidedRefOrCreate<ContainerElement | null>(settings?.containerRef)
+  const initialFocusRef = useProvidedRefOrCreate<InitialFocusElement | null>(settings?.initialFocusRef)
   const disabled = settings?.disabled
   const abortController = React.useRef<AbortController | null>(null)
   const previousFocusedElement = React.useRef<Element | null>(null)
@@ -50,8 +56,8 @@ export function useFocusTrap(
     abortController.current?.abort()
     if (settings?.restoreFocusOnCleanUp && previousFocusedElement.current instanceof HTMLElement) {
       previousFocusedElement.current.focus()
-      previousFocusedElement.current = null
     }
+    previousFocusedElement.current = null
   }
 
   React.useEffect(