base waitFor on RTL implementation

Author: mikeduminy

src/waitFor.js

@@ -1,4 +1,5 @@
 // @flow
+/* globals jest */
 
 import * as React from 'react';
 import act from './act';
@@ -7,9 +8,14 @@ import {
   throwRemovedFunctionError,
   copyStackTrace,
 } from './helpers/errors';
-import { setTimeout } from './helpers/getTimerFuncs';
+import {
+  setTimeout,
+  clearTimeout,
+  setImmediate,
+  jestFakeTimersAreEnabled,
+} from './helpers/timers';
 
-const DEFAULT_TIMEOUT = 4500;
+const DEFAULT_TIMEOUT = 1000;
 const DEFAULT_INTERVAL = 50;
 
 function checkReactVersionAtLeast(major: number, minor: number): boolean {
@@ -22,56 +28,169 @@ function checkReactVersionAtLeast(major: number, minor: number): boolean {
 export type WaitForOptions = {
   timeout?: number,
   interval?: number,
+  stackTraceError?: ErrorWithStack,
 };
 
 function waitForInternal<T>(
   expectation: () => T,
-  options?: WaitForOptions
+  {
+    timeout = DEFAULT_TIMEOUT,
+    interval = DEFAULT_INTERVAL,
+    stackTraceError,
+  }: WaitForOptions
 ): Promise<T> {
-  const timeout = options?.timeout ?? DEFAULT_TIMEOUT;
-  let interval = options?.interval ?? DEFAULT_INTERVAL;
-  // Being able to display a useful stack trace requires generating it before doing anything async
-  const stackTraceError = new ErrorWithStack('STACK_TRACE_ERROR', waitFor);
+  if (typeof expectation !== 'function') {
+    throw new TypeError('Received `expectation` arg must be a function');
+  }
+
+  // eslint-disable-next-line no-async-promise-executor
+  return new Promise(async (resolve, reject) => {
+    let lastError, intervalId;
+    let finished = false;
+    let promiseStatus = 'idle';
+
+    const overallTimeoutTimer = setTimeout(handleTimeout, timeout);
 
-  if (interval < 1) interval = 1;
-  const maxTries = Math.ceil(timeout / interval);
-  let tries = 0;
+    const usingFakeTimers = jestFakeTimersAreEnabled();
+    if (usingFakeTimers) {
+      checkExpectation();
+      // 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) {
+        if (!jestFakeTimersAreEnabled()) {
+          const error = new Error(
+            `Changed from using fake timers to real timers while using waitFor. This is not allowed and will result in very strange behavior. Please ensure you're awaiting all async things your test is doing before changing to real timers. For more info, please go to https://github.com/testing-library/dom-testing-library/issues/830`
+          );
+          if (stackTraceError) {
+            copyStackTrace(error, stackTraceError);
+          }
+          reject(error);
+          return;
+        }
+        // we *could* (maybe should?) use `advanceTimersToNextTimer` but it's
+        // possible that could make this loop go on forever if someone is using
+        // third party code that's setting up recursive timers so rapidly that
+        // the user's timer's don't get a chance to resolve. So we'll advance
+        // by an interval instead. (We have a test for this case).
+        jest.advanceTimersByTime(interval);
+
+        // It's really important that checkExpectation is run *before* we flush
+        // in-flight promises. To be honest, I'm not sure why, and I can't quite
+        // think of a way to reproduce the problem in a test, but I spent
+        // an entire day banging my head against a wall on this.
+        checkExpectation();
+
+        // 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((resolve) => setImmediate(resolve));
+      }
+    } else {
+      intervalId = setInterval(checkRealTimersCallback, interval);
+      checkExpectation();
+    }
+
+    function onDone(error, result) {
+      finished = true;
+      clearTimeout(overallTimeoutTimer);
+
+      if (!usingFakeTimers) {
+        clearInterval(intervalId);
+      }
 
-  return new Promise((resolve, reject) => {
-    const rejectOrRerun = (error) => {
-      if (tries > maxTries) {
-        copyStackTrace(error, stackTraceError);
+      if (error) {
         reject(error);
-        return;
+      } else {
+        // $FlowIgnore[incompatible-return] error and result are mutually exclusive
+        resolve(result);
+      }
+    }
+
+    function checkRealTimersCallback() {
+      if (jestFakeTimersAreEnabled()) {
+        const error = new Error(
+          `Changed from using real timers to fake timers while using waitFor. This is not allowed and will result in very strange behavior. Please ensure you're awaiting all async things your test is doing before changing to fake timers. For more info, please go to https://github.com/testing-library/dom-testing-library/issues/830`
+        );
+        if (stackTraceError) {
+          copyStackTrace(error, stackTraceError);
+        }
+        return reject(error);
+      } else {
+        return checkExpectation();
       }
-      setTimeout(runExpectation, interval);
-    };
-    function runExpectation() {
-      tries += 1;
+    }
+
+    function checkExpectation() {
+      if (promiseStatus === 'pending') return;
       try {
-        const result = Promise.resolve(expectation());
-        resolve(result);
+        const result = expectation();
+
+        // $FlowIgnore[incompatible-type]
+        if (typeof result?.then === 'function') {
+          promiseStatus = 'pending';
+          // eslint-disable-next-line promise/catch-or-return
+          result.then(
+            (resolvedValue) => {
+              promiseStatus = 'resolved';
+              onDone(null, resolvedValue);
+              return;
+            },
+            (rejectedValue) => {
+              promiseStatus = 'rejected';
+              lastError = rejectedValue;
+              return;
+            }
+          );
+        } else {
+          onDone(null, result);
+        }
+        // If `callback` throws, wait for the next mutation, interval, or timeout.
       } catch (error) {
-        rejectOrRerun(error);
+        // Save the most recent callback error to reject the promise with it in the event of a timeout
+        lastError = error;
+      }
+    }
+
+    function handleTimeout() {
+      let error;
+      if (lastError) {
+        error = lastError;
+        if (stackTraceError) {
+          copyStackTrace(error, stackTraceError);
+        }
+      } else {
+        error = new Error('Timed out in waitFor.');
+        if (stackTraceError) {
+          copyStackTrace(error, stackTraceError);
+        }
       }
+      onDone(error, null);
     }
-    setTimeout(runExpectation, 0);
   });
 }
 
 export default async function waitFor<T>(
   expectation: () => T,
   options?: WaitForOptions
 ): Promise<T> {
+  // Being able to display a useful stack trace requires generating it before doing anything async
+  const stackTraceError = new ErrorWithStack('STACK_TRACE_ERROR', waitFor);
+  const optionsWithStackTrace = { stackTraceError, ...options };
+
   if (!checkReactVersionAtLeast(16, 9)) {
-    return waitForInternal(expectation, options);
+    return waitForInternal(expectation, optionsWithStackTrace);
   }
 
   let result: T;
 
   //$FlowFixMe: `act` has incorrect flow typing
   await act(async () => {
-    result = await waitForInternal(expectation, options);
+    result = await waitForInternal(expectation, optionsWithStackTrace);
   });
 
   //$FlowFixMe: either we have result or `waitFor` threw error