ADFA-6795 | Let chat test runs finish before Main is reset

plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatViewModelCompactionTest.kt

@@ -71,6 +71,9 @@ class ChatViewModelCompactionTest {
 
     private val configScope = CoroutineScope(SupervisorJob() + Dispatchers.Default)
 
+    /** Every ViewModel a test made, cleared before Main is reset; see [clearAndAwait]. */
+    private val viewModels = mutableListOf<ChatViewModel>()
+
     /** The history each model turn was sent with, in order. */
     private val sentHistories = CopyOnWriteArrayList<List<LlmInferenceService.ChatMessage>>()
 
@@ -99,6 +102,7 @@ class ChatViewModelCompactionTest {
 
     @After
     fun tearDown() {
+        clearAndAwait(viewModels)
         SharedServices.clear()
         Dispatchers.resetMain()
         sharedPromptConfig.clear()
@@ -381,7 +385,7 @@ class ChatViewModelCompactionTest {
             initializeStorage(context, TEST_PROJECT_KEY)
             checkBackendAvailability()
             awaitCondition { isBackendAvailable.value }
-        }
+        }.also(viewModels::add)
         return viewModel to service
     }
 

plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatViewModelPromptEditTest.kt

@@ -62,6 +62,9 @@ class ChatViewModelPromptEditTest {
 
     private val configScope = CoroutineScope(SupervisorJob() + Dispatchers.Default)
 
+    /** Every ViewModel a test made, cleared before Main is reset; see [clearAndAwait]. */
+    private val viewModels = mutableListOf<ChatViewModel>()
+
     @Before
     fun setUp() {
         // ChatViewModel's stateIn() calls run on viewModelScope, i.e. Dispatchers.Main.
@@ -87,6 +90,7 @@ class ChatViewModelPromptEditTest {
 
     @After
     fun tearDown() {
+        clearAndAwait(viewModels)
         SharedServices.clear()
         Dispatchers.resetMain()
         sharedPromptConfig.clear()
@@ -384,7 +388,7 @@ class ChatViewModelPromptEditTest {
 
     private fun newViewModel() = ChatViewModel { null }.apply {
         initializeStorage(context, TEST_PROJECT_KEY)
-    }
+    }.also(viewModels::add)
 
     /**
      * A ViewModel whose selected backend is ready and generates by calling [reply] with the

plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatViewModelTitleQueueTest.kt

@@ -62,6 +62,9 @@ class ChatViewModelTitleQueueTest {
     /** Runs the prompt config load that activation starts on device; the title is worded from it. */
     private val configScope = CoroutineScope(SupervisorJob() + Dispatchers.Default)
 
+    /** Every ViewModel a test made, cleared before Main is reset; see [clearAndAwait]. */
+    private val viewModels = mutableListOf<ChatViewModel>()
+
     @Before
     fun setUp() {
         // ChatViewModel's stateIn() calls run on viewModelScope, i.e. Dispatchers.Main.
@@ -90,6 +93,7 @@ class ChatViewModelTitleQueueTest {
 
     @After
     fun tearDown() {
+        clearAndAwait(viewModels)
         Dispatchers.resetMain()
         SharedServices.clear()
         sharedPromptConfig.clear()
@@ -98,7 +102,7 @@ class ChatViewModelTitleQueueTest {
 
     private fun newViewModel() = ChatViewModel { null }.apply {
         initializeStorage(context, TEST_PROJECT_KEY)
-    }
+    }.also(viewModels::add)
 
     /** A ViewModel whose selected backend is ready; its chat turns never answer. */
     private fun newViewModelWithBackend(): ChatViewModel {

plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ViewModelTeardown.kt

@@ -0,0 +1,37 @@
+package com.itsaky.androidide.plugins.aicore.viewmodel
+
+import androidx.lifecycle.ViewModel
+import androidx.lifecycle.ViewModelProvider
+import androidx.lifecycle.ViewModelStore
+import androidx.lifecycle.viewModelScope
+import kotlinx.coroutines.job
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.withTimeout
+
+/** Upper bound on waiting for a cleared ViewModel's coroutines to end. */
+private const val CLEAR_TIMEOUT_MS = 5_000L
+
+/**
+ * Clears each of [viewModels] as its owner would, then waits until every coroutine it started has
+ * ended.
+ *
+ * Call it before `Dispatchers.resetMain()`. A stopped run unwinds on an IO thread and switches to
+ * Main for its cleanup; when that switch overlaps resetMain(), kotlinx-coroutines-test throws
+ * "Dispatchers.Main is used concurrently with setting it" and fails whichever test is tearing
+ * down at that moment.
+ */
+internal fun clearAndAwait(viewModels: Iterable<ViewModel>) {
+    viewModels.forEach { viewModel ->
+        val scopeJob = viewModel.viewModelScope.coroutineContext.job
+        val store = ViewModelStore()
+        ViewModelProvider(store, Owning(viewModel))[viewModel.javaClass]
+        store.clear()
+        runBlocking { withTimeout(CLEAR_TIMEOUT_MS) { scopeJob.join() } }
+    }
+}
+
+/** Hands an existing ViewModel to a store, the public way to put one there. */
+private class Owning(private val viewModel: ViewModel) : ViewModelProvider.Factory {
+    @Suppress("UNCHECKED_CAST")
+    override fun <T : ViewModel> create(modelClass: Class<T>): T = viewModel as T
+}