diff --git a/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/workspace/KiloBackendWorkspaceTest.kt b/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/workspace/KiloBackendWorkspaceTest.kt index e0d33eeea98..44c2330303e 100644 --- a/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/workspace/KiloBackendWorkspaceTest.kt +++ b/packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/workspace/KiloBackendWorkspaceTest.kt @@ -2,11 +2,15 @@ package ai.kilocode.backend.workspace import ai.kilocode.backend.app.KiloAppState import ai.kilocode.backend.app.KiloBackendAppService +import ai.kilocode.backend.app.KiloBackendSessionManager +import ai.kilocode.backend.app.SseEvent +import ai.kilocode.backend.cli.KiloBackendHttpClients import ai.kilocode.backend.workspace.KiloBackendWorkspace import ai.kilocode.backend.workspace.KiloWorkspaceState import ai.kilocode.backend.testing.FakeCliServer import ai.kilocode.backend.testing.MockCliServer import ai.kilocode.backend.testing.TestLog +import ai.kilocode.jetbrains.api.client.DefaultApi import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.SupervisorJob @@ -14,6 +18,7 @@ import kotlinx.coroutines.async import kotlinx.coroutines.awaitAll import kotlinx.coroutines.cancel import kotlinx.coroutines.delay +import kotlinx.coroutines.flow.MutableSharedFlow import kotlinx.coroutines.flow.first import kotlinx.coroutines.launch import kotlinx.coroutines.runBlocking @@ -388,20 +393,27 @@ class KiloBackendWorkspaceTest { @Test fun `concurrent get for same directory returns same instance`() = runBlocking { - val app = setup() - app.connect() - withTimeout(10_000) { app.appState.first { it is KiloAppState.Ready } } - - // Launch many concurrent get() calls for the same directory - val results = (1..10).map { - async(Dispatchers.Default) { - app.workspaces.get("/same/dir") - } - }.awaitAll() + val port = mock.start() + val http = KiloBackendHttpClients.api(mock.password) + val api = DefaultApi(basePath = "http://127.0.0.1:$port", client = http) + val events = MutableSharedFlow() + val sessions = KiloBackendSessionManager(scope, log) + val manager = KiloBackendWorkspaceManager(scope, sessions, log) + manager.start(api, http, port, events) + + try { + val results = (1..10).map { + async(Dispatchers.Default) { + manager.get("/same/dir") + } + }.awaitAll() - // All must return the exact same instance - val first = results[0] - results.forEach { assertTrue(it === first) } + val first = results[0] + results.forEach { assertTrue(it === first) } + } finally { + manager.stop() + KiloBackendHttpClients.shutdown(http) + } } @Test