Use async act

Author: eps1lon

src/act-compat.js

@@ -33,55 +33,22 @@ function getIsReactActEnvironment() {
   return getGlobalThis().IS_REACT_ACT_ENVIRONMENT
 }
 
-function withGlobalActEnvironment(actImplementation) {
-  return callback => {
-    const previousActEnvironment = getIsReactActEnvironment()
-    setIsReactActEnvironment(true)
-    try {
-      // The return value of `act` is always a thenable.
-      let callbackNeedsToBeAwaited = false
-      const actResult = actImplementation(() => {
-        const result = callback()
-        if (
-          result !== null &&
-          typeof result === 'object' &&
-          typeof result.then === 'function'
-        ) {
-          callbackNeedsToBeAwaited = true
-        }
-        return result
-      })
-      if (callbackNeedsToBeAwaited) {
-        const thenable = actResult
-        return {
-          then: (resolve, reject) => {
-            thenable.then(
-              returnValue => {
-                setIsReactActEnvironment(previousActEnvironment)
-                resolve(returnValue)
-              },
-              error => {
-                setIsReactActEnvironment(previousActEnvironment)
-                reject(error)
-              },
-            )
-          },
-        }
-      } else {
-        setIsReactActEnvironment(previousActEnvironment)
-        return actResult
-      }
-    } catch (error) {
-      // Can't be a `finally {}` block since we don't know if we have to immediately restore IS_REACT_ACT_ENVIRONMENT
-      // or if we have to await the callback first.
-      setIsReactActEnvironment(previousActEnvironment)
-      throw error
-    }
+async function act(scope) {
+  const previousActEnvironment = getIsReactActEnvironment()
+  setIsReactActEnvironment(true)
+  try {
+    // scope passed to domAct needs to be `async` until React.act treats every scope as async.
+    // We already enforce `await act()` (regardless of scope) to flush microtasks
+    // inside the act scope.
+    const result = await reactAct(async () => {
+      return scope()
+    })
+    return result
+  } finally {
+    setIsReactActEnvironment(previousActEnvironment)
   }
 }
 
-const act = withGlobalActEnvironment(reactAct)
-
 export default act
 export {
   setIsReactActEnvironment as setReactActEnvironment,

src/fire-event.js

@@ -15,29 +15,29 @@ Object.keys(dtlFireEvent).forEach(key => {
 // @link https://github.com/facebook/react/blob/b87aabdfe1b7461e7331abb3601d9e6bb27544bc/packages/react-dom/src/events/EnterLeaveEventPlugin.js#L24-L31
 const mouseEnter = fireEvent.mouseEnter
 const mouseLeave = fireEvent.mouseLeave
-fireEvent.mouseEnter = (...args) => {
-  mouseEnter(...args)
+fireEvent.mouseEnter = async (...args) => {
+  await mouseEnter(...args)
   return fireEvent.mouseOver(...args)
 }
-fireEvent.mouseLeave = (...args) => {
-  mouseLeave(...args)
+fireEvent.mouseLeave = async (...args) => {
+  await mouseLeave(...args)
   return fireEvent.mouseOut(...args)
 }
 
 const pointerEnter = fireEvent.pointerEnter
 const pointerLeave = fireEvent.pointerLeave
-fireEvent.pointerEnter = (...args) => {
-  pointerEnter(...args)
+fireEvent.pointerEnter = async (...args) => {
+  await pointerEnter(...args)
   return fireEvent.pointerOver(...args)
 }
-fireEvent.pointerLeave = (...args) => {
-  pointerLeave(...args)
+fireEvent.pointerLeave = async (...args) => {
+  await pointerLeave(...args)
   return fireEvent.pointerOut(...args)
 }
 
 const select = fireEvent.select
-fireEvent.select = (node, init) => {
-  select(node, init)
+fireEvent.select = async (node, init) => {
+  await select(node, init)
   // React tracks this event only on focused inputs
   node.focus()
 
@@ -49,20 +49,20 @@ fireEvent.select = (node, init) => {
   // - keyDown
   // so we can use any here
   // @link https://github.com/facebook/react/blob/b87aabdfe1b7461e7331abb3601d9e6bb27544bc/packages/react-dom/src/events/SelectEventPlugin.js#L203-L224
-  fireEvent.keyUp(node, init)
+  await fireEvent.keyUp(node, init)
 }
 
 // React event system tracks native focusout/focusin events for
 // running blur/focus handlers
 // @link https://github.com/facebook/react/pull/19186
 const blur = fireEvent.blur
 const focus = fireEvent.focus
-fireEvent.blur = (...args) => {
-  fireEvent.focusOut(...args)
+fireEvent.blur = async (...args) => {
+  await fireEvent.focusOut(...args)
   return blur(...args)
 }
-fireEvent.focus = (...args) => {
-  fireEvent.focusIn(...args)
+fireEvent.focus = async (...args) => {
+  await fireEvent.focusIn(...args)
   return focus(...args)
 }
 

src/index.js

@@ -10,15 +10,15 @@ if (typeof process === 'undefined' || !process.env?.RTL_SKIP_AUTO_CLEANUP) {
   // ignore teardown() in code coverage because Jest does not support it
   /* istanbul ignore else */
   if (typeof afterEach === 'function') {
-    afterEach(() => {
-      cleanup()
+    afterEach(async () => {
+      await cleanup()
     })
   } else if (typeof teardown === 'function') {
     // Block is guarded by `typeof` check.
     // eslint does not support `typeof` guards.
     // eslint-disable-next-line no-undef
-    teardown(() => {
-      cleanup()
+    teardown(async () => {
+      await cleanup()
     })
   }
 

src/pure.js

@@ -58,11 +58,7 @@ configureDTL({
     }
   },
   eventWrapper: cb => {
-    let result
-    act(() => {
-      result = cb()
-    })
-    return result
+    return act(cb)
   },
 })
 
@@ -89,13 +85,13 @@ function wrapUiIfNeeded(innerElement, wrapperComponent) {
     : innerElement
 }
 
-function createConcurrentRoot(
+async function createConcurrentRoot(
   container,
   {hydrate, ui, wrapper: WrapperComponent},
 ) {
   let root
   if (hydrate) {
-    act(() => {
+    await act(() => {
       root = ReactDOMClient.hydrateRoot(
         container,
         strictModeIfNeeded(wrapUiIfNeeded(ui, WrapperComponent)),
@@ -116,45 +112,53 @@ function createConcurrentRoot(
       // Nothing to do since hydration happens when creating the root object.
     },
     render(element) {
-      root.render(element)
+      return act(() => {
+        root.render(element)
+      })
     },
     unmount() {
-      root.unmount()
+      return act(() => {
+        root.unmount()
+      })
     },
   }
 }
 
-function createLegacyRoot(container) {
+async function createLegacyRoot(container) {
   return {
     hydrate(element) {
-      ReactDOM.hydrate(element, container)
+      return act(() => {
+        ReactDOM.hydrate(element, container)
+      })
     },
     render(element) {
-      ReactDOM.render(element, container)
+      return act(() => {
+        ReactDOM.render(element, container)
+      })
     },
     unmount() {
-      ReactDOM.unmountComponentAtNode(container)
+      return act(() => {
+        ReactDOM.unmountComponentAtNode(container)
+      })
     },
   }
 }
 
-function renderRoot(
+async function renderRoot(
   ui,
   {baseElement, container, hydrate, queries, root, wrapper: WrapperComponent},
 ) {
-  act(() => {
-    if (hydrate) {
-      root.hydrate(
-        strictModeIfNeeded(wrapUiIfNeeded(ui, WrapperComponent)),
-        container,
-      )
-    } else {
-      root.render(
-        strictModeIfNeeded(wrapUiIfNeeded(ui, WrapperComponent)),
-        container,
-      )
-    }
-  })
+  if (hydrate) {
+    await root.hydrate(
+      strictModeIfNeeded(wrapUiIfNeeded(ui, WrapperComponent)),
+      container,
+    )
+  } else {
+    await root.render(
+      strictModeIfNeeded(wrapUiIfNeeded(ui, WrapperComponent)),
+      container,
+    )
+  }
 
   return {
     container,
@@ -166,12 +170,10 @@ function renderRoot(
         : // eslint-disable-next-line no-console,
           console.log(prettyDOM(el, maxLength, options)),
     unmount: () => {
-      act(() => {
-        root.unmount()
-      })
+      return root.unmount()
     },
-    rerender: rerenderUi => {
-      renderRoot(rerenderUi, {
+    rerender: async rerenderUi => {
+      await renderRoot(rerenderUi, {
         container,
         baseElement,
         root,
@@ -196,7 +198,7 @@ function renderRoot(
   }
 }
 
-function render(
+async function render(
   ui,
   {
     container,
@@ -230,7 +232,7 @@ function render(
   // eslint-disable-next-line no-negated-condition -- we want to map the evolution of this over time. The root is created first. Only later is it re-used so we don't want to read the case that happens later first.
   if (!mountedContainers.has(container)) {
     const createRootImpl = legacyRoot ? createLegacyRoot : createConcurrentRoot
-    root = createRootImpl(container, {hydrate, ui, wrapper})
+    root = await createRootImpl(container, {hydrate, ui, wrapper})
 
     mountedRootEntries.push({container, root})
     // we'll add it to the mounted containers regardless of whether it's actually
@@ -258,20 +260,22 @@ function render(
   })
 }
 
-function cleanup() {
-  mountedRootEntries.forEach(({root, container}) => {
-    act(() => {
-      root.unmount()
-    })
-    if (container.parentNode === document.body) {
-      document.body.removeChild(container)
-    }
-  })
+async function cleanup() {
+  await Promise.all(
+    mountedRootEntries.map(async ({root, container}) => {
+      await act(() => {
+        root.unmount()
+      })
+      if (container.parentNode === document.body) {
+        document.body.removeChild(container)
+      }
+    }),
+  )
   mountedRootEntries.length = 0
   mountedContainers.clear()
 }
 
-function renderHook(renderCallback, options = {}) {
+async function renderHook(renderCallback, options = {}) {
   const {initialProps, ...renderOptions} = options
 
   if (renderOptions.legacyRoot && typeof ReactDOM.render !== 'function') {
@@ -296,7 +300,7 @@ function renderHook(renderCallback, options = {}) {
     return null
   }
 
-  const {rerender: baseRerender, unmount} = render(
+  const {rerender: baseRerender, unmount} = await render(
     <TestComponent renderCallbackProps={initialProps} />,
     renderOptions,
   )