From cebe239b902e338630945691f42bd906a8380f05 Mon Sep 17 00:00:00 2001
From: Alex LaFroscia <alex@lafroscia.com>
Date: Mon, 1 Mar 2021 18:59:33 -0500
Subject: [PATCH] fix: re-try assertions that only temporarily throw an
 exception

---
 src/index.ts           | 10 ++++++++++
 tests/wait-for-test.js | 15 ++++++++++++++-
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/src/index.ts b/src/index.ts
index 3714eae..bb02dac 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -16,12 +16,14 @@ export function installWaitFor(QUnit: QUnit) {
       { timeout = 1000 }: Options = { timeout: undefined }
     ) {
       let originalPushResult = this.pushResult;
+      let lastException: unknown;
       let lastResults: Result[];
 
       try {
         await waitUntil(async () => {
           // Reset to capture the most recent round of results
           lastResults = [];
+          lastException = undefined;
 
           // Stub out `pushResult` as late as possible
           originalPushResult = this.pushResult;
@@ -31,6 +33,9 @@ export function installWaitFor(QUnit: QUnit) {
 
           try {
             await assertionCallback();
+          } catch (e) {
+            lastException = e;
+            return false;
           } finally {
             // Restore original `pushResult` implementation right after callback
             this.pushResult = originalPushResult;
@@ -45,6 +50,11 @@ export function installWaitFor(QUnit: QUnit) {
         }
       }
 
+      // If we caught an error on our last attempt, re-throw it
+      if (lastException) {
+        throw lastException;
+      }
+
       for (const result of lastResults) {
         this.pushResult(result);
       }
diff --git a/tests/wait-for-test.js b/tests/wait-for-test.js
index 07b5a36..f52d5b1 100644
--- a/tests/wait-for-test.js
+++ b/tests/wait-for-test.js
@@ -24,7 +24,7 @@ test("it can wait for a condition to be successful", async function (assert) {
   );
 });
 
-test("it throws a non-timeout error", async function (assert) {
+test("it re-throws an exception that is generated on every assertion invocation", async function (assert) {
   assert.expect(2);
 
   const originalPushResult = assert.pushResult;
@@ -47,6 +47,19 @@ test("it throws a non-timeout error", async function (assert) {
   );
 });
 
+test("it re-tries an assertion that throws an exception temporarily", async function (assert) {
+  const stub = td.function();
+  td.when(stub(1)).thenThrow(new Error("Error thrown the first time"));
+  td.when(stub(2)).thenReturn(true);
+
+  let count = 0;
+
+  await assert.waitFor(() => {
+    count++;
+    assert.ok(stub(count));
+  });
+});
+
 test("it can time out while waiting", async function (assert) {
   const stub = td.when(td.function()()).thenReturn(1);