Kotlin Coroutine Unit Test Flow collection with viewModelScope - kotlin

Question

I want to test a method of my ViewModel that collects a Flow. Inside the collector a LiveData object is mutated, which I want to check in the end. This is roughly how the setup looks:
//Outside viewmodel
val f = flow { emit("Test") }.flowOn(Dispatchers.IO)
//Inside viewmodel
val liveData = MutableLiveData<String>()
fun action() {
viewModelScope.launch { privateAction() }
}
suspend fun privateAction() {
f.collect {
liveData.value = it
}
}
When I now call the action() method in my unit test, the test finishes before the flow is collected. This is how the test might look:
#Test
fun example() = runBlockingTest {
viewModel.action()
assertEquals(viewModel.liveData.value, "Test")
}
I am using the TestCoroutineDispatcher via this Junit5 extension and also the instant executor extension for LiveData:
class TestCoroutineDispatcherExtension : BeforeEachCallback, AfterEachCallback, ParameterResolver {
#SuppressLint("NewApi") // Only used in unit tests
override fun supportsParameter(parameterContext: ParameterContext?, extensionContext: ExtensionContext?): Boolean {
return parameterContext?.parameter?.type === testDispatcher.javaClass
}
override fun resolveParameter(parameterContext: ParameterContext?, extensionContext: ExtensionContext?): Any {
return testDispatcher
}
private val testDispatcher = TestCoroutineDispatcher()
override fun beforeEach(context: ExtensionContext?) {
Dispatchers.setMain(testDispatcher)
}
override fun afterEach(context: ExtensionContext?) {
Dispatchers.resetMain()
testDispatcher.cleanupTestCoroutines()
}
}
class InstantExecutorExtension : BeforeEachCallback, AfterEachCallback {
override fun beforeEach(context: ExtensionContext?) {
ArchTaskExecutor.getInstance()
.setDelegate(object : TaskExecutor() {
override fun executeOnDiskIO(runnable: Runnable) = runnable.run()
override fun postToMainThread(runnable: Runnable) = runnable.run()
override fun isMainThread(): Boolean = true
})
}
override fun afterEach(context: ExtensionContext?) {
ArchTaskExecutor.getInstance().setDelegate(null)
}
}

Answer 1

You can try either,
fun action() = viewModelScope.launch { privateAction() }
suspend fun privateAction() {
f.collect {
liveData.value = it
}
}
#Test
fun example() = runBlockingTest {
viewModel.action().join()
assertEquals(viewModel.liveData.value, "Test")
}
or
fun action() {
viewModelScope.launch { privateAction()
}
suspend fun privateAction() {
f.collect {
liveData.value = it
}
}
#Test
fun example() = runBlockingTest {
viewModel.action()
viewModel.viewModelScope.coroutineContext[Job]!!.join()
assertEquals(viewModel.liveData.value, "Test")
}
You could also try this,
suspend fun <T> LiveData<T>.awaitValue(): T? {
return suspendCoroutine { cont ->
val observer = object : Observer<T> {
override fun onChanged(t: T?) {
removeObserver(this)
cont.resume(t)
}
}
observeForever(observer)
}
}
#Test
fun example() = runBlockingTest {
viewModel.action()
assertEquals(viewModel.liveData.awaitValue(), "Test")
}

Answer 2

So what I ended up doing is just passing the Dispatcher to the viewmodel constructor:
class MyViewModel(..., private val dispatcher = Dispatchers.Main)
and then using it like this:
viewModelScope.launch(dispatcher) {}
So now I can override this when I instantiate the ViewModel in my test with a TestCoroutineDispatcher and then advance the time, use testCoroutineDispatcher.runBlockingTest {}, etc.