Refactor to make focusOnMount=true work as expected.

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,25 +67,18 @@ export default function useFocusOnMount( focusOnMount = 'firstElement' ) {
 			return;
 		}
 
-		if ( focusOnMountRef.current !== 'firstElement' ) {
-			node.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,
-			} );
+		if ( focusOnMountRef.current === 'firstElement' ) {
+			timerId.current = setTimeout( () => {
+				const firstTabbable = focus.tabbable.find( node )[ 0 ];
+
+				if ( firstTabbable ) {
+					setFocus( /** @type {HTMLElement} */ ( firstTabbable ) );
+				}
+			}, 0 );
 
 			return;
 		}
 
-		const focusTimeout = setTimeout( () => {
-			const firstTabbable = focus.tabbable.find( node )[ 0 ];
-
-			if ( firstTabbable ) {
-				/** @type {HTMLElement} */ ( firstTabbable ).focus();
-			}
-		}, 0 );
-
-		return () => clearTimeout( focusTimeout );
+		setFocus( node );
 	}, [] );
 }