From f9af3579826697efc3682b1144a7a60b2b90ec99 Mon Sep 17 00:00:00 2001
From: Yang Gu <yang.gu@intel.com>
Date: Sat, 6 Aug 2022 23:34:34 +0800
Subject: [PATCH] Address Matt's comments

---
 .../src/platforms/platform_browser_test.ts     | 18 +++++-------------
 tfjs-core/src/util_base.ts                     | 10 ++++------
 2 files changed, 9 insertions(+), 19 deletions(-)

diff --git a/tfjs-core/src/platforms/platform_browser_test.ts b/tfjs-core/src/platforms/platform_browser_test.ts
index 97405eaacf6..a4dc4b457dd 100644
--- a/tfjs-core/src/platforms/platform_browser_test.ts
+++ b/tfjs-core/src/platforms/platform_browser_test.ts
@@ -90,7 +90,7 @@ describeWithFlags('PlatformBrowser', BROWSER_ENVS, async () => {
   });
 });
 
-describe('setTimeout', () => {
+describeWithFlags('setTimeout', BROWSER_ENVS, () => {
   const totalCount = 100;
   // Skip the first few samples because the browser does not clamp the timeout
   const skipCount = 5;
@@ -99,11 +99,7 @@ describe('setTimeout', () => {
     let count = 0;
     let startTime = performance.now();
     let totalTime = 0;
-    if (env().platformName === 'browser') {
-      setTimeout(_testSetTimeout, 0);
-    } else {
-      expect().nothing();
-    }
+    setTimeout(_testSetTimeout, 0);
 
     function _testSetTimeout() {
       const endTime = performance.now();
@@ -129,13 +125,9 @@ describe('setTimeout', () => {
     let totalTime = 0;
     let originUseSettimeoutcustom: boolean;
 
-    if (env().platformName === 'browser') {
-      originUseSettimeoutcustom = env().getBool('USE_SETTIMEOUTCUSTOM');
-      env().set('USE_SETTIMEOUTCUSTOM', true);
-      env().platform.setTimeoutCustom(_testSetTimeoutCustom, 0);
-    } else {
-      expect().nothing();
-    }
+    originUseSettimeoutcustom = env().getBool('USE_SETTIMEOUTCUSTOM');
+    env().set('USE_SETTIMEOUTCUSTOM', true);
+    env().platform.setTimeoutCustom(_testSetTimeoutCustom, 0);
 
     function _testSetTimeoutCustom() {
       const endTime = performance.now();
diff --git a/tfjs-core/src/util_base.ts b/tfjs-core/src/util_base.ts
index db1af99de5d..d2e33246f29 100644
--- a/tfjs-core/src/util_base.ts
+++ b/tfjs-core/src/util_base.ts
@@ -303,7 +303,9 @@ export function rightPad(a: string, size: number): string {
 
 export function repeatedTry(
     checkFn: () => boolean, delayFn = (counter: number) => 0,
-    maxCounter?: number, scheduleFn?: Function): Promise<void> {
+    maxCounter?: number,
+    scheduleFn: (functionRef: Function, delay: number) => void =
+        setTimeout): Promise<void> {
   return new Promise<void>((resolve, reject) => {
     let tryCount = 0;
 
@@ -321,11 +323,7 @@ export function repeatedTry(
         reject();
         return;
       }
-      if (typeof scheduleFn === 'undefined') {
-        setTimeout(tryFn, nextBackoff);
-      } else {
-        scheduleFn(tryFn, nextBackoff);
-      }
+      scheduleFn(tryFn, nextBackoff);
     };
 
     tryFn();