refactor: chai retry plugin

wixplosives · Jul 19, 2023 · e40f364 · e40f364
1 parent ddbb6df
commit e40f364
Show file tree

Hide file tree

Showing 2 changed files with 94 additions and 54 deletions.
diff --git a/packages/testing/src/chai-retry-plugin/helpers.ts b/packages/testing/src/chai-retry-plugin/helpers.ts
@@ -1,39 +1,35 @@
 import Chai from 'chai';
-import * as promiseHelpers from 'promise-assist';
-import { adjustTestTime, mochaCtx } from '../mocha-ctx';
-
 import { chaiMethodsThatHandleFunction } from './constants';
-import type { RetryAndAssertArguments } from './types';
-
-// Add ms to the current test timeout
-const addTimeoutSafetyMargin = (ms: number) => mochaCtx() && adjustTestTime(ms);
-
-function sleepWithSafetyMargin(ms: number): Promise<void> {
-    addTimeoutSafetyMargin(ms);
-    return promiseHelpers.sleep(ms);
-}
+import type { AssertionMethod, RetryAndAssertArguments } from './types';
+import { adjustTestTime } from '../mocha-ctx';
+import { deferred, timeout } from 'promise-assist';
 
-function timeoutWithSafetyMargin(promise: Promise<void>, ms: number, getTimeoutError: () => string): Promise<void> {
-    addTimeoutSafetyMargin(ms);
-    return promiseHelpers.timeout(promise, ms, getTimeoutError);
-}
-
-export const retryFunctionAndAssertions = async (retryAndAssertArguments: RetryAndAssertArguments): Promise<void> => {
+export const retryFunctionAndAssertions = async ({
+    functionToRetry,
+    options,
+    assertionStack,
+    description,
+}: RetryAndAssertArguments): Promise<void> => {
     let assertionError: Error | undefined;
-    let isTimeoutExceeded = false;
+    let didTimeout = false;
+    let cancel = () => void 0;
 
-    const performRetries = async ({
-        functionToRetry,
-        options,
-        assertionStack,
-        description,
-    }: RetryAndAssertArguments) => {
+    function sleep(ms: number) {
+        const { promise, resolve, reject } = deferred();
+        const timeoutId = setTimeout(resolve, ms);
+        cancel = () => {
+            clearTimeout(timeoutId);
+            reject();
+        };
+        return promise;
+    }
+
+    const performRetries = async () => {
         const { retries, delay } = options;
-        let retriesCount = 0;
+        let time = Date.now();
 
-        while (retriesCount < retries && !isTimeoutExceeded) {
+        for (let retriesCount = 0; retriesCount < retries && !didTimeout; retriesCount++) {
             try {
-                retriesCount++;
                 /**
                  * If assertion chain includes such method as `change`, `decrease` or `increase` that means function passed to
                  * the `expect` will be called by Chai itself
@@ -45,35 +41,49 @@ export const retryFunctionAndAssertions = async (retryAndAssertArguments: RetryA
                 let assertion = Chai.expect(valueToAssert, description);
 
                 for (const { propertyName, method, args } of assertionStack) {
-                    if (method && args) {
-                        assertion = method.apply(assertion, args);
-                    } else {
-                        assertion = assertion[propertyName] as Chai.Assertion;
-                    }
+                    assertion = updateAssertion(method, args, assertion, propertyName);
                 }
 
                 return;
-            } catch (error: unknown) {
-                assertionError = error as Error;
-                await sleepWithSafetyMargin(delay);
+            } catch (error: any) {
+                if (!didTimeout) {
+                    assertionError = error as Error;
+                    time = adjustTest(time, delay);
+                    await sleep(delay);
+                }
             }
         }
 
-        if (!isTimeoutExceeded) {
-            throw new Error(`Limit of ${retries} retries exceeded! ${assertionError}`);
-        }
+        throw new Error(`Limit of ${retries} retries exceeded! ${assertionError}`);
     };
 
-    const getTimeoutError = () =>
-        `Timed out after ${retryAndAssertArguments.options.timeout}ms. ${assertionError ?? ''}`;
-
-    setTimeout(() => {
-        isTimeoutExceeded = true;
-    }, retryAndAssertArguments.options.timeout);
+    const getTimeoutError = () => `Timed out after ${options.timeout}ms. ${assertionError ?? ''}`;
 
-    return timeoutWithSafetyMargin(
-        performRetries(retryAndAssertArguments),
-        retryAndAssertArguments.options.timeout,
-        getTimeoutError
-    );
+    return timeout(performRetries(), options.timeout, getTimeoutError).catch((err) => {
+        cancel();
+        didTimeout = true;
+        throw err;
+    });
 };
+
+function updateAssertion(
+    method: AssertionMethod | undefined,
+    args: unknown[] | undefined,
+    assertion: Chai.Assertion,
+    propertyName: keyof Chai.Assertion
+) {
+    if (method && args) {
+        assertion = method.apply(assertion, args);
+    } else {
+        assertion = assertion[propertyName] as Chai.Assertion;
+    }
+    return assertion;
+}
+
+function adjustTest(time: number, delay: number): number {
+    const now = Date.now();
+    const diff = now - time;
+
+    adjustTestTime(diff + delay);
+    return now + delay;
+}
diff --git a/packages/testing/src/test/chai-retry-plugin.unit.ts b/packages/testing/src/test/chai-retry-plugin.unit.ts
@@ -21,18 +21,18 @@ describe('chai-retry-plugin', () => {
         expect(getCallCount()).to.equal(3);
     });
 
-    describe('options should work correctly:', () => {
+    describe('options', () => {
         it('timeout after the specified duration', async () => {
             const funcToRetry = async () => {
-                await sleep(250);
+                await sleep(150);
                 return 'Success';
             };
 
             try {
-                await expect(funcToRetry).retry({ timeout: 700 }).to.equal('Failure');
+                await expect(funcToRetry).retry({ timeout: 100 }).to.equal('Failure');
                 throw new Error('This should not be called');
             } catch (error: unknown) {
-                expect((error as Error).message).includes('Timed out after 700ms');
+                expect((error as Error).message).includes('Timed out after 100ms');
             }
         });
 
@@ -255,6 +255,36 @@ describe('chai-retry-plugin', () => {
             expect(getCallCount()).to.equal(4);
         });
     });
+
+    describe('mocha test timeout adjustment', () => {
+        const BASE_TIMEOUT = 200;
+        const RETRY_TIMEOUT = 100;
+
+        it('upon failure', async function () {
+            this.timeout(BASE_TIMEOUT);
+            try {
+                await expect(() => sleep(20))
+                    .retry({ timeout: RETRY_TIMEOUT })
+                    .to.equal(false);
+            } catch {
+                //
+            }
+            expect(this.timeout(), 'test timeout').to.be.approximately(BASE_TIMEOUT + RETRY_TIMEOUT, 50);
+        });
+
+        it('upon success', async function () {
+            this.timeout(BASE_TIMEOUT);
+            const SUCCESS_TIME = 10;
+            try {
+                await expect(() => sleep(SUCCESS_TIME).then(() => true))
+                    .retry({ timeout: RETRY_TIMEOUT })
+                    .to.equal(true);
+            } catch {
+                //
+            }
+            expect(this.timeout(), 'test timeout').to.be.approximately(BASE_TIMEOUT + SUCCESS_TIME, 10);
+        });
+    });
 });
 
 const withCallCount = (func: (callCount: number) => unknown) => {