From 62ddb32e25140c0e0a83d9a0d0785ca94a4e3afc Mon Sep 17 00:00:00 2001
From: Chris Rohr <51920+chrisrohr@users.noreply.github.com>
Date: Thu, 21 Mar 2024 15:34:48 -0400
Subject: [PATCH] Add support for waitForIt callback to be async (#255)

* Bump to version 0.6.0

Closes #254
---
 __tests__/waitForIt/waitForIt.t.ts |  4 ++++
 package-lock.json                  |  4 ++--
 package.json                       |  2 +-
 src/waitForIt/waitForIt.ts         | 17 ++++++++++++++---
 4 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/__tests__/waitForIt/waitForIt.t.ts b/__tests__/waitForIt/waitForIt.t.ts
index 41c5a4f..a84d0e7 100644
--- a/__tests__/waitForIt/waitForIt.t.ts
+++ b/__tests__/waitForIt/waitForIt.t.ts
@@ -10,6 +10,10 @@ describe("WaitForIt", () => {
     wait().until(() => true);
   });
 
+  it("should resolve when condition passes and callback returns a promise", () => {
+    wait().until(() => Promise.resolve(true));
+  });
+
   it("should reject when condition runs out of tries", async () => {
     await expect(wait().until(() => false)).rejects.toEqual(
       new Error("Condition was not met after 40 tries"),
diff --git a/package-lock.json b/package-lock.json
index 64e1d3f..33fdb66 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "@kiwiproject/kiwi-test-js",
-  "version": "0.5.1",
+  "version": "0.6.0",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "@kiwiproject/kiwi-test-js",
-      "version": "0.5.1",
+      "version": "0.6.0",
       "dependencies": {
         "@elastic/elasticsearch": "8.12.2",
         "@jest/globals": "29.7.0",
diff --git a/package.json b/package.json
index 526a300..391e63b 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@kiwiproject/kiwi-test-js",
-  "version": "0.5.1",
+  "version": "0.6.0",
   "description": "kiwi-test-js is a test utility library. Most of these utilities are ports from the Java Kiwi-test library",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
diff --git a/src/waitForIt/waitForIt.ts b/src/waitForIt/waitForIt.ts
index f5639fa..e595265 100644
--- a/src/waitForIt/waitForIt.ts
+++ b/src/waitForIt/waitForIt.ts
@@ -77,7 +77,7 @@ class WaitFor {
    * will be returned.
    * @param cb The function to check if the condition is true or false.
    */
-  until(cb: () => boolean): Promise<string> {
+  async until(cb: () => boolean | Promise<boolean>): Promise<string> {
     const totalTries = this.timeoutMs / this.pollIntervalMs;
 
     const aliasText = this.alias ? `with alias ${this.alias} ` : "";
@@ -91,8 +91,8 @@ class WaitFor {
     return new Promise((resolve, reject) => {
       let tries = 1;
 
-      const loop = () => {
-        if (cb.apply(this)) {
+      const processResult = (result: boolean) => {
+        if (result) {
           if (this.verbose) {
             console.log(
               `Condition ${aliasText}met after ${tries} of ${totalTries} tries`,
@@ -118,6 +118,17 @@ class WaitFor {
         }
       };
 
+      const loop = async () => {
+        const cbResult = cb.apply(this);
+
+        if (typeof cbResult === "boolean") {
+          processResult(cbResult);
+        } else {
+          const cbPromiseResult = await cbResult;
+          processResult(cbPromiseResult);
+        }
+      };
+
       if (this.pollDelayMs > 0) {
         setTimeout(loop, this.pollDelayMs);
       } else {