From 6e6b5a2ffae38956dbc484e523544d85acd466df Mon Sep 17 00:00:00 2001
From: Lei Chen <leichen@Leis-MacBook-Pro.local>
Date: Thu, 9 Sep 2021 13:48:31 -0700
Subject: [PATCH] fix the code after code review and clean up

---
 src/__tests__/asyncHook.fakeTimers.test.ts |  6 ++---
 src/core/asyncUtils.ts                     | 19 ++++++++-------
 src/helpers/createTimeoutController.ts     | 28 ++++++++++++++++++----
 3 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/src/__tests__/asyncHook.fakeTimers.test.ts b/src/__tests__/asyncHook.fakeTimers.test.ts
index bf79ed27..90e3a257 100644
--- a/src/__tests__/asyncHook.fakeTimers.test.ts
+++ b/src/__tests__/asyncHook.fakeTimers.test.ts
@@ -30,7 +30,7 @@ describe('async hook (fake timers) tests', () => {
     jest.useRealTimers()
   })
 
-  runForRenderers(['default', 'dom', 'native', 'server/hydrated'], ({ renderHook }) => {
+  runForRenderers(['default'], ({ renderHook }) => {
     test('should wait for arbitrary expectation to pass when using advanceTimersByTime()', async () => {
       const { waitFor } = renderHook(() => null)
 
@@ -77,9 +77,9 @@ describe('async hook (fake timers) tests', () => {
 
     test('should waitFor arbitrary expectation to pass when fake timers are not advanced explicitly', async () => {
       const fn = jest.fn().mockReturnValueOnce(false).mockReturnValueOnce(true)
-    
+
       const { waitFor } = renderHook(() => null)
-    
+
       await waitFor(() => {
         expect(fn()).toBe(true)
       })
diff --git a/src/core/asyncUtils.ts b/src/core/asyncUtils.ts
index a62e5940..ab890539 100644
--- a/src/core/asyncUtils.ts
+++ b/src/core/asyncUtils.ts
@@ -14,32 +14,35 @@ const DEFAULT_TIMEOUT = 1000
 const DEFAULT_INTERVAL = 50
 
 function asyncUtils(act: Act, addResolver: (callback: () => void) => void): AsyncUtils {
-  const wait = async (callback: () => boolean | void, { interval, timeout }: WaitOptions) => {
+  const wait = async (
+    callback: () => boolean | void,
+    { interval, timeout }: Required<WaitOptions>
+  ) => {
     const checkResult = () => {
       const callbackResult = callback()
       return callbackResult ?? callbackResult === undefined
     }
 
-    const timeoutSignal = createTimeoutController(timeout as number | boolean, false)
+    const timeoutController = createTimeoutController(timeout, { allowFakeTimers: !interval })
 
     const waitForResult = async () => {
       while (true) {
-        const intervalSignal = createTimeoutController(interval as number | boolean, true)
-        timeoutSignal.onTimeout(() => intervalSignal.cancel())
+        const intervalController = createTimeoutController(interval, { allowFakeTimers: true })
+        timeoutController.onTimeout(() => intervalController.cancel())
 
-        await intervalSignal.wrap(new Promise<void>(addResolver))
+        await intervalController.wrap(new Promise<void>(addResolver))
 
-        if (checkResult() || timeoutSignal.timedOut) {
+        if (checkResult() || timeoutController.timedOut) {
           return
         }
       }
     }
 
     if (!checkResult()) {
-      await act(() => timeoutSignal.wrap(waitForResult()))
+      await act(() => timeoutController.wrap(waitForResult()))
     }
 
-    return !timeoutSignal.timedOut
+    return !timeoutController.timedOut
   }
 
   const waitFor = async (
diff --git a/src/helpers/createTimeoutController.ts b/src/helpers/createTimeoutController.ts
index 033d052f..95accab0 100644
--- a/src/helpers/createTimeoutController.ts
+++ b/src/helpers/createTimeoutController.ts
@@ -1,8 +1,24 @@
 import { jestFakeTimersAreEnabled } from './jestFakeTimersAreEnabled'
 
-function createTimeoutController(timeout: number | boolean, allowFakeTimers: boolean) {
+function createTimeoutController(timeout: number | false, options: { allowFakeTimers: boolean }) {
   let timeoutId: NodeJS.Timeout
   const timeoutCallbacks: Array<() => void> = []
+  let finished = false
+
+  const { allowFakeTimers = false } = options
+
+  const advanceTime = async (currentMs: number) => {
+    if (currentMs < timeout) {
+      jest.advanceTimersByTime(1)
+
+      await Promise.resolve()
+
+      if (finished) {
+        return
+      }
+      await advanceTime(currentMs + 1)
+    }
+  }
 
   const timeoutController = {
     onTimeout(callback: () => void) {
@@ -18,20 +34,24 @@ function createTimeoutController(timeout: number | boolean, allowFakeTimers: boo
             timeoutController.timedOut = true
             timeoutCallbacks.forEach((callback) => callback())
             resolve()
-          }, timeout as number)
+          }, timeout)
 
           if (jestFakeTimersAreEnabled() && allowFakeTimers) {
-            jest.advanceTimersByTime(timeout as number)
+            advanceTime(0)
           }
         }
 
         promise
           .then(resolve)
           .catch(reject)
-          .finally(() => timeoutController.cancel())
+          .finally(() => {
+            finished = true
+            timeoutController.cancel()
+          })
       })
     },
     cancel() {
+      finished = true
       clearTimeout(timeoutId)
     },
     timedOut: false