support async retry function

coveooss · Aug 2, 2020 · 2211c8e · 2211c8e
1 parent a263496
commit 2211c8e
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 33 deletions.
diff --git a/src/backoff.spec.ts b/src/backoff.spec.ts
@@ -109,33 +109,49 @@ describe("BackOff", () => {
       return request.catch(err => expect(err).toBe(mockFailResponse));
     });
 
-    it("retries the request as many times as specified in #BackOffOptions.numOfAttempts", () => {
-      backOffOptions.numOfAttempts = 2;
+    it("retries the request as many times as specified in #BackOffOptions.numOfAttempts", async () => {
+      const numOfAttemps = 2;
+      backOffOptions.numOfAttempts = numOfAttemps;
       backOffRequest = jest.fn(() => Promise.reject(mockFailResponse));
 
-      const request = initBackOff();
-
-      return request.catch(() => {
-        expect(backOffRequest).toHaveBeenCalledTimes(
-          backOffOptions.numOfAttempts as number
-        );
-      });
+      try {
+        await initBackOff();
+      } catch {
+        expect(backOffRequest).toHaveBeenCalledTimes(numOfAttemps);
+      }
     });
 
-    it(`when the #BackOffOptions.retry method is set to always return false,
-    it only calls request function one time`, () => {
+    it(`when the #BackOffOptions.retry function is set to always return false,
+    it only calls request function one time`, async () => {
       backOffOptions.retry = () => false;
       backOffOptions.numOfAttempts = 2;
       backOffRequest = jest.fn(() => Promise.reject(mockFailResponse));
 
-      const request = initBackOff();
-
-      return request.catch(() =>
-        expect(backOffRequest).toHaveBeenCalledTimes(1)
-      );
+      try {
+        await initBackOff();
+      } catch {
+        expect(backOffRequest).toHaveBeenCalledTimes(1);
+      }
     });
   });
 
+  it("when the #BackOffOptions.retry function returns a promise, it awaits it", async () => {
+    const retryDuration = 100;
+    backOffOptions.retry = () =>
+      new Promise(resolve => setTimeout(() => resolve(true), retryDuration));
+    backOffRequest = promiseThatFailsOnceThenSucceeds();
+
+    const start = Date.now();
+    await initBackOff();
+    const end = Date.now();
+
+    const duration = end - start;
+    const roundedDuration =
+      Math.round(duration / retryDuration) * retryDuration;
+
+    expect(roundedDuration).toBe(retryDuration);
+  });
+
   describe(`when calling #backOff with a function that throws an error the first time, and succeeds the second time`, () => {
     beforeEach(
       () => (backOffRequest = jest.fn(promiseThatFailsOnceThenSucceeds()))
@@ -148,16 +164,14 @@ describe("BackOff", () => {
       );
     });
 
-    it("calls the request function two times", () => {
-      const request = initBackOff();
-      return request.then(() =>
-        expect(backOffRequest).toHaveBeenCalledTimes(2)
-      );
+    it("calls the request function two times", async () => {
+      await initBackOff();
+      expect(backOffRequest).toHaveBeenCalledTimes(2);
     });
 
     it(`when setting the #BackOffOption.timeMultiple to a value,
     when setting the #BackOffOption.delayFirstAttempt to true,
-    it applies a delay between the first and the second call`, () => {
+    it applies a delay between the first and the second call`, async () => {
       const startingDelay = 100;
       const timeMultiple = 3;
       const totalExpectedDelay = startingDelay + timeMultiple * startingDelay;
@@ -166,17 +180,15 @@ describe("BackOff", () => {
       backOffOptions.timeMultiple = timeMultiple;
       backOffOptions.delayFirstAttempt = true;
 
-      const startTime = Date.now();
-      const request = initBackOff();
+      const start = Date.now();
+      await initBackOff();
+      const end = Date.now();
 
-      return request.then(() => {
-        const endTime = Date.now();
-        const duration = endTime - startTime;
-        const roundedDuration =
-          Math.round(duration / startingDelay) * startingDelay;
+      const duration = end - start;
+      const roundedDuration =
+        Math.round(duration / startingDelay) * startingDelay;
 
-        expect(roundedDuration).toBe(totalExpectedDelay);
-      });
+      expect(roundedDuration).toBe(totalExpectedDelay);
     });
   });
 });
diff --git a/src/backoff.ts b/src/backoff.ts
@@ -28,7 +28,7 @@ class BackOff<T> {
         return await this.request();
       } catch (e) {
         this.attemptNumber++;
-        const shouldRetry = this.options.retry(e, this.attemptNumber);
+        const shouldRetry = await this.options.retry(e, this.attemptNumber);
 
         if (!shouldRetry || this.attemptLimitReached) {
           throw e;

diff --git a/src/options.ts b/src/options.ts
@@ -5,7 +5,7 @@ export interface IBackOffOptions {
   jitter: JitterType;
   maxDelay: number;
   numOfAttempts: number;
-  retry: (e: any, attemptNumber: number) => boolean;
+  retry: (e: any, attemptNumber: number) => boolean | Promise<boolean>;
   startingDelay: number;
   timeMultiple: number;
 }