Unit test the new Kotlin coroutine StateFlow

Question

Recently, the class StateFlow was introduced as part of Kotlin coroutines.

I'm currently trying it and encountered an issue while trying to unit test my ViewModel. What I want to achieve: testing that my StateFlow is receiving all the state values in the correct order in my ViewModel.

My code is as follows.

ViewModel:

class WalletViewModel(private val getUserWallets: GetUersWallets) : ViewModel() {

val userWallet: StateFlow<State<UserWallets>> get() = _userWallets
private val _userWallets: MutableStateFlow<State<UserWallets>> =
        MutableStateFlow(State.Init)

fun getUserWallets() {
    viewModelScope.launch {
        getUserWallets.getUserWallets()
            .onStart { _userWallets.value = State.Loading }
            .collect { _userWallets.value = it }
    }
}

My test:

@Test
fun `observe user wallets ok`() = runBlockingTest {
    Mockito.`when`(api.getAssetWallets()).thenReturn(TestUtils.getAssetsWalletResponseOk())
    Mockito.`when`(api.getFiatWallets()).thenReturn(TestUtils.getFiatWalletResponseOk())

    viewModel.getUserWallets()
        
    val res = arrayListOf<State<UserWallets>>()
    viewModel.userWallet.toList(res) //doesn't works

    Assertions.assertThat(viewModel.userWallet.value is State.Success).isTrue() //works, last value enmited
}

Accessing the last value emitted works. But what I want to test is that all the emitted values are emitted in the correct order.

With this piece of code: viewModel.userWallet.toList(res) I'm getting the following error:

java.lang.IllegalStateException: This job has not completed yet
    at kotlinx.coroutines.JobSupport.getCompletionExceptionOrNull(JobSupport.kt:1189)
    at kotlinx.coroutines.test.TestBuildersKt.runBlockingTest(TestBuilders.kt:53)
    at kotlinx.coroutines.test.TestBuildersKt.runBlockingTest$default(TestBuilders.kt:45)
    at WalletViewModelTest.observe user wallets ok(WalletViewModelTest.kt:52)
....

I guess I'm missing something obvious. But not sure why as I'm just getting started with coroutines and Flow and this error seems to happen when not using runBlockingTest, which I use already.

EDIT:

As a temporary solution, I'm testing it as a live data:

@Captor
lateinit var captor: ArgumentCaptor<State<UserWallets>>
    
@Mock
lateinit var walletsObserver: Observer<State<UserWallets>>

@Test
fun `observe user wallets ok`() = runBlockingTest {
    viewModel.userWallet.asLiveData().observeForever(walletsObserver)
    
    viewModel.getUserWallets()
    captor.run {
        Mockito.verify(walletsObserver, Mockito.times(3)).onChanged(capture())
        Assertions.assertThat(allValues[0] is State.Init).isTrue()
        Assertions.assertThat(allValues[1] is State.Loading).isTrue()
        Assertions.assertThat(allValues[2] is State.Success).isTrue()
    }
}

Answer 1

It seems that the Android team changed the API and documentation after this thread. You can check it here: Continuous collection

SharedFlow/StateFlow is a hot flow, and, as described in the docs, A shared flow is called hot because its active instance exists independently of the presence of collectors. It means the scope that launches the collection of your flow won't complete by itself.

To solve this issue, you need to cancel the scope in which collect is called, and as the scope of your test is the test itself, its not ok to cancel the test, so what you need is to launch it in a different job.

@Test
fun `Testing a integer state flow`() = runTest {
    val _intSharedFlow = MutableStateFlow(0)
    val intSharedFlow = _intSharedFlow.asStateFlow()
    val testResults = mutableListOf<Int>()

    val job = launch(UnconfinedTestDispatcher(testScheduler)) {
        intSharedFlow.toList(testResults)
    }
    _intSharedFlow.value = 5

    assertEquals(2, testResults.size)
    assertEquals(0, testResults.first())
    assertEquals(5, testResults.last())
    job.cancel()
}

Improved case: TestScope.backgroundScope ensures that the coroutine gets cancelled before the end of the test.

@Test
fun `Testing an integer state flow`() = runTest {
    val _intSharedFlow = MutableStateFlow(0)
    val intSharedFlow = _intSharedFlow.asStateFlow()
    val testResults = mutableListOf<Int>()

    backgroundScope.launch(UnconfinedTestDispatcher(testScheduler)) {
        intSharedFlow.toList(testResults)
    }
    _intSharedFlow.value = 5

    assertEquals(2, testResults.size)
    assertEquals(0, testResults.first())
    assertEquals(5, testResults.last())
}

A few important things:

1. Always cancel your created job to avoid java.lang.IllegalStateException: This job has not completed yet
2. As this is a StateFlow, when you start collecting (inside toList), you receive the last state. But if you first start collecting and then call your function viewModel.getUserWallets(), the result list will then have all the states, in case you want to test it too.
3. The runTest API changed a little bit, and we need to use UnconfinedTestDispatcher(testScheduler) in the context of the launch call. The documentation says: Notice how UnconfinedTestDispatcher is used for the collecting coroutine here. This ensures that the collecting coroutine is launched eagerly and is ready to receive values after launch returns.

Answer 2

Another way which I derived from this solution in Kotlin coroutines GitHub repository:

@Test fun `The StateFlow should emit all expected values`() = runTest {
    val dispatcher = UnconfinedTestDispatcher(testScheduler)
    val viewModel = MyViewModel(dispatcher)
    val results = mutableListOf<Int>()
    val job = launch(dispatcher) { viewModel.numbers.toList(results) }

    viewModel.addNumber(5)
    viewModel.addNumber(8)
    runCurrent() // Important

    assertThat(results).isEqualTo(listOf(0, 5, 8))
    job.cancel() // Important
}

And this is my ViewModel class:

class MyViewModel(private val dispatcher: CoroutineDispatcher) : ViewModel() {

    private val _numbers = MutableStateFlow(0)
    val numbers: StateFlow<Int> = _numbers

    fun addNumber(number: Int) {
        viewModelScope.launch(dispatcher) {
            _numbers.value = number
        }
    }
}

Please note that I'm using Kotlin 1.6.10 and kotlinx.coroutines-test 1.6.1:

testImplementation("org.jetbrains.kotlinx:kotlinx-coroutines-test:1.6.1")

Also, see the official Kotlin coroutines migration guide to the new test API.

Answer 3

runBlockingTest just skips the delays in your case but not override the dispatcher used in the ViewModel with your test dispatcher. You need to inject TestCoroutineDispatcher to your ViewModel or since you are using viewModelScope.launch {} which already uses Dispatchers.Main by default, you need to override the main dispatcher via Dispatchers.setMain(testCoroutineDispatcher). You can create and add the following rule to your test file.

class MainCoroutineRule(
        val testDispatcher: TestCoroutineDispatcher = TestCoroutineDispatcher()
) : TestWatcher() {

    override fun starting(description: Description?) {
        super.starting(description)
        Dispatchers.setMain(testDispatcher)
    }

    override fun finished(description: Description?) {
        super.finished(description)
        Dispatchers.resetMain()
        testDispatcher.cleanupTestCoroutines()
    }
} 

And in your test file

@get:Rule
var mainCoroutineRule = MainCoroutineRule()

@Test
fun `observe user wallets ok`() = mainCoroutineRule.testDispatcher.runBlockingTest {
}

Btw it is always a good practice to inject dispatchers. For instance if you would have been using a dispatcher other than Dispatchers.Main in your coroutine scope like viewModelScope.launch(Dispatchers.Default), then your test will fail again even if you are using a test dispatcher. The reason is you can only override main dispatcher with Dispatchers.setMain() as it can be understood from its name but not Dispatchers.IO or Dispatchers.Default. In that case you need to inject mainCoroutineRule.testDispatcher to your view model and use the injected dispatcher rather than hardcoding it.

Answer 4

We can create a coroutine for given and one coroutine for whenever

after whenever code, we can use yield so our given code would be done and ready to assert!

to do that you need to extend CouroutinScope as you can see :

done!

• you can use emit instead of tryEmit

Answer 5

The issue you are facing is because toList() needs the flow to complete and "State flow never completes" as per the documentation.

https://kotlin.github.io/kotlinx.coroutines/kotlinx-coroutines-core/kotlinx.coroutines.flow/-state-flow/

https://kotlin.github.io/kotlinx.coroutines/kotlinx-coroutines-core/kotlinx.coroutines.flow/to-list.html

Answer 6

This is what I'm using (without the need to customize the VM dispatcher):

...

@get:Rule
val coroutineRule = MainCoroutineRule()
...

@Test
fun `blablabla`() = runTest {
    val event = mutableListOf<SealedCustomEvent>()
    viewModel.screenEvent
        .onEach { event.add(it) }
        .launchIn(CoroutineScope(UnconfinedTestDispatcher(testScheduler)))
    
    viewModel.onCtaClick()
    advanceUntilIdle()

    Assertions.assertThat(event.last()).isInstanceOf(SealedCustomEvent.OnCtaClick::class.java)

    ...more checks
}

Using launchIn and advanceUntilIdle might solve your testing issues.

Answer 7

Using this with some minor improvement https://github.com/Kotlin/kotlinx.coroutines/issues/3143#issuecomment-1097428912

import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.Job
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.launch
import kotlinx.coroutines.test.TestCoroutineScheduler
import kotlinx.coroutines.test.TestScope
import kotlinx.coroutines.test.UnconfinedTestDispatcher
import org.junit.Assert.assertEquals

@OptIn(ExperimentalCoroutinesApi::class)
/**
 * Test observer for Flow to be able to capture and verify all states.
 */
class TestObserver<T>(
  scope: CoroutineScope,
  testScheduler: TestCoroutineScheduler,
  flow: Flow<T>
) {
  private val values = mutableListOf<T>()

  private val job: Job = scope.launch(UnconfinedTestDispatcher(testScheduler)) {
    flow.collect { values.add(it) }
  }

  /**
   * Assert no values
   */
  fun assertNoValues(): TestObserver<T> {
    assertEquals(emptyList<T>(), this.values)
    return this
  }

  /**
   * Assert the values. Important [TestObserver.finish] needs to be called at the end of the test.
   */
  fun assertValues(vararg values: T): TestObserver<T> {
    assertEquals(values.toList(), this.values)
    return this
  }

  /**
   * Assert the values and finish. Convenient to avoid having to call finish if done last in the test.
   */
  fun assertValuesAndFinish(vararg values: T): TestObserver<T> {
    assertEquals(values.toList(), this.values)
    finish()
    return this
  }

  /**
   * Finish the job
   */
  fun finish() {
    job.cancel()
  }
}

@OptIn(ExperimentalCoroutinesApi::class)
/**
 * Test function for the [TestObserver]
 */
fun <T> Flow<T>.test(
  scope: TestScope
): TestObserver<T> {
  return TestObserver(scope, scope.testScheduler, this)
}

I can now do the following in my test

@Test
fun `test some states`() = runTest {
  val viewModel = ViewModel(
    repository = repository
  )
  val observer = viewModel.state.test(this)
  advanceUntilIdle()
  verify(repository).getData()
  observer.assertValuesAndFinish(
    defaultState,
    defaultState.copy(isLoading = true),
    defaultState.copy(title = "Some title")
  )
}

And my ViewModel

@HiltViewModel
internal class ViewModel @Inject constructor(
  private val repository: Repository
) : ViewModel() {

  private val _state = MutableStateFlow(State())
  val state: StateFlow<State> = _state

  init {
    fetch()
  }

  private fun fetch() {
    _state.value = state.value.copy(
      isLoading = true
    )
    val someData = repository.getData()
    _state.value = state.value.copy(
      isLoading = false,
      title = someData.title
    )
  }
}

Answer 8

This also might help, let's try to mock a full scenario:

// On your VM
class MyViewModel(
    ...
) : ViewModel() {
    val listItems = combine(...) {
        ...
    }.stateIn(
        scope = viewModelScope, 
        started = SharingStarted.WhileSubscribed(), 
        initialValue = null,
    )
}

// On the test file
class MyViewModelTest {
    private val testCoroutineDispatcher = StandardTestDispatcher(TestCoroutineScheduler())

    @BeforeEach
    fun setUp() {
        // in case you're initializing your VM here just make sure
        // to setMain before VM initialization
        Dispatchers.setMain(testCoroutineDispatcher)
    }

    @AfterEach
    fun tearDown() {
        Dispatchers.resetMain()
    }

    @Test
    fun `given X When Y Then Z`() = runTest {
        // Given
        ...

        // When
        // Attach your observer
        val observer = viewModel.listItems.attachFlowObserver(scope = this)
        viewModel.fetchData()
        advanceUntilIdle() // this is needed!

        // Then
        // If you're interested in all values check:
        // observer.getEmittedValues()
        // Do your assertions here
        val results = observer.getLastValue()
        assertEquals(...)

        // Finish the observer
        observer.finish()
    }
}

I hope this is helpful.