Address Matt's comments

tensorflow · Aug 6, 2022 · f9af357 · f9af357
1 parent b87bf71
commit f9af357
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 19 deletions.
diff --git 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
@@ -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();