From 899a28a8fd493c4b26b796753a27f09f6fd11a5d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pawe=C5=82=20Bator?=
 <720354+jembezmamy@users.noreply.github.com>
Date: Wed, 6 Nov 2024 13:55:09 +0100
Subject: [PATCH] feat: support assert.expect with reruns

---
 src/assert-result-handler.js | 13 +++++++++++--
 test/retry.js                | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/src/assert-result-handler.js b/src/assert-result-handler.js
index 733acb5..e1aff43 100644
--- a/src/assert-result-handler.js
+++ b/src/assert-result-handler.js
@@ -4,8 +4,11 @@ export default class AssertResultHandler {
   }
 
   get (target, prop, receiver) {
-    if (prop === 'pushResult') {
-      return this.pushResultFn(target)
+    switch (prop) {
+      case 'pushResult':
+        return this.pushResultFn(target)
+      case 'expect':
+        return this.expectFn(target)
     }
     return Reflect.get(target, prop, receiver)
   }
@@ -24,4 +27,10 @@ export default class AssertResultHandler {
       }
     }
   }
+
+  expectFn (target) {
+    return count => {
+      target.expect(count + target.test.assertions.length)
+    }
+  }
 }
diff --git a/test/retry.js b/test/retry.js
index a720a35..64b343e 100644
--- a/test/retry.js
+++ b/test/retry.js
@@ -167,6 +167,41 @@ QUnit.module('default max runs', function () {
   })
 })
 
+QUnit.module('assert.expect', function () {
+  const calls = []
+
+  retry('should pass with retries', function (assert, currentRun) {
+    calls.push(['with', currentRun])
+
+    assert.expect(3)
+    assert.ok(true)
+    if (currentRun === 1) { throw new Error('fail') }
+    assert.ok(true)
+    if (currentRun === 2) { throw new Error('fail') }
+    assert.ok(true)
+  }, 3)
+
+  retry('should pass without retries', function (assert, currentRun) {
+    calls.push(['without', currentRun])
+
+    assert.expect(2)
+    assert.ok(true)
+    assert.ok(true)
+  })
+
+  retry.todo('should fail with incorrect count', function (assert, currentRun) {
+    calls.push(['incorrect', currentRun])
+
+    assert.expect(2)
+    assert.ok(true)
+    if (currentRun === 1) { throw new Error('fail') }
+  })
+
+  QUnit.test('verify calls', function (assert) {
+    assert.deepEqual(calls, [['with', 1], ['with', 2], ['with', 3], ['without', 1], ['incorrect', 1], ['incorrect', 2]])
+  })
+})
+
 QUnit.module('retry.todo', function () {
   const calls = []