diff --git a/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/app/KiloBackendAppServiceTest.kt b/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/app/KiloBackendAppServiceTest.kt index f63b25cb5d9..c0101b0e7cd 100644 --- a/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/app/KiloBackendAppServiceTest.kt +++ b/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/app/KiloBackendAppServiceTest.kt @@ -285,10 +285,12 @@ class KiloBackendAppServiceTest { svc.connect() withTimeout(10_000) { - svc.appState.first { it is KiloAppState.Ready } + svc.appState.first { state -> + state is KiloAppState.Ready && state.data.warnings.any { it.path == ".kilo/kilo.json" } + } } - assertTrue(log.messages.any { + assertTrue(log.awaitMessage { it.contains("App warnings:") && it.contains(".kilo/kilo.json: Invalid JSON") }) } @@ -516,13 +518,14 @@ class KiloBackendAppServiceTest { // Change the config response and push an SSE event mock.config = """{"model":"updated"}""" + val before = mock.requestCount("/global/config") mock.awaitSseConnection() mock.pushEvent("global.config.updated", """{"type":"global.config.updated"}""") - // Wait for config to be refreshed + assertTrue(mock.awaitRequestCount("/global/config", before + 1)) withTimeout(5_000) { - while (svc.config?.model != "updated") { - delay(100) + svc.appState.first { state -> + state is KiloAppState.Ready && state.data.config.model == "updated" } } @@ -542,12 +545,14 @@ class KiloBackendAppServiceTest { assertEquals(1, (svc.appState.value as KiloAppState.Ready).data.warnings.size) mock.warnings = "[]" + val before = mock.requestCount("/config/warnings") mock.awaitSseConnection() mock.pushEvent("global.config.updated", """{"type":"global.config.updated"}""") + assertTrue(mock.awaitRequestCount("/config/warnings", before + 1)) withTimeout(5_000) { - while ((svc.appState.value as? KiloAppState.Ready)?.data?.warnings?.isNotEmpty() == true) { - delay(100) + svc.appState.first { state -> + state is KiloAppState.Ready && state.data.warnings.isEmpty() } } diff --git a/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/app/KiloConnectionServiceTest.kt b/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/app/KiloConnectionServiceTest.kt index ee55db274b3..335b91e2e87 100644 --- a/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/app/KiloConnectionServiceTest.kt +++ b/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/app/KiloConnectionServiceTest.kt @@ -8,11 +8,11 @@ import ai.kilocode.backend.testing.MockCliServer import ai.kilocode.backend.testing.TestLog import ai.kilocode.log.KiloLog import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.CoroutineStart import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.SupervisorJob import kotlinx.coroutines.async import kotlinx.coroutines.cancel -import kotlinx.coroutines.delay import kotlinx.coroutines.flow.first import kotlinx.coroutines.flow.take import kotlinx.coroutines.flow.toList @@ -89,13 +89,11 @@ class KiloConnectionServiceTest { } // Use first{} on the flow to capture the event — avoids race with SharedFlow subscription - val deferred = scope.launch { + val deferred = scope.launch(start = CoroutineStart.UNDISPATCHED) { val event = svc.events.first { it.type == "global.config.updated" } assertEquals("global.config.updated", event.type) } - // Small delay to ensure the collector subscription is active - delay(200) mock.pushEvent("global.config.updated", """{"type":"global.config.updated"}""") withTimeout(5_000) { @@ -114,13 +112,12 @@ class KiloConnectionServiceTest { } val count = 100 - val received = async { + val received = async(start = CoroutineStart.UNDISPATCHED) { withTimeout(5_000) { svc.events.take(count).toList() } } - delay(200) repeat(count) { idx -> mock.pushEvent("test.event", idx.toString()) } @@ -144,18 +141,23 @@ class KiloConnectionServiceTest { @Suppress("UNCHECKED_CAST") val ref = sourceField.get(svc) as AtomicReference ref.set(source) - val received = scope.async { + val received = scope.async(start = CoroutineStart.UNDISPATCHED) { withTimeout(5_000) { svc.events.take(2).toList() } } - delay(200) + val entered = CountDownLatch(1) val first = Thread { listener.onEvent(source, null, "first.event", "first") } - val second = Thread { listener.onEvent(source, null, "second.event", "second") } + val second = Thread { + entered.countDown() + listener.onEvent(source, null, "second.event", "second") + } first.start() assertTrue(blocked.started.await(1, TimeUnit.SECONDS)) second.start() + assertTrue(entered.await(1, TimeUnit.SECONDS)) + blocked.release.countDown() first.join() second.join() @@ -247,9 +249,11 @@ class KiloConnectionServiceTest { } // Restart should tear down and re-open + val seen = mock.sseConnectionCount svc.restart() // Wait for reconnection + assertTrue(mock.awaitSseConnections(seen + 1)) withTimeout(10_000) { svc.state.first { it is ConnectionState.Connected } } @@ -281,13 +285,14 @@ class KiloConnectionServiceTest { private class BlockingLog : KiloLog { val started = CountDownLatch(1) + val release = CountDownLatch(1) override var isDebugEnabled: Boolean = true override fun debug(block: () -> String) { val msg = block() if (!msg.contains("evt=first.event")) return started.countDown() - Thread.sleep(250) + release.await(1, TimeUnit.SECONDS) } override fun info(msg: String) {} diff --git a/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/testing/MockCliServer.kt b/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/testing/MockCliServer.kt index 12ff956be7e..1913c7d2320 100644 --- a/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/testing/MockCliServer.kt +++ b/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/testing/MockCliServer.kt @@ -99,10 +99,40 @@ class MockCliServer : AutoCloseable { /** Request counts by bare path (e.g. "/session" or "/global/config"). Thread-safe. */ private val counts = ConcurrentHashMap() + private val requests = Object() + private val streams = Object() + private val sse = AtomicInteger(0) + + val sseConnectionCount: Int + get() = sse.get() /** Return the number of requests received for [path] (bare, no query). */ fun requestCount(path: String): Int = counts[path]?.get() ?: 0 + fun awaitRequestCount(path: String, target: Int, timeout: Long = 5_000): Boolean { + val end = System.currentTimeMillis() + timeout + synchronized(requests) { + while (requestCount(path) < target) { + val wait = end - System.currentTimeMillis() + if (wait <= 0) return false + requests.wait(wait) + } + return true + } + } + + fun awaitSseConnections(target: Int, timeout: Long = 5_000): Boolean { + val end = System.currentTimeMillis() + timeout + synchronized(streams) { + while (sse.get() < target) { + val wait = end - System.currentTimeMillis() + if (wait <= 0) return false + streams.wait(wait) + } + return true + } + } + @Volatile var lastExperimentalSessionPath: String? = null /** Reset all request counters. */ @@ -127,7 +157,8 @@ class MockCliServer : AutoCloseable { // Clean up any previous instance shutdownServer() - sseLatch = CountDownLatch(1) + val latch = CountDownLatch(1) + sseLatch = latch sseConnected = CountDownLatch(1) sseWriter = null @@ -218,9 +249,11 @@ class MockCliServer : AutoCloseable { val output = BufferedWriter(OutputStreamWriter(socket.getOutputStream())) val bare = path.substringBefore("?") + val latch = sseLatch // Track request counts counts.computeIfAbsent(bare) { AtomicInteger(0) }.incrementAndGet() + synchronized(requests) { requests.notifyAll() } // Optional delay for race condition testing val delay = responseDelay @@ -255,7 +288,7 @@ class MockCliServer : AutoCloseable { lastOrganizationSetBody = body respond(output, organizationSetStatus, "true") } - path == "/global/event" -> handleSse(output) + path == "/global/event" -> handleSse(output, latch) path == "/path" -> respond(output, 200, this.path) bare == "/provider" -> respond(output, providersStatus, providers) bare == "/agent" -> respond(output, agentsStatus, agents) @@ -319,7 +352,7 @@ class MockCliServer : AutoCloseable { writer.flush() } - private fun handleSse(writer: BufferedWriter) { + private fun handleSse(writer: BufferedWriter, latch: CountDownLatch) { writer.write("HTTP/1.1 200 OK\r\n") writer.write("Content-Type: text/event-stream\r\n") writer.write("Cache-Control: no-cache\r\n") @@ -327,8 +360,10 @@ class MockCliServer : AutoCloseable { writer.write("\r\n") writer.flush() sseWriter = writer + sse.incrementAndGet() + synchronized(streams) { streams.notifyAll() } sseConnected.countDown() // Block until SSE is closed or server shuts down - sseLatch.await() + latch.await() } } diff --git a/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/testing/TestLog.kt b/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/testing/TestLog.kt index 1219a20fbd8..b9efae14924 100644 --- a/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/testing/TestLog.kt +++ b/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/testing/TestLog.kt @@ -7,31 +7,51 @@ import ai.kilocode.log.KiloLog */ class TestLog : KiloLog { private val items = mutableListOf() + private val lock = Object() val messages: List - get() = synchronized(items) { items.toList() } + get() = synchronized(lock) { items.toList() } override var isDebugEnabled: Boolean = true + fun awaitMessage(timeout: Long = 5_000, predicate: (String) -> Boolean): Boolean { + val end = System.currentTimeMillis() + timeout + synchronized(lock) { + while (items.none(predicate)) { + val wait = end - System.currentTimeMillis() + if (wait <= 0) return false + lock.wait(wait) + } + return true + } + } + override fun debug(block: () -> String) { if (!isDebugEnabled) return val msg = block() - synchronized(items) { items.add("DEBUG: $msg") } + add("DEBUG: $msg") println("[test] DEBUG: $msg") } override fun info(msg: String) { - synchronized(items) { items.add("INFO: $msg") } + add("INFO: $msg") println("[test] INFO: $msg") } override fun warn(msg: String, t: Throwable?) { - synchronized(items) { items.add("WARN: $msg") } + add("WARN: $msg") println("[test] WARN: $msg") t?.printStackTrace() } override fun error(msg: String, t: Throwable?) { - synchronized(items) { items.add("ERROR: $msg") } + add("ERROR: $msg") System.err.println("[test] ERROR: $msg") t?.printStackTrace() } + + private fun add(msg: String) { + synchronized(lock) { + items.add(msg) + lock.notifyAll() + } + } }