diff --git a/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatViewModelCompactionTest.kt b/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatViewModelCompactionTest.kt index a2406644..2a265f0a 100644 --- a/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatViewModelCompactionTest.kt +++ b/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() + /** The history each model turn was sent with, in order. */ private val sentHistories = CopyOnWriteArrayList>() @@ -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 } diff --git a/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatViewModelPromptEditTest.kt b/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatViewModelPromptEditTest.kt index c3882c1d..d78dc846 100644 --- a/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatViewModelPromptEditTest.kt +++ b/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() + @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 diff --git a/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatViewModelTitleQueueTest.kt b/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatViewModelTitleQueueTest.kt index 481c4cf7..f5adbf69 100644 --- a/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ChatViewModelTitleQueueTest.kt +++ b/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() + @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 { diff --git a/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ViewModelTeardown.kt b/plugins/AI-Core/src/test/kotlin/com/itsaky/androidide/plugins/aicore/viewmodel/ViewModelTeardown.kt new file mode 100644 index 00000000..73695c7a --- /dev/null +++ b/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) { + 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 create(modelClass: Class): T = viewModel as T +}