fix: cleanup of waitFor timers when an error occurs during onTimeout printout (#1918)

Co-authored-by: Maciej Jastrzębski <mdjastrzebski@gmail.com>

Author: ggilles

src/__tests__/wait-for.test.tsx

@@ -2,6 +2,7 @@ import * as React from 'react';
 import { Pressable, Text, TouchableOpacity, View } from 'react-native';
 
 import { configure, fireEvent, render, screen, waitFor } from '..';
+import { cleanup } from '../pure';
 import type { TimerType } from '../test-utils/timers';
 import { setupTimeType } from '../test-utils/timers';
 
@@ -175,6 +176,39 @@ describe('timeout errors', () => {
       }),
     ).rejects.toThrow('Unable to find an element with text: Never appears');
   });
+
+  test('rejects with onTimeout error when onTimeout callback throws', async () => {
+    await expect(
+      waitFor(
+        () => {
+          throw new Error('Original timeout error');
+        },
+        {
+          timeout: 10,
+          onTimeout: () => {
+            throw new Error('onTimeout failed');
+          },
+        },
+      ),
+    ).rejects.toThrow('onTimeout failed');
+  });
+
+  test('rejects with onTimeout value when onTimeout throws a non-Error value', async () => {
+    await expect(
+      waitFor(
+        () => {
+          throw new Error('Original timeout error');
+        },
+        {
+          timeout: 10,
+          onTimeout: () => {
+            // eslint-disable-next-line @typescript-eslint/only-throw-error
+            throw 'onTimeout failed with string';
+          },
+        },
+      ),
+    ).rejects.toBe('onTimeout failed with string');
+  });
 });
 
 describe('error handling', () => {
@@ -235,6 +269,46 @@ describe('error handling', () => {
       'Changed from using fake timers to real timers while using waitFor',
     );
   });
+
+  test('cleanup rejects pending waitFor calls and stops real timer polling', async () => {
+    const expectation = jest.fn(() => {
+      throw new Error('Not ready yet');
+    });
+
+    const waitForPromise = waitFor(expectation, { timeout: 300, interval: 20 });
+
+    await new Promise((resolve) => setTimeout(resolve, 60));
+    expect(expectation).toHaveBeenCalled();
+
+    await cleanup();
+    await expect(waitForPromise).rejects.toThrow('waitFor was aborted by cleanup');
+
+    const callsAfterCleanup = expectation.mock.calls.length;
+
+    await new Promise((resolve) => setTimeout(resolve, 60));
+    expect(expectation).toHaveBeenCalledTimes(callsAfterCleanup);
+  });
+
+  test('async expectation resolving after cleanup abort has no effect', async () => {
+    let resolveExpectation!: (value: string) => void;
+    const expectationPromise = new Promise<string>((resolve) => {
+      resolveExpectation = resolve;
+    });
+
+    const waitForPromise = waitFor(() => expectationPromise, { timeout: 300, interval: 20 });
+
+    // Wait for at least one interval tick so the expectation is called and its promise is pending
+    await new Promise((resolve) => setTimeout(resolve, 30));
+
+    // Start cleanup while the expectation promise is still pending
+    const cleanupPromise = cleanup();
+
+    // Resolve the expectation promise while cleanup is in progress
+    resolveExpectation('success');
+
+    await cleanupPromise;
+    await expect(waitForPromise).rejects.toThrow('waitFor was aborted by cleanup');
+  });
 });
 
 describe('configuration', () => {

src/cleanup.ts

@@ -17,3 +17,7 @@ export async function cleanup() {
 export function addToCleanupQueue(fn: CleanUpFunction) {
   cleanupQueue.add(fn);
 }
+
+export function removeFromCleanupQueue(fn: CleanUpFunction) {
+  cleanupQueue.delete(fn);
+}

src/wait-for.ts

@@ -1,5 +1,6 @@
 /* globals jest */
 import { act } from './act';
+import { addToCleanupQueue, removeFromCleanupQueue } from './cleanup';
 import { getConfig } from './config';
 import { flushMicroTasks } from './flush-micro-tasks';
 import { copyStackTraceIfNeeded, ErrorWithStack } from './helpers/errors';
@@ -35,11 +36,13 @@ function waitForInternal<T>(
 
   // eslint-disable-next-line no-async-promise-executor
   return new Promise(async (resolve, reject) => {
-    let lastError: unknown, intervalId: ReturnType<typeof setTimeout>;
+    let lastError: unknown;
+    let intervalId: ReturnType<typeof setInterval> | null = null;
     let finished = false;
     let promiseStatus = 'idle';
 
-    let overallTimeoutTimer: NodeJS.Timeout | null = null;
+    let overallTimeoutTimer: ReturnType<typeof setTimeout> | null = null;
+    const cleanupQueueCallback = () => finalizeWaitFor({ rejectOnAbort: true });
 
     const fakeTimersType = getJestFakeTimersType();
 
@@ -94,19 +97,42 @@ function waitForInternal<T>(
     } else {
       overallTimeoutTimer = setTimeout(handleTimeout, timeout);
       intervalId = setInterval(checkRealTimersCallback, interval);
+      addToCleanupQueue(cleanupQueueCallback);
       checkExpectation();
     }
 
-    function onDone(done: { type: 'result'; result: T } | { type: 'error'; error: unknown }) {
+    function finalizeWaitFor({ rejectOnAbort = false } = {}) {
+      /* istanbul ignore next */
+      if (finished) {
+        return;
+      }
+
       finished = true;
+
+      removeFromCleanupQueue(cleanupQueueCallback);
+
+      if (rejectOnAbort) {
+        reject(new Error('waitFor was aborted by cleanup'));
+      }
+
       if (overallTimeoutTimer) {
         clearTimeout(overallTimeoutTimer);
+        overallTimeoutTimer = null;
       }
 
-      if (!fakeTimersType) {
+      if (intervalId) {
         clearInterval(intervalId);
+        intervalId = null;
+      }
+    }
+
+    function onDone(done: { type: 'result'; result: T } | { type: 'error'; error: unknown }) {
+      if (finished) {
+        return;
       }
 
+      finalizeWaitFor();
+
       if (done.type === 'error') {
         reject(done.error);
       } else {
@@ -120,7 +146,7 @@ function waitForInternal<T>(
           `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`,
         );
         copyStackTraceIfNeeded(error, stackTraceError);
-        return reject(error);
+        return onDone({ type: 'error', error });
       } else {
         return checkExpectation();
       }
@@ -176,13 +202,19 @@ function waitForInternal<T>(
         error = new Error('Timed out in waitFor.');
         copyStackTraceIfNeeded(error, stackTraceError);
       }
+
+      let errorForRejection: unknown = error;
       if (typeof onTimeout === 'function') {
-        const result = onTimeout(error);
-        if (result) {
-          error = result;
+        try {
+          const result = onTimeout(error);
+          if (result) {
+            errorForRejection = result;
+          }
+        } catch (onTimeoutError) {
+          errorForRejection = onTimeoutError;
         }
       }
-      onDone({ type: 'error', error });
+      onDone({ type: 'error', error: errorForRejection });
     }
   });
 }