Make sure useFocusOnMount runs when all the children tabbable elements have mounted

packages/components/src/popover/test/index.tsx

@@ -138,9 +138,12 @@ describe( 'Popover', () => {
 		} );
 
 		describe( 'focus behavior', () => {
-			it( 'should focus the popover by default when opened', async () => {
+			it( 'should focus the popover container when opened', async () => {
 				render(
-					<Popover data-testid="popover-element">
+					<Popover
+						focusOnMount={ true }
+						data-testid="popover-element"
+					>
 						Popover content
 					</Popover>
 				);

packages/compose/src/hooks/use-focus-on-mount/index.js

@@ -27,10 +27,37 @@ import { focus } from '@wordpress/dom';
  */
 export default function useFocusOnMount( focusOnMount = 'firstElement' ) {
 	const focusOnMountRef = useRef( focusOnMount );
+
+	/**
+	 * Sets focus on a DOM element.
+	 *
+	 * @param {HTMLElement} target The DOM element to set focus to.
+	 * @return {void}
+	 */
+	const setFocus = ( target ) => {
+		target.focus( {
+			// When focusing newly mounted dialogs,
+			// the position of the popover is often not right on the first render
+			// This prevents the layout shifts when focusing the dialogs.
+			preventScroll: true,
+		} );
+	};
+
+	/** @type {import('react').MutableRefObject<ReturnType<setTimeout> | undefined>} */
+	const timerId = useRef();
+
 	useEffect( () => {
 		focusOnMountRef.current = focusOnMount;
 	}, [ focusOnMount ] );
 
+	useEffect( () => {
+		return () => {
+			if ( timerId.current ) {
+				clearTimeout( timerId.current );
+			}
+		};
+	}, [] );
+
 	return useCallback( ( node ) => {
 		if ( ! node || focusOnMountRef.current === false ) {
 			return;
@@ -40,21 +67,18 @@ export default function useFocusOnMount( focusOnMount = 'firstElement' ) {
 			return;
 		}
 
-		let target = node;
-
 		if ( focusOnMountRef.current === 'firstElement' ) {
-			const firstTabbable = focus.tabbable.find( node )[ 0 ];
+			timerId.current = setTimeout( () => {
+				const firstTabbable = focus.tabbable.find( node )[ 0 ];
 
-			if ( firstTabbable ) {
-				target = /** @type {HTMLElement} */ ( firstTabbable );
-			}
+				if ( firstTabbable ) {
+					setFocus( /** @type {HTMLElement} */ ( firstTabbable ) );
+				}
+			}, 0 );
+
+			return;
 		}
 
-		target.focus( {
-			// When focusing newly mounted dialogs,
-			// the position of the popover is often not right on the first render
-			// This prevents the layout shifts when focusing the dialogs.
-			preventScroll: true,
-		} );
+		setFocus( node );
 	}, [] );
 }

packages/e2e-tests/specs/editor/various/rich-text.test.js

@@ -412,13 +412,13 @@ describe( 'RichText', () => {
 		await page.waitForXPath(
 			'//button[@role="tab"][@aria-selected="true"][text()="Text"]'
 		);
-		// Tab to the "Text" tab.
+		// Initial focus is on the "Text" tab.
+		// Tab to the "Custom color picker".
 		await page.keyboard.press( 'Tab' );
 		// Tab to black.
 		await page.keyboard.press( 'Tab' );
 		// Select color other than black.
 		await page.keyboard.press( 'Tab' );
-		await page.keyboard.press( 'Tab' );
 		await page.keyboard.press( 'Enter' );
 
 		expect( await getEditedPostContent() ).toMatchSnapshot();