feat: Add Awaitility extensions for coroutine support in test-utils (#1331)

# feat: Add Awaitility extensions for coroutine support in test-utils

- Introduced custom Awaitility extensions to support suspending
assertions.
- Updated `a2a-test` to use new `test-utils` dependencies for Awaitility
utilities.
- Upgraded `mockk` and `mockito` dependencies to newer versions.

## Motivation and Context
[A2AServerJsonRpcIntegrationTest.kt#L74 is
flaky](https://github.com/JetBrains/koog/pull/1116/files#annotation_43754744518)
It must use awatility to verify eventually consistent conditions

## Breaking Changes
No

---

#### Type of the changes
- [x] New feature (non-breaking change which adds functionality)
- [x] Bug fix (non-breaking change which fixes an issue)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [x] Tests improvement
- [ ] Refactoring

#### Checklist
- [x] The pull request has a description of the proposed change
- [x] I read the [Contributing
Guidelines](https://github.com/JetBrains/koog/blob/main/CONTRIBUTING.md)
before opening the pull request
- [x] The pull request uses **`develop`** as the base branch
- [x] Tests for the changes have been added
- [x] All new and existing tests passed

##### Additional steps for pull requests adding a new feature
- [ ] An issue describing the proposed change exists
- [ ] The pull request includes a link to the issue
- [ ] The change was discussed and approved in the issue
- [ ] Docs have been added / updated

Author: kpavlov

a2a/a2a-test/build.gradle.kts

@@ -21,7 +21,7 @@ kotlin {
 
         jvmMain {
             dependencies {
-                api(kotlin("test-junit5"))
+                api(project(":test-utils"))
             }
         }
 

a2a/a2a-test/src/jvmMain/kotlin/ai/koog/a2a/test/BaseA2AProtocolTest.kt

@@ -21,6 +21,7 @@ import ai.koog.a2a.model.TaskStatusUpdateEvent
 import ai.koog.a2a.model.TextPart
 import ai.koog.a2a.model.TransportProtocol
 import ai.koog.a2a.transport.Request
+import ai.koog.test.utils.untilAsserted
 import io.kotest.assertions.throwables.shouldThrowExactly
 import io.kotest.inspectors.shouldForAll
 import io.kotest.matchers.collections.shouldHaveSize
@@ -32,6 +33,7 @@ import io.kotest.matchers.string.shouldStartWith
 import io.kotest.matchers.types.shouldBeInstanceOf
 import kotlinx.coroutines.flow.toList
 import kotlinx.coroutines.test.runTest
+import org.awaitility.kotlin.await
 import kotlin.time.Duration
 import kotlin.uuid.ExperimentalUuidApi
 import kotlin.uuid.Uuid
@@ -265,15 +267,19 @@ abstract class BaseA2AProtocolTest {
             )
         )
 
-        val response = client.getTask(getTaskRequest)
+        await
+            .ignoreExceptions()
+            .untilAsserted(this) {
+                val response = client.getTask(getTaskRequest)
 
-        response.data should {
-            it.id shouldBe taskId
-            it.contextId shouldBe "test-context"
-            it.status should {
-                it.state shouldBe TaskState.Completed
+                response.data should { task ->
+                    task.id shouldBe taskId
+                    task.contextId shouldBe "test-context"
+                    task.status should { status ->
+                        status.state shouldBe TaskState.Completed
+                    }
+                }
             }
-        }
     }
 
     open fun `test cancel task`() = runTest(timeout = testTimeout) {

gradle/libs.versions.toml

@@ -25,8 +25,8 @@ ktor3 = "3.2.2"
 lettuce = "6.5.5.RELEASE"
 logback = "1.5.13"
 mcp = "0.8.1"
-mockito = "5.19.0"
-mockk = "1.13.8"
+mockito = "5.21.0"
+mockk = "1.14.7"
 mokksy = "0.5.1"
 mysql = "8.0.33"
 netty = "4.2.6.Final"

test-utils/build.gradle.kts

@@ -23,6 +23,7 @@ kotlin {
                 api(kotlin("test-junit5"))
                 api(libs.junit.jupiter.params)
                 api(libs.testcontainers)
+                api(libs.awaitility)
                 runtimeOnly(libs.slf4j.simple)
             }
         }

test-utils/src/jvmMain/kotlin/ai/koog/test/utils/AwaitilityExtensions.kt

@@ -0,0 +1,79 @@
+package ai.koog.test.utils
+
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.runBlocking
+import org.awaitility.core.ConditionFactory
+import kotlin.coroutines.ContinuationInterceptor
+import kotlin.time.Duration
+import kotlin.time.toJavaDuration
+
+/**
+ * Repeatedly evaluates the given block until it does not throw any exceptions,
+ * providing support for asynchronous or coroutine-based operations. This method
+ * is particularly helpful in cases where conditions are expected to eventually
+ * become true due to the nature of concurrent or delayed behavior.
+ *
+ * @param scope The [CoroutineScope] used to provide the context for the suspendable block.
+ * @param block A suspendable lambda function containing the assertions or conditions to be verified.
+ */
+public fun ConditionFactory.untilAsserted(scope: CoroutineScope, block: suspend () -> Unit) {
+    val self: ConditionFactory = this
+    self.untilAsserted {
+        runBlocking(scope.coroutineContext.minusKey(ContinuationInterceptor)) {
+            block()
+        }
+    }
+}
+
+/**
+ * Repeatedly evaluates the given block until it does not throw any exceptions.
+ * This method is useful for asserting conditions that may eventually become true,
+ * particularly in scenarios involving asynchronous or concurrent operations.
+ *
+ * @param block A suspendable lambda function containing the assertions or conditions to be verified.
+ */
+public fun ConditionFactory.untilAsserted(block: suspend () -> Unit) {
+    val self: ConditionFactory = this
+    self.untilAsserted {
+        runBlocking {
+            block()
+        }
+    }
+}
+
+/**
+ * Repeatedly evaluates the given block until it does not throw any exceptions.
+ * This method is useful for asserting conditions that may eventually become true,
+ * particularly in scenarios involving asynchronous or concurrent operations.
+ *
+ * @param block A suspendable lambda function containing the assertions or conditions to be verified.
+ */
+public fun ConditionFactory.atMost(duration: Duration): ConditionFactory =
+    this.atMost(duration.toJavaDuration())
+
+/**
+ * Specifies the minimum amount of time that the condition should be evaluated for.
+ *
+ * @param duration The minimum duration to evaluate the condition, represented as a [Duration].
+ * @return The updated [ConditionFactory] instance with the specified minimum duration applied.
+ */
+public fun ConditionFactory.atLeast(duration: Duration): ConditionFactory =
+    this.atLeast(duration.toJavaDuration())
+
+/**
+ * Specifies the delay before polling begins when evaluating a condition.
+ *
+ * @param duration The duration to wait before the first polling attempt, represented as a [Duration].
+ * @return The updated [ConditionFactory] instance with the specified polling delay applied.
+ */
+public fun ConditionFactory.pollDelay(duration: Duration): ConditionFactory =
+    this.pollDelay(duration.toJavaDuration())
+
+/**
+ * Specifies the interval between consecutive polling attempts when evaluating a condition.
+ *
+ * @param duration The interval duration between polling attempts, represented as a [Duration].
+ * @return The updated [ConditionFactory] instance with the specified polling interval applied.
+ */
+public fun ConditionFactory.pollInterval(duration: Duration): ConditionFactory =
+    this.pollInterval(duration.toJavaDuration())

test-utils/src/jvmTest/kotlin/ai/koog/test/utils/AwaitilityExtensionsTest.kt

@@ -0,0 +1,80 @@
+package ai.koog.test.utils
+
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.test.runTest
+import org.awaitility.Awaitility.await
+import org.awaitility.core.ConditionTimeoutException
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.assertThrows
+import java.time.Duration
+import java.util.concurrent.atomic.AtomicInteger
+import kotlin.test.fail
+import kotlin.time.Duration.Companion.milliseconds
+import kotlin.time.Duration.Companion.seconds
+
+class AwaitilityExtensionsTest {
+
+    @Test
+    fun `untilAsserted with suspend block eventually succeeds`() = runTest {
+        val counter = AtomicInteger(0)
+
+        await()
+            .atMost(2.seconds)
+            .untilAsserted(block = {
+                delay(30)
+                val current = counter.incrementAndGet()
+                if (current < 5) {
+                    throw AssertionError("Not enough yet: $current")
+                }
+            })
+
+        assertEquals(5, counter.get())
+    }
+
+    @Test
+    fun `untilAsserted with suspend block fails on timeout`() = runTest {
+        assertThrows<ConditionTimeoutException> {
+            await()
+                .atMost(200.milliseconds)
+                .pollInterval(10.milliseconds)
+                .untilAsserted(block = {
+                    delay(50)
+                    fail("Always failing")
+                })
+        }
+    }
+
+    @Test
+    fun `untilAsserted with scope and suspend block should eventually succeed`() = runTest {
+        val counter = AtomicInteger(0)
+
+        await()
+            .atMost(2.seconds)
+            .atLeast(40.milliseconds)
+            .pollInterval(10.milliseconds)
+            .untilAsserted(this) {
+                val current = counter.incrementAndGet()
+                if (current < 5) {
+                    throw AssertionError("Not enough yet: $current")
+                }
+            }
+
+        assertEquals(5, counter.get())
+    }
+
+    @Test
+    fun `untilAsserted with scope and suspend block fails on timeout`() = runTest {
+        assertThrows<ConditionTimeoutException> {
+            await()
+                .pollDelay(0.milliseconds)
+                .atLeast(10.milliseconds)
+                .atMost(200.milliseconds)
+                .pollInterval(Duration.ofMillis(10))
+                .untilAsserted(this) {
+                    delay(50)
+                    throw AssertionError("Always failing")
+                }
+        }
+    }
+}