From 066b3dfc1f891cad09b1889d845f1e4a9995f167 Mon Sep 17 00:00:00 2001
From: Anton Baliasnikov <Anton.baliasnikov@iohk.io>
Date: Wed, 6 Dec 2023 12:37:02 +0000
Subject: [PATCH] fix: add error message parameter to Wait.until function (#12)

Signed-off-by: Anton Baliasnikov <anton.baliasnikov@iohk.io>
---
 .../io/iohk/atala/automation/utils/Wait.kt    | 29 +++++++++++++------
 1 file changed, 20 insertions(+), 9 deletions(-)

diff --git a/src/main/kotlin/io/iohk/atala/automation/utils/Wait.kt b/src/main/kotlin/io/iohk/atala/automation/utils/Wait.kt
index e880117..cd86eb8 100644
--- a/src/main/kotlin/io/iohk/atala/automation/utils/Wait.kt
+++ b/src/main/kotlin/io/iohk/atala/automation/utils/Wait.kt
@@ -1,6 +1,7 @@
 package io.iohk.atala.automation.utils
 
 import org.awaitility.Awaitility
+import org.awaitility.core.ConditionTimeoutException
 import kotlin.time.Duration
 import kotlin.time.Duration.Companion.milliseconds
 import kotlin.time.Duration.Companion.seconds
@@ -12,7 +13,7 @@ object Wait {
      *
      * Usage example:
      * ```
-     * Utils.waitUntil() {
+     * Wait.until() {
      *     customRequest().statusCode == HttpStatus.SC_OK
      * }
      * ```
@@ -20,19 +21,29 @@ object Wait {
      * @param timeout maximum time to wait
      * @param pollInterval polling interval
      * @param condition lambda expression to run the condition
+     * @param errorMessage error message to throw if the condition is not met
      */
     fun until(
         timeout: Duration = 5.seconds,
         pollInterval: Duration = 500.milliseconds,
-        condition: () -> Boolean
+        errorMessage: String? = null,
+        condition: () -> Boolean,
     ) {
-        Awaitility.await()
-            .pollInSameThread()
-            .with()
-            .pollInterval(pollInterval.toJavaDuration())
-            .atMost(timeout.toJavaDuration())
-            .until {
-                condition()
+        try {
+            Awaitility.await()
+                .pollInSameThread()
+                .with()
+                .pollInterval(pollInterval.toJavaDuration())
+                .atMost(timeout.toJavaDuration())
+                .until {
+                    condition()
+                }
+        } catch (err: ConditionTimeoutException) {
+            if (errorMessage != null) {
+                throw ConditionTimeoutException(errorMessage)
+            } else {
+                throw err
             }
+        }
     }
 }