tests: Refactor A2A test, bump Testcontainers, improve Awaitility extension & tests (JetBrains#1338)

Replace Testcontainers JUnit dependency and refactor Docker test setup

1. Replace Testcontainers JUnit dependency and refactor Docker test
setup

- Removed `testcontainers-junit` dependency and replaced usage with
direct `testcontainers` library.
    - Don't build container in Gradle
    - Upgraded `testcontainers` to the latest version.
- Refactored Python A2A server Docker handling into
`TestA2AServerContainer` object for cleaner test configuration.

2. Enhance Awaitility extensions to support return values

- Added support for returning values from `untilAsserted` extensions.
    - Updated test cases to verify return value functionality.

3. Replace direct assertions with Awaitility in A2A tests

- Refactored tests to use Awaitility's coroutine extensions for
asserting eventually consistent conditions.
- Simplified `sendMessageStreaming`, `resubscribeTask`, and notification
configuration logic.
    
4. Use lazy/cached initialization for Docker availability check

## Motivation and Context
A2AClientJsonRpcIntegrationTest is flaky

## Breaking Changes
No

---

#### Type of the changes
- [ ] New feature (non-breaking change which adds functionality)
- [ ] 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
- [x] Refactoring
- [x] Dependency changes

#### 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
- [ ] Tests for the changes have been added
- [ ] All new and existing tests passed

Author: kpavlov

.github/pull_request_template.md

@@ -20,6 +20,8 @@ Also, please tick the appropriate points in the checklist below.
 - [ ] Documentation update
 - [ ] Tests improvement
 - [ ] Refactoring
+- [ ] CI/CD changes
+- [ ] Dependencies update
 
 #### Checklist
 - [ ] The pull request has a description of the proposed change

a2a/a2a-client/build.gradle.kts

@@ -1,7 +1,4 @@
 import ai.koog.gradle.publish.maven.Publishing.publishToMaven
-import org.gradle.internal.os.OperatingSystem
-import org.gradle.kotlin.dsl.support.serviceOf
-import java.io.ByteArrayOutputStream
 
 group = rootProject.group
 version = rootProject.version
@@ -41,7 +38,7 @@ kotlin {
 
                 implementation(libs.ktor.client.cio)
                 implementation(libs.ktor.client.logging)
-                implementation(libs.testcontainers.junit)
+                implementation(libs.testcontainers)
                 runtimeOnly(libs.logback.classic)
             }
         }
@@ -57,36 +54,3 @@ kotlin {
 }
 
 publishToMaven()
-
-tasks.register<Exec>("dockerBuildTestPythonA2AServer") {
-    group = "docker"
-    description = "Build Python A2A test server image"
-    workingDir = file("../test-python-a2a-server")
-    commandLine = listOf("docker", "build", "-t", "test-python-a2a-server", ".")
-
-    onlyIf {
-        // do not attempt to check for docker on windows
-        if (OperatingSystem.current().isWindows) {
-            return@onlyIf false
-        }
-
-        try {
-            val buffer = ByteArrayOutputStream()
-
-            serviceOf<ExecOperations>().exec {
-                commandLine = listOf("docker", "--version")
-                standardOutput = buffer
-                errorOutput = buffer
-            }
-
-            true
-        } catch (_: Exception) {
-            logger.warn("Docker not available. Skipping task 'dockerBuildTestPythonA2AServer'")
-
-            false
-        }
-    }
-}
-tasks.named("jvmTest") {
-    dependsOn("dockerBuildTestPythonA2AServer")
-}

a2a/a2a-client/src/jvmTest/kotlin/ai/koog/a2a/client/A2AClientJsonRpcIntegrationTest.kt

@@ -6,17 +6,13 @@ import ai.koog.test.utils.DockerAvailableCondition
 import io.ktor.client.HttpClient
 import io.ktor.client.plugins.logging.LogLevel
 import io.ktor.client.plugins.logging.Logging
-import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.runBlocking
 import org.junit.jupiter.api.AfterAll
 import org.junit.jupiter.api.BeforeAll
 import org.junit.jupiter.api.TestInstance
 import org.junit.jupiter.api.extension.ExtendWith
 import org.junit.jupiter.api.parallel.Execution
 import org.junit.jupiter.api.parallel.ExecutionMode
-import org.testcontainers.containers.GenericContainer
-import org.testcontainers.containers.wait.strategy.Wait
-import org.testcontainers.junit.jupiter.Container
-import org.testcontainers.junit.jupiter.Testcontainers
 import kotlin.test.Test
 import kotlin.time.Duration.Companion.seconds
 
@@ -26,17 +22,10 @@ import kotlin.time.Duration.Companion.seconds
  * using the JSON-RPC standard.
  */
 @TestInstance(TestInstance.Lifecycle.PER_CLASS)
-@Testcontainers
 @ExtendWith(DockerAvailableCondition::class)
 @Execution(ExecutionMode.SAME_THREAD, reason = "Working with the same instance of test server.")
 class A2AClientJsonRpcIntegrationTest : BaseA2AProtocolTest() {
-    companion object {
-        @Container
-        val testA2AServer: GenericContainer<*> =
-            GenericContainer("test-python-a2a-server")
-                .withExposedPorts(9999)
-                .waitingFor(Wait.forListeningPort())
-    }
+    val testA2AServer = TestA2AServerContainer
 
     override val testTimeout = 10.seconds
 
@@ -47,14 +36,14 @@ class A2AClientJsonRpcIntegrationTest : BaseA2AProtocolTest() {
     }
 
     @Suppress("HttpUrlsUsage")
-    private val agentUrl by lazy { "http://${testA2AServer.host}:${testA2AServer.getMappedPort(9999)}" }
+    private val agentUrl by lazy { "http://${testA2AServer.host}:${testA2AServer.port}" }
 
     private lateinit var transport: HttpJSONRPCClientTransport
 
     override lateinit var client: A2AClient
 
     @BeforeAll
-    fun setUp() = runTest {
+    fun setUp() = runBlocking {
         transport = HttpJSONRPCClientTransport(
             url = agentUrl,
             baseHttpClient = httpClient
@@ -72,8 +61,9 @@ class A2AClientJsonRpcIntegrationTest : BaseA2AProtocolTest() {
     }
 
     @AfterAll
-    fun tearDown() = runTest {
+    fun tearDown() {
         transport.close()
+        testA2AServer.shutdown()
     }
 
     @Test

a2a/a2a-client/src/jvmTest/kotlin/ai/koog/a2a/client/TestA2AServerContainer.kt

@@ -0,0 +1,57 @@
+package ai.koog.a2a.client
+
+import org.slf4j.LoggerFactory
+import org.testcontainers.containers.GenericContainer
+import org.testcontainers.containers.output.Slf4jLogConsumer
+import org.testcontainers.containers.wait.strategy.Wait
+import org.testcontainers.images.builder.ImageFromDockerfile
+import kotlin.io.path.Path
+import kotlin.time.Duration.Companion.seconds
+import kotlin.time.toJavaDuration
+
+/**
+ * Object representing a Dockerized test server for Python-based A2A (Agent-to-Agent) communication.
+ * This server is instantiated using Testcontainers and is primarily utilized for integration testing
+ * purposes, such as validating JSON-RPC HTTP communication in an A2A client context.
+ *
+ * The server is built from a Dockerfile located in the `../test-python-a2a-server` directory and runs
+ * on a predefined exposed port. It provides runtime flexibility to inspect the server's dynamically
+ * assigned host and port.
+ *
+ * This `object` ensures the server is only started once for use in test environments, with the
+ * capability of shutting it down after the tests are completed.
+ *
+ * The server should be initialized before tests, retrieving its host and port, and
+ * using them to configure a test client.
+ */
+object TestA2AServerContainer {
+
+    private const val EXPOSED_PORT = 9999
+    private val STARTUP_TIMEOUT = 20.seconds.toJavaDuration()
+
+    private val image =
+        ImageFromDockerfile("test-python-a2a-server:latest", false) // "false" prevents deleting intermediate images
+            .withFileFromPath(".", Path("../test-python-a2a-server")) // Specify Dockerfile context path
+    private val container: GenericContainer<*> =
+        GenericContainer(image)
+            .withTmpFs(mapOf("/tmp" to "rw,noexec,size=16m"))
+            .withExposedPorts(EXPOSED_PORT)
+            .withReuse(true)
+            .waitingFor(Wait.forListeningPort())
+            .withLogConsumer(Slf4jLogConsumer(LoggerFactory.getLogger(TestA2AServerContainer::class.java)))
+            .withStartupTimeout(STARTUP_TIMEOUT)
+
+    init {
+        container.start()
+    }
+
+    fun shutdown() = runCatching { container.stop() }
+
+    val host: String by lazy {
+        container.host
+    }
+
+    val port: Int by lazy {
+        container.getMappedPort(EXPOSED_PORT)
+    }
+}

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

@@ -5,6 +5,7 @@ import ai.koog.a2a.exceptions.A2AInternalErrorException
 import ai.koog.a2a.model.AgentCapabilities
 import ai.koog.a2a.model.AgentCard
 import ai.koog.a2a.model.AgentSkill
+import ai.koog.a2a.model.Event
 import ai.koog.a2a.model.Message
 import ai.koog.a2a.model.MessageSendConfiguration
 import ai.koog.a2a.model.MessageSendParams
@@ -101,7 +102,7 @@ abstract class BaseA2AProtocolTest {
     }
 
     open fun `test get authenticated extended agent card`() = runTest(timeout = testTimeout) {
-        val request = Request<Nothing?>(data = null)
+        val request = Request(data = null)
 
         val response = client.getAuthenticatedExtendedAgentCard(request)
 
@@ -172,10 +173,10 @@ abstract class BaseA2AProtocolTest {
 
         val response = client.sendMessage(request)
 
-        response should {
-            it.id shouldBe request.id
+        response should { response ->
+            response.id shouldBe request.id
 
-            it.data.shouldBeInstanceOf<Message> {
+            response.data.shouldBeInstanceOf<Message> {
                 it.role shouldBe Role.Agent
                 it.parts shouldBe listOf(TextPart("Hello World"))
                 it.contextId shouldBe "test-context"
@@ -197,19 +198,23 @@ abstract class BaseA2AProtocolTest {
             ),
         )
 
-        val events = client
-            .sendMessageStreaming(createTaskRequest)
-            .toList()
-            .map { it.data }
+        val events: List<Event> = await.untilAsserted(this) {
+            val list = client
+                .sendMessageStreaming(createTaskRequest)
+                .toList()
+                .map { it.data }
 
-        events shouldHaveSize 3
-        events[0].shouldBeInstanceOf<Task> {
-            it.contextId shouldBe "test-context"
-            it.status should {
+            list shouldHaveSize 3
+            return@untilAsserted list
+        }!!
+
+        events[0].shouldBeInstanceOf<Task> { task ->
+            task.contextId shouldBe "test-context"
+            task.status should {
                 it.state shouldBe TaskState.Submitted
             }
 
-            it.history shouldNotBeNull {
+            task.history shouldNotBeNull {
                 this shouldHaveSize 1
 
                 this[0] should {
@@ -344,12 +349,15 @@ abstract class BaseA2AProtocolTest {
             )
         )
 
-        val events = client
-            .resubscribeTask(resubscribeTaskRequest)
-            .toList()
-            .map { it.data }
-
-        events.shouldNotBeEmpty()
+        val events =
+            await.ignoreExceptions().untilAsserted(this) {
+                val list = client
+                    .resubscribeTask(resubscribeTaskRequest)
+                    .toList()
+                    .map { it.data }
+                list.shouldNotBeEmpty()
+                return@untilAsserted list
+            }!!
 
         events.shouldForAll {
             it.shouldBeInstanceOf<TaskStatusUpdateEvent> {
@@ -415,17 +423,21 @@ abstract class BaseA2AProtocolTest {
             )
         )
 
-        val getPushConfigResponse = client.getTaskPushNotificationConfig(getPushConfigRequest)
-        getPushConfigResponse.data shouldBe pushConfig
+        await.untilAsserted(this) {
+            val response = client.getTaskPushNotificationConfig(getPushConfigRequest)
+            response.data shouldBe pushConfig
+        }
 
         val listPushConfigRequest = Request(
             data = TaskIdParams(
                 id = taskId,
             )
         )
 
-        val listPushConfigResponse = client.listTaskPushNotificationConfig(listPushConfigRequest)
-        listPushConfigResponse.data shouldBe listOf(pushConfig)
+        await.untilAsserted(this) {
+            val listPushConfigResponse = client.listTaskPushNotificationConfig(listPushConfigRequest)
+            listPushConfigResponse.data shouldBe listOf(pushConfig)
+        }
 
         val deletePushConfigRequest = Request(
             data = TaskPushNotificationConfigParams(
@@ -436,8 +448,10 @@ abstract class BaseA2AProtocolTest {
 
         client.deleteTaskPushNotificationConfig(deletePushConfigRequest)
 
-        shouldThrowExactly<A2AInternalErrorException> {
-            client.getTaskPushNotificationConfig(getPushConfigRequest)
+        await.untilAsserted(this) {
+            shouldThrowExactly<A2AInternalErrorException> {
+                client.getTaskPushNotificationConfig(getPushConfigRequest)
+            }
         }
     }
 }

a2a/test-python-a2a-server/pyproject.toml

@@ -4,6 +4,6 @@ version = "0.1.0"
 description = "Python A2A server for integration tests."
 requires-python = ">=3.12"
 dependencies = [
-    "a2a-sdk[http-server]==0.3.5",
+    "a2a-sdk[http-server]==0.3.22",
     "uvicorn==0.35.0",
 ]

a2a/test-python-a2a-server/uv.lock

@@ -29,7 +29,7 @@ mockito = "5.21.0"
 mockk = "1.14.7"
 mokksy = "0.5.1"
 mysql = "8.0.33"
-netty = "4.2.6.Final"
+netty = "4.2.9.Final"
 okhttp = "5.3.2"
 opentelemetry = "1.51.0"
 oshai-logging = "7.0.7"
@@ -38,7 +38,7 @@ slf4j = "2.0.17"
 spring-boot = "3.5.9"
 spring-management = "1.1.7"
 sqlite = "3.51.1.0"
-testcontainers = "1.19.7"
+testcontainers = "1.21.4"
 
 [libraries]
 jetbrains-annotations = { module = "org.jetbrains:annotations", version.ref = "annotations" }
@@ -90,7 +90,6 @@ mcp-server = { module = "io.modelcontextprotocol:kotlin-sdk-server", version.ref
 slf4j-simple = { module = "org.slf4j:slf4j-simple", version.ref = "slf4j" }
 jetsign-gradle-plugin = { module = "com.jetbrains:jet-sign", version.ref = "jetsign" }
 testcontainers = { module = "org.testcontainers:testcontainers", version.ref = "testcontainers" }
-testcontainers-junit = { module = "org.testcontainers:junit-jupiter", version.ref = "testcontainers" }
 testcontainers-postgresql = { module = "org.testcontainers:postgresql", version.ref = "testcontainers" }
 testcontainers-mysql = { module = "org.testcontainers:mysql", version.ref = "testcontainers" }
 exposed-core = { module = "org.jetbrains.exposed:exposed-core", version.ref = "exposed" }