feat(waitFor): allow returning a promise (#782)

kentcdodds · web-flow · commit 4e5e533e7922 · 2020-10-12T16:36:09.000-06:00
You can now return a promise in the waitFor callback and your callback will not be called again until the promise rejects. If the promise rejects, then the callback will be called again until it gets a promise which resolves or it times out.
diff --git a/src/__tests__/wait-for.js b/src/__tests__/wait-for.js
@@ -2,6 +2,15 @@ import {waitFor} from '../'
 import {configure, getConfig} from '../config'
 import {renderIntoDocument} from './helpers/test-utils'
 
+function deferred() {
+  let resolve, reject
+  const promise = new Promise((res, rej) => {
+    resolve = res
+    reject = rej
+  })
+  return {promise, resolve, reject}
+}
+
 test('waits callback to not throw an error', async () => {
   const spy = jest.fn()
   // we are using random timeout here to simulate a real-time example
@@ -145,3 +154,41 @@ test('should delegate to config.getElementError', async () => {
   expect(error.message).toMatchInlineSnapshot(`"Custom element error"`)
   configure(originalConfig)
 })
+
+test('when a promise is returned, it does not call the callback again until that promise rejects', async () => {
+  const sleep = t => new Promise(r => setTimeout(r, t))
+  const p1 = deferred()
+  const waitForCb = jest.fn(() => p1.promise)
+  const waitForPromise = waitFor(waitForCb, {interval: 1})
+  expect(waitForCb).toHaveBeenCalledTimes(1)
+  waitForCb.mockClear()
+  await sleep(50)
+  expect(waitForCb).toHaveBeenCalledTimes(0)
+
+  const p2 = deferred()
+  waitForCb.mockImplementation(() => p2.promise)
+
+  p1.reject('p1 rejection (should not fail this test)')
+  await sleep(50)
+
+  expect(waitForCb).toHaveBeenCalledTimes(1)
+  p2.resolve()
+
+  await waitForPromise
+})
+
+test('when a promise is returned, if that is not resolved within the timeout, then waitFor is rejected', async () => {
+  const sleep = t => new Promise(r => setTimeout(r, t))
+  const {promise} = deferred()
+  const waitForPromise = waitFor(() => promise, {timeout: 1}).catch(e => e)
+  await sleep(5)
+
+  expect((await waitForPromise).message).toMatchInlineSnapshot(`
+    "Timed out in waitFor.
+
+    <html>
+      <head />
+      <body />
+    </html>"
+  `)
+})
diff --git a/src/wait-for.js b/src/wait-for.js
@@ -47,6 +47,7 @@ function waitFor(
   return new Promise(async (resolve, reject) => {
     let lastError, intervalId, observer
     let finished = false
+    let promiseStatus = 'idle'
 
     const overallTimeoutTimer = setTimeout(handleTimeout, timeout)
 
@@ -104,8 +105,24 @@ function waitFor(
     }
 
     function checkCallback() {
+      if (promiseStatus === 'pending') return
       try {
-        onDone(null, runWithExpensiveErrorDiagnosticsDisabled(callback))
+        const result = runWithExpensiveErrorDiagnosticsDisabled(callback)
+        if (typeof result?.then === 'function') {
+          promiseStatus = 'pending'
+          result.then(
+            resolvedValue => {
+              promiseStatus = 'resolved'
+              onDone(null, resolvedValue)
+            },
+            rejectedValue => {
+              promiseStatus = 'rejected'
+              lastError = rejectedValue
+            },
+          )
+        } else {
+          onDone(null, result)
+        }
         // If `callback` throws, wait for the next mutation, interval, or timeout.
       } catch (error) {
         // Save the most recent callback error to reject the promise with it in the event of a timeout
diff --git a/types/__tests__/type-tests.ts b/types/__tests__/type-tests.ts
@@ -174,6 +174,5 @@ async function testWaitFors() {
   await waitForElementToBeRemoved(getByText(element, 'apple'))
   await waitForElementToBeRemoved(getAllByText(element, 'apple'))
 
-  // $ExpectError
   await waitFor(async () => {})
 }
diff --git a/types/wait-for.d.ts b/types/wait-for.d.ts
@@ -7,6 +7,6 @@ export interface waitForOptions {
 }
 
 export function waitFor<T>(
-  callback: () => T extends Promise<any> ? never : T,
+  callback: () => T | Promise<T>,
   options?: waitForOptions,
 ): Promise<T>

Original file line number	Diff line number	Diff line change
`@@ -174,6 +174,5 @@ async function testWaitFors() {`
`174`	`174`	`await waitForElementToBeRemoved(getByText(element, 'apple'))`
`175`	`175`	`await waitForElementToBeRemoved(getAllByText(element, 'apple'))`
`176`	`176`
`177`		`- // $ExpectError`
`178`	`177`	`await waitFor(async () => {})`
`179`	`178`	`}`
Original file line number	Diff line number	Diff line change
`@@ -7,6 +7,6 @@ export interface waitForOptions {`
`7`	`7`	`}`
`8`	`8`
`9`	`9`	`export function waitFor<T>(`
`10`		`- callback: () => T extends Promise<any> ? never : T,`
	`10`	`+ callback: () => T \| Promise<T>,`
`11`	`11`	`options?: waitForOptions,`
`12`	`12`	`): Promise<T>`