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 42879a11459..b3cfc1b486f 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 @@ -12,9 +12,9 @@ import kotlinx.coroutines.SupervisorJob import kotlinx.coroutines.cancel import kotlinx.coroutines.delay import kotlinx.coroutines.flow.first -import kotlinx.coroutines.launch import kotlinx.coroutines.runBlocking import kotlinx.coroutines.withTimeout +import java.util.concurrent.CountDownLatch import kotlin.test.AfterTest import kotlin.test.Test import kotlin.test.assertEquals @@ -366,25 +366,26 @@ class KiloBackendAppServiceTest { @Test fun `loading tracks progress through Loading state`() = runBlocking { + val gate = CountDownLatch(1) + mock.responseGate = gate val svc = create() - val states = mutableListOf() - val collector = scope.launch { - svc.appState.collect { states.add(it) } + try { + svc.connect() + + val loading = withTimeout(10_000) { + svc.appState.first { it is KiloAppState.Loading } + } + assertIs(loading) + + gate.countDown() + val ready = withTimeout(10_000) { + svc.appState.first { it is KiloAppState.Ready } + } + assertIs(ready) + } finally { + gate.countDown() } - - svc.connect() - - withTimeout(10_000) { - svc.appState.first { it is KiloAppState.Ready } - } - - collector.cancel() - - // Should have passed through Loading at least once - assertTrue(states.any { it is KiloAppState.Loading }) - // Should have reached Ready - assertTrue(states.any { it is KiloAppState.Ready }) } @Test 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 e98c52c0429..f40013e7ade 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 @@ -81,6 +81,9 @@ class MockCliServer : AutoCloseable { /** Configurable delay for all endpoint responses (ms). 0 = no delay. */ @Volatile var responseDelay: Long = 0 + /** Optional gate for REST responses; SSE stays unblocked so the app can enter Loading. */ + @Volatile var responseGate: CountDownLatch? = null + /** Request counts by bare path (e.g. "/session" or "/global/config"). Thread-safe. */ private val counts = ConcurrentHashMap() @@ -209,6 +212,7 @@ class MockCliServer : AutoCloseable { // Optional delay for race condition testing val delay = responseDelay if (delay > 0) Thread.sleep(delay) + if (bare != "/global/event") responseGate?.await() when { path == "/global/health" -> respond(output, 200, health)