feat(waitFor): add complete and transparent support for fake timers

Closes #661

Author: kentcdodds

src/__tests__/wait-for-element-to-be-removed.js

@@ -1,17 +1,6 @@
+import {waitForElementToBeRemoved} from '..'
 import {renderIntoDocument} from './helpers/test-utils'
 
-function importModule() {
-  return require('../').waitForElementToBeRemoved
-}
-
-let waitForElementToBeRemoved
-
-beforeEach(() => {
-  jest.useRealTimers()
-  jest.resetModules()
-  waitForElementToBeRemoved = importModule()
-})
-
 test('resolves on mutation only when the element is removed', async () => {
   const {queryAllByTestId} = renderIntoDocument(`
     <div data-testid="div"></div>
@@ -89,55 +78,6 @@ test('after successful removal, fullfills promise with empty value (undefined)',
   return expect(waitResult).resolves.toBeUndefined()
 })
 
-describe('timers', () => {
-  const expectElementToBeRemoved = async () => {
-    const importedWaitForElementToBeRemoved = importModule()
-
-    const {queryAllByTestId} = renderIntoDocument(`
-  <div data-testid="div"></div>
-  <div data-testid="div"></div>
-`)
-    const divs = queryAllByTestId('div')
-    // first mutation
-    setTimeout(() => {
-      divs.forEach(d => d.setAttribute('id', 'mutated'))
-    })
-    // removal
-    setTimeout(() => {
-      divs.forEach(div => div.parentElement.removeChild(div))
-    }, 100)
-
-    const promise = importedWaitForElementToBeRemoved(
-      () => queryAllByTestId('div'),
-      {
-        timeout: 200,
-      },
-    )
-
-    if (setTimeout._isMockFunction) {
-      jest.advanceTimersByTime(110)
-    }
-
-    await promise
-  }
-
-  it('works with real timers', async () => {
-    jest.useRealTimers()
-    await expectElementToBeRemoved()
-  })
-  it('works with fake timers', async () => {
-    jest.useFakeTimers()
-    await expectElementToBeRemoved()
-  })
-})
-
-test("doesn't change jest's timers value when importing the module", () => {
-  jest.useFakeTimers()
-  importModule()
-
-  expect(window.setTimeout._isMockFunction).toEqual(true)
-})
-
 test('rethrows non-testing-lib errors', () => {
   let throwIt = false
   const div = document.createElement('div')

src/__tests__/wait-for.js

@@ -95,3 +95,44 @@ test('throws nice error if provided callback is not a function', () => {
     'Received `callback` arg must be a function',
   )
 })
+
+describe('works the same with fake timers', () => {
+  afterEach(() => jest.useRealTimers())
+
+  async function runTest(options) {
+    const doAsyncThing = () =>
+      new Promise(r => setTimeout(() => r('data'), 300))
+    let result
+    doAsyncThing().then(r => (result = r))
+
+    await waitFor(() => expect(result).toBe('data'), options)
+  }
+
+  test('real timers', async () => {
+    // the only difference when not using fake timers is this test will
+    // have to wait the full length of the timeout
+    await runTest()
+  })
+
+  test('legacy', async () => {
+    jest.useFakeTimers('legacy')
+    await runTest()
+  })
+
+  test('modern', async () => {
+    jest.useFakeTimers()
+    await runTest()
+  })
+
+  test('fake timer timeout', async () => {
+    jest.useFakeTimers()
+    await expect(
+      waitFor(
+        () => {
+          throw new Error('always throws')
+        },
+        {timeout: 10},
+      ),
+    ).rejects.toMatchInlineSnapshot(`[Error: always throws]`)
+  })
+})

src/helpers.js

@@ -2,44 +2,52 @@ const globalObj = typeof window === 'undefined' ? global : window
 
 // Currently this fn only supports jest timers, but it could support other test runners in the future.
 function runWithRealTimers(callback) {
-  const usingJestAndTimers =
-    typeof jest !== 'undefined' && typeof globalObj.setTimeout !== 'undefined'
-  const usingLegacyJestFakeTimers =
-    usingJestAndTimers &&
+  const fakeTimersType = getJestFakeTimersType()
+  if (fakeTimersType) {
+    jest.useRealTimers()
+  }
+
+  const callbackReturnValue = callback()
+
+  if (fakeTimersType) {
+    jest.useFakeTimers(fakeTimersType)
+  }
+
+  return callbackReturnValue
+}
+
+function getJestFakeTimersType() {
+  if (
+    typeof jest === 'undefined' ||
+    typeof globalObj.setTimeout === 'undefined'
+  ) {
+    return null
+  }
+
+  if (
     typeof globalObj.setTimeout._isMockFunction !== 'undefined' &&
     globalObj.setTimeout._isMockFunction
+  ) {
+    return 'legacy'
+  }
 
-  let usingModernJestFakeTimers = false
   if (
-    usingJestAndTimers &&
     typeof globalObj.setTimeout.clock !== 'undefined' &&
     typeof jest.getRealSystemTime !== 'undefined'
   ) {
     try {
       // jest.getRealSystemTime is only supported for Jest's `modern` fake timers and otherwise throws
       jest.getRealSystemTime()
-      usingModernJestFakeTimers = true
+      return 'modern'
     } catch {
       // not using Jest's modern fake timers
     }
   }
-
-  const usingJestFakeTimers =
-    usingLegacyJestFakeTimers || usingModernJestFakeTimers
-
-  if (usingJestFakeTimers) {
-    jest.useRealTimers()
-  }
-
-  const callbackReturnValue = callback()
-
-  if (usingJestFakeTimers) {
-    jest.useFakeTimers(usingModernJestFakeTimers ? 'modern' : 'legacy')
-  }
-
-  return callbackReturnValue
+  return null
 }
 
+const jestFakeTimersAreEnabled = () => Boolean(getJestFakeTimersType())
+
 // we only run our tests in node, and setImmediate is supported in node.
 // istanbul ignore next
 function setImmediatePolyfill(fn) {
@@ -117,4 +125,5 @@ export {
   setTimeoutFn as setTimeout,
   runWithRealTimers,
   checkContainerType,
+  jestFakeTimersAreEnabled,
 }

src/wait-for.js

@@ -5,6 +5,7 @@ import {
   setTimeout,
   clearTimeout,
   runWithRealTimers,
+  jestFakeTimersAreEnabled,
 } from './helpers'
 import {getConfig, runWithExpensiveErrorDiagnosticsDisabled} from './config'
 
@@ -35,22 +36,50 @@ function waitFor(
   }
 
   if (interval < 1) interval = 1
-  return new Promise((resolve, reject) => {
-    let lastError
-    const overallTimeoutTimer = setTimeout(onTimeout, timeout)
-    const intervalId = setInterval(checkCallback, interval)
+  return new Promise(async (resolve, reject) => {
+    let lastError, overallTimeoutTimer, intervalId, observer
+    let finished = false
 
-    const {MutationObserver} = getWindowFromNode(container)
-    const observer = new MutationObserver(checkCallback)
-    runWithRealTimers(() =>
-      observer.observe(container, mutationObserverOptions),
-    )
-    checkCallback()
+    // whether we're using fake timers or not, we still want the timeout support
+    runWithRealTimers(() => {
+      overallTimeoutTimer = setTimeout(onTimeout, timeout)
+    })
+
+    const usingFakeTimers = jestFakeTimersAreEnabled()
+    if (usingFakeTimers) {
+      // this is a dangerous rule to disable because it could lead to an
+      // infinite loop. However, eslint isn't smart enough to know that we're
+      // setting finished inside `onDone` which will be called when we're done
+      // waiting or when we've timed out.
+      // eslint-disable-next-line no-unmodified-loop-condition
+      while (!finished) {
+        jest.advanceTimersByTime(interval)
+        // in this rare case, we *need* to wait for in-flight promises
+        // to resolve before continuing. We don't need to take advantage
+        // of parallelization so we're fine.
+        // https://stackoverflow.com/a/59243586/971592
+        // eslint-disable-next-line no-await-in-loop
+        await new Promise(r => setImmediate(r))
+        checkCallback()
+      }
+    } else {
+      intervalId = setInterval(checkCallback, interval)
+      const {MutationObserver} = getWindowFromNode(container)
+      observer = new MutationObserver(checkCallback)
+      observer.observe(container, mutationObserverOptions)
+      checkCallback()
+    }
 
     function onDone(error, result) {
-      clearTimeout(overallTimeoutTimer)
-      clearInterval(intervalId)
-      setImmediate(() => observer.disconnect())
+      finished = true
+      runWithRealTimers(() => {
+        clearTimeout(overallTimeoutTimer)
+      })
+
+      if (!usingFakeTimers) {
+        clearInterval(intervalId)
+        setImmediate(() => observer.disconnect())
+      }
 
       if (error) {
         reject(error)
@@ -62,9 +91,9 @@ function waitFor(
     function checkCallback() {
       try {
         onDone(null, runWithExpensiveErrorDiagnosticsDisabled(callback))
-        // If `callback` throws, wait for the next mutation or timeout.
+        // If `callback` throws, wait for the next mutation, interval, or timeout.
       } catch (error) {
-        // Save the callback error to reject the promise with it.
+        // Save the most recent callback error to reject the promise with it in the event of a timeout
         lastError = error
       }
     }