From c02db2eb52bd9004d91ada500ac7b4631949ef1f Mon Sep 17 00:00:00 2001
From: eps1lon <sebastian.silbermann@vercel.com>
Date: Tue, 26 Nov 2024 00:04:00 +0100
Subject: [PATCH] Add async render APIs

---
 src/__tests__/renderAsync.js |  25 +++++
 src/act-compat.js            |  14 +++
 src/pure.js                  | 189 ++++++++++++++++++++++++++++++++++-
 types/index.d.ts             |  81 +++++++++++++++
 4 files changed, 308 insertions(+), 1 deletion(-)
 create mode 100644 src/__tests__/renderAsync.js

diff --git a/src/__tests__/renderAsync.js b/src/__tests__/renderAsync.js
new file mode 100644
index 00000000..4e580b82
--- /dev/null
+++ b/src/__tests__/renderAsync.js
@@ -0,0 +1,25 @@
+import * as React from 'react'
+import {act, renderAsync} from '../'
+
+test('async data requires async APIs', async () => {
+  const {promise, resolve} = Promise.withResolvers()
+
+  function Component() {
+    const value = React.use(promise)
+    return <div>{value}</div>
+  }
+
+  const {container} = await renderAsync(
+    <React.Suspense fallback="loading...">
+      <Component />
+    </React.Suspense>,
+  )
+
+  expect(container).toHaveTextContent('loading...')
+
+  await act(async () => {
+    resolve('Hello, Dave!')
+  })
+
+  expect(container).toHaveTextContent('Hello, Dave!')
+})
diff --git a/src/act-compat.js b/src/act-compat.js
index 6eaec0fb..8d5da94b 100644
--- a/src/act-compat.js
+++ b/src/act-compat.js
@@ -82,8 +82,22 @@ function withGlobalActEnvironment(actImplementation) {
 
 const act = withGlobalActEnvironment(reactAct)
 
+async function actAsync(scope) {
+  const previousActEnvironment = getIsReactActEnvironment()
+  setIsReactActEnvironment(true)
+  try {
+    // React.act isn't async yet so we need to force it.
+    return await reactAct(async () => {
+      scope()
+    })
+  } finally {
+    setIsReactActEnvironment(previousActEnvironment)
+  }
+}
+
 export default act
 export {
+  actAsync,
   setIsReactActEnvironment as setReactActEnvironment,
   getIsReactActEnvironment,
 }
diff --git a/src/pure.js b/src/pure.js
index f546af98..a25bcd88 100644
--- a/src/pure.js
+++ b/src/pure.js
@@ -7,6 +7,7 @@ import {
   configure as configureDTL,
 } from '@testing-library/dom'
 import act, {
+  actAsync,
   getIsReactActEnvironment,
   setReactActEnvironment,
 } from './act-compat'
@@ -196,6 +197,64 @@ function renderRoot(
   }
 }
 
+async function renderRootAsync(
+  ui,
+  {baseElement, container, hydrate, queries, root, wrapper: WrapperComponent},
+) {
+  await actAsync(() => {
+    if (hydrate) {
+      root.hydrate(
+        strictModeIfNeeded(wrapUiIfNeeded(ui, WrapperComponent)),
+        container,
+      )
+    } else {
+      root.render(
+        strictModeIfNeeded(wrapUiIfNeeded(ui, WrapperComponent)),
+        container,
+      )
+    }
+  })
+
+  return {
+    container,
+    baseElement,
+    debug: (el = baseElement, maxLength, options) =>
+      Array.isArray(el)
+        ? // eslint-disable-next-line no-console
+          el.forEach(e => console.log(prettyDOM(e, maxLength, options)))
+        : // eslint-disable-next-line no-console,
+          console.log(prettyDOM(el, maxLength, options)),
+    unmount: async () => {
+      await actAsync(() => {
+        root.unmount()
+      })
+    },
+    rerender: async rerenderUi => {
+      await renderRootAsync(rerenderUi, {
+        container,
+        baseElement,
+        root,
+        wrapper: WrapperComponent,
+      })
+      // Intentionally do not return anything to avoid unnecessarily complicating the API.
+      // folks can use all the same utilities we return in the first place that are bound to the container
+    },
+    asFragment: () => {
+      /* istanbul ignore else (old jsdom limitation) */
+      if (typeof document.createRange === 'function') {
+        return document
+          .createRange()
+          .createContextualFragment(container.innerHTML)
+      } else {
+        const template = document.createElement('template')
+        template.innerHTML = container.innerHTML
+        return template.content
+      }
+    },
+    ...getQueriesForElement(baseElement, queries),
+  }
+}
+
 function render(
   ui,
   {
@@ -258,6 +317,68 @@ function render(
   })
 }
 
+function renderAsync(
+  ui,
+  {
+    container,
+    baseElement = container,
+    legacyRoot = false,
+    queries,
+    hydrate = false,
+    wrapper,
+  } = {},
+) {
+  if (legacyRoot && typeof ReactDOM.render !== 'function') {
+    const error = new Error(
+      '`legacyRoot: true` is not supported in this version of React. ' +
+        'If your app runs React 19 or later, you should remove this flag. ' +
+        'If your app runs React 18 or earlier, visit https://react.dev/blog/2022/03/08/react-18-upgrade-guide for upgrade instructions.',
+    )
+    Error.captureStackTrace(error, render)
+    throw error
+  }
+
+  if (!baseElement) {
+    // default to document.body instead of documentElement to avoid output of potentially-large
+    // head elements (such as JSS style blocks) in debug output
+    baseElement = document.body
+  }
+  if (!container) {
+    container = baseElement.appendChild(document.createElement('div'))
+  }
+
+  let root
+  // 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})
+
+    mountedRootEntries.push({container, root})
+    // we'll add it to the mounted containers regardless of whether it's actually
+    // added to document.body so the cleanup method works regardless of whether
+    // they're passing us a custom container or not.
+    mountedContainers.add(container)
+  } else {
+    mountedRootEntries.forEach(rootEntry => {
+      // Else is unreachable since `mountedContainers` has the `container`.
+      // Only reachable if one would accidentally add the container to `mountedContainers` but not the root to `mountedRootEntries`
+      /* istanbul ignore else */
+      if (rootEntry.container === container) {
+        root = rootEntry.root
+      }
+    })
+  }
+
+  return renderRootAsync(ui, {
+    container,
+    baseElement,
+    queries,
+    hydrate,
+    wrapper,
+    root,
+  })
+}
+
 function cleanup() {
   mountedRootEntries.forEach(({root, container}) => {
     act(() => {
@@ -271,6 +392,21 @@ function cleanup() {
   mountedContainers.clear()
 }
 
+async function cleanupAsync() {
+  for (const {root, container} of mountedRootEntries) {
+    // eslint-disable-next-line no-await-in-loop -- act calls can't overlap
+    await actAsync(() => {
+      root.unmount()
+    })
+    if (container.parentNode === document.body) {
+      document.body.removeChild(container)
+    }
+  }
+
+  mountedRootEntries.length = 0
+  mountedContainers.clear()
+}
+
 function renderHook(renderCallback, options = {}) {
   const {initialProps, ...renderOptions} = options
 
@@ -310,8 +446,59 @@ function renderHook(renderCallback, options = {}) {
   return {result, rerender, unmount}
 }
 
+async function renderHookAsync(renderCallback, options = {}) {
+  const {initialProps, ...renderOptions} = options
+
+  if (renderOptions.legacyRoot && typeof ReactDOM.render !== 'function') {
+    const error = new Error(
+      '`legacyRoot: true` is not supported in this version of React. ' +
+        'If your app runs React 19 or later, you should remove this flag. ' +
+        'If your app runs React 18 or earlier, visit https://react.dev/blog/2022/03/08/react-18-upgrade-guide for upgrade instructions.',
+    )
+    Error.captureStackTrace(error, renderHookAsync)
+    throw error
+  }
+
+  const result = React.createRef()
+
+  function TestComponent({renderCallbackProps}) {
+    const pendingResult = renderCallback(renderCallbackProps)
+
+    React.useEffect(() => {
+      result.current = pendingResult
+    })
+
+    return null
+  }
+
+  const {rerender: baseRerender, unmount} = await renderAsync(
+    <TestComponent renderCallbackProps={initialProps} />,
+    renderOptions,
+  )
+
+  function rerender(rerenderCallbackProps) {
+    return baseRerender(
+      <TestComponent renderCallbackProps={rerenderCallbackProps} />,
+    )
+  }
+
+  return {result, rerender, unmount}
+}
+
 // just re-export everything from dom-testing-library
 export * from '@testing-library/dom'
-export {render, renderHook, cleanup, act, fireEvent, getConfig, configure}
+export {
+  render,
+  renderAsync,
+  renderHook,
+  renderHookAsync,
+  cleanup,
+  cleanupAsync,
+  act,
+  fireEvent,
+  // TODO: fireEventAsync
+  getConfig,
+  configure,
+}
 
 /* eslint func-name-matching:0 */
diff --git a/types/index.d.ts b/types/index.d.ts
index 3ad8cf46..71a8d60b 100644
--- a/types/index.d.ts
+++ b/types/index.d.ts
@@ -46,6 +46,27 @@ export type RenderResult<
   asFragment: () => DocumentFragment
 } & {[P in keyof Q]: BoundFunction<Q[P]>}
 
+export type RenderAsyncResult<
+  Q extends Queries = typeof queries,
+  Container extends RendererableContainer | HydrateableContainer = HTMLElement,
+  BaseElement extends RendererableContainer | HydrateableContainer = Container,
+> = {
+  container: Container
+  baseElement: BaseElement
+  debug: (
+    baseElement?:
+      | RendererableContainer
+      | HydrateableContainer
+      | Array<RendererableContainer | HydrateableContainer>
+      | undefined,
+    maxLength?: number | undefined,
+    options?: prettyFormat.OptionsReceived | undefined,
+  ) => void
+  rerender: (ui: React.ReactNode) => Promise<void>
+  unmount: () => Promise<void>
+  asFragment: () => DocumentFragment
+} & {[P in keyof Q]: BoundFunction<Q[P]>}
+
 /** @deprecated */
 export type BaseRenderOptions<
   Q extends Queries,
@@ -152,6 +173,22 @@ export function render(
   options?: Omit<RenderOptions, 'queries'> | undefined,
 ): RenderResult
 
+/**
+ * Render into a container which is appended to document.body. It should be used with cleanup.
+ */
+export function renderAsync<
+  Q extends Queries = typeof queries,
+  Container extends RendererableContainer | HydrateableContainer = HTMLElement,
+  BaseElement extends RendererableContainer | HydrateableContainer = Container,
+>(
+  ui: React.ReactNode,
+  options: RenderOptions<Q, Container, BaseElement>,
+): Promise<RenderAsyncResult<Q, Container, BaseElement>>
+export function renderAsync(
+  ui: React.ReactNode,
+  options?: Omit<RenderOptions, 'queries'> | undefined,
+): Promise<RenderAsyncResult>
+
 export interface RenderHookResult<Result, Props> {
   /**
    * Triggers a re-render. The props will be passed to your renderHook callback.
@@ -174,6 +211,28 @@ export interface RenderHookResult<Result, Props> {
   unmount: () => void
 }
 
+export interface RenderHookAsyncResult<Result, Props> {
+  /**
+   * Triggers a re-render. The props will be passed to your renderHook callback.
+   */
+  rerender: (props?: Props) => Promise<void>
+  /**
+   * This is a stable reference to the latest value returned by your renderHook
+   * callback
+   */
+  result: {
+    /**
+     * The value returned by your renderHook callback
+     */
+    current: Result
+  }
+  /**
+   * Unmounts the test component. This is useful for when you need to test
+   * any cleanup your useEffects have.
+   */
+  unmount: () => Promise<void>
+}
+
 /** @deprecated */
 export type BaseRenderHookOptions<
   Props,
@@ -242,11 +301,31 @@ export function renderHook<
   options?: RenderHookOptions<Props, Q, Container, BaseElement> | undefined,
 ): RenderHookResult<Result, Props>
 
+/**
+ * Allows you to render a hook within a test React component without having to
+ * create that component yourself.
+ */
+export function renderHookAsync<
+  Result,
+  Props,
+  Q extends Queries = typeof queries,
+  Container extends RendererableContainer | HydrateableContainer = HTMLElement,
+  BaseElement extends RendererableContainer | HydrateableContainer = Container,
+>(
+  render: (initialProps: Props) => Result,
+  options?: RenderHookOptions<Props, Q, Container, BaseElement> | undefined,
+): Promise<RenderHookResult<Result, Props>>
+
 /**
  * Unmounts React trees that were mounted with render.
  */
 export function cleanup(): void
 
+/**
+ * Unmounts React trees that were mounted with render.
+ */
+export function cleanupAsync(): Promise<void>
+
 /**
  * Simply calls React.act(cb)
  * If that's not available (older version of react) then it
@@ -256,3 +335,5 @@ export function cleanup(): void
 export const act: 0 extends 1 & typeof reactAct
   ? typeof reactDeprecatedAct
   : typeof reactAct
+
+export function actAsync(scope: () => void | Promise<void>): Promise<void>