From 174a19449656b4c7d467ce4fb3e7294b98dc6ca0 Mon Sep 17 00:00:00 2001
From: jmbeach <jaredbeachdesign@gmail.com>
Date: Mon, 27 Dec 2021 11:27:02 -0600
Subject: [PATCH] feat: adds onTimeout option to waitFor

---
 src/__tests__/waitFor.test.js | 26 ++++++++++++++++++++++++++
 src/waitFor.js                |  5 +++++
 typings/index.d.ts            |  1 +
 3 files changed, 32 insertions(+)

diff --git a/src/__tests__/waitFor.test.js b/src/__tests__/waitFor.test.js
index 75e3dd5af..793c45f2c 100644
--- a/src/__tests__/waitFor.test.js
+++ b/src/__tests__/waitFor.test.js
@@ -115,6 +115,32 @@ test.each([TimerMode.Legacy, TimerMode.Modern])(
   }
 );
 
+test.each([TimerMode.Legacy, TimerMode.Modern])(
+  'waits for assertion until timeout is met with %s fake timers',
+  async (fakeTimerType) => {
+    jest.useFakeTimers(fakeTimerType);
+
+    const mockErrorFn = jest.fn(() => {
+      throw Error('test');
+    });
+
+    const mockHandleFn = jest.fn((e) => e);
+
+    try {
+      await waitFor(() => mockErrorFn(), {
+        timeout: 400,
+        interval: 200,
+        onTimeout: mockHandleFn,
+      });
+    } catch (error) {
+      // suppress
+    }
+
+    expect(mockErrorFn).toHaveBeenCalledTimes(3);
+    expect(mockHandleFn).toHaveBeenCalledTimes(1);
+  }
+);
+
 test.each([TimerMode.Legacy, TimerMode.Legacy])(
   'awaiting something that succeeds before timeout works with %s fake timers',
   async (fakeTimerType) => {
diff --git a/src/waitFor.js b/src/waitFor.js
index 79285c357..3e5a881f0 100644
--- a/src/waitFor.js
+++ b/src/waitFor.js
@@ -29,6 +29,7 @@ export type WaitForOptions = {
   timeout?: number,
   interval?: number,
   stackTraceError?: ErrorWithStack,
+  onTimeout?: (error: Error) => Error,
 };
 
 function waitForInternal<T>(
@@ -37,6 +38,7 @@ function waitForInternal<T>(
     timeout = DEFAULT_TIMEOUT,
     interval = DEFAULT_INTERVAL,
     stackTraceError,
+    onTimeout,
   }: WaitForOptions
 ): Promise<T> {
   if (typeof expectation !== 'function') {
@@ -179,6 +181,9 @@ function waitForInternal<T>(
           copyStackTrace(error, stackTraceError);
         }
       }
+      if (typeof onTimeout === 'function') {
+        onTimeout(error);
+      }
       onDone(error, null);
     }
   });
diff --git a/typings/index.d.ts b/typings/index.d.ts
index c0b6ea4bf..89ecc59d7 100644
--- a/typings/index.d.ts
+++ b/typings/index.d.ts
@@ -388,6 +388,7 @@ type TextMatchOptions = {
 type WaitForOptions = {
   timeout?: number;
   interval?: number;
+  onTimeout?: (error: Error) => Error;
 };
 
 export type WaitForFunction = <T = any>(