diff --git a/src/Dom/focus.ts b/src/Dom/focus.ts index f6a6af98..556d79ef 100644 --- a/src/Dom/focus.ts +++ b/src/Dom/focus.ts @@ -1,4 +1,5 @@ -import { useEffect } from 'react'; +import { useEffect, useRef, useState } from 'react'; +import type { DependencyList } from 'react'; import isVisible from './isVisible'; import useId from '../hooks/useId'; @@ -211,6 +212,41 @@ export function lockFocus(element: HTMLElement, id: string): VoidFunction { }; } +/** + * Retry an effect until it reports ready. + * When `ready` is `false`, it will schedule one more effect cycle and call `func` again + * with the next `retryTimes`. + */ +type RetryEffectResult = readonly [ + clearFunc: VoidFunction | undefined, + ready: boolean, +]; + +function useRetryEffect( + func: (retryTimes: number) => RetryEffectResult, + deps: DependencyList, +): void { + /* eslint-disable react-hooks/exhaustive-deps */ + const retryTimesRef = useRef(0); + const [retryMark, setRetryMark] = useState(0); + + useEffect(() => { + retryTimesRef.current = 0; + }, deps); + + useEffect(() => { + const [clearFn, ready] = func(retryTimesRef.current); + + if (!ready) { + retryTimesRef.current += 1; + setRetryMark(count => count + 1); + } + + return clearFn; + }, [...deps, retryMark]); + /* eslint-enable react-hooks/exhaustive-deps */ +} + /** * Lock focus within an element. * When locked, focus will be restricted to focusable elements within the specified element. @@ -222,15 +258,24 @@ export function useLockFocus( getElement: () => HTMLElement | null, ): [ignoreElement: (ele: HTMLElement) => void] { const id = useId(); + const getElementRef = useRef(getElement); - useEffect(() => { - if (lock) { - const element = getElement(); - if (element) { - return lockFocus(element, id); - } + getElementRef.current = getElement; + + const lockEffect = (retryTimes: number): RetryEffectResult => { + if (!lock) { + return [undefined, true]; } - }, [lock, id]); + + const element = getElementRef.current(); + if (element) { + return [lockFocus(element, id), true]; + } + + return [undefined, retryTimes >= 1]; + }; + + useRetryEffect(lockEffect, [id, lock]); const ignoreElement = (ele: HTMLElement) => { if (ele) { diff --git a/tests/focus.test.tsx b/tests/focus.test.tsx index 78378425..80142537 100644 --- a/tests/focus.test.tsx +++ b/tests/focus.test.tsx @@ -1,6 +1,6 @@ /* eslint-disable class-methods-use-this */ -import React, { useRef } from 'react'; -import { render } from '@testing-library/react'; +import React, { useEffect, useRef, useState } from 'react'; +import { render, waitFor } from '@testing-library/react'; import { spyElementPrototype } from '../src/test/domHook'; import { getFocusNodeList, triggerFocus, useLockFocus } from '../src/Dom/focus'; @@ -95,6 +95,39 @@ describe('focus', () => { outerButton.focus(); expect(document.activeElement).toBe(input1); }); + + it('should retry lock once when element is filled after lock starts', async () => { + const DelayedElementComponent: React.FC = () => { + const elementRef = useRef(null); + const [element, setElement] = useState(null); + + useLockFocus(true, () => element); + + useEffect(() => { + setElement(elementRef.current); + }, []); + + return ( + <> + +
+ + +
+ + ); + }; + + const { getByTestId } = render(); + + const focusContainer = getByTestId('focus-container'); + + await waitFor(() => { + expect(document.activeElement).toBe(focusContainer); + }); + }); }); it('ignoreElement should allow focus on ignored elements', () => {