Introduction to unit testing ViewModel with Kotlin Coroutine & Flow

Introduction 

In this blog post, we are going to write unit test of ViewModel which has kotlin coroutine. We will use the example of our previous article Jetpack Compose: MVVM State management in a simple way.

If you haven’t read our previous articles on MVVM and testing, I suggest you read them out.

Sponsored

We are what we repeatedly do. Excellence, then, is not an act, but a habit. Try out Justly and start building your habits today!

Let's Get Started!

Let’s have a quick recap of our our previous article's ViewModel. We have used https://jsonplaceholder.typicode.com/users web API to fetch users. We have created a Sealed class to manage the state of our main screen in a very simple way. Here’s our MainViewModel.

@HiltViewModel
class MainViewModel
@Inject constructor(private val userServices: UserServices) : ViewModel() {

    val state = MutableStateFlow<State>(State.START)

    init {
        loadUser()
    }

    private fun loadUser() = viewModelScope.launch {
        state.value = State.LOADING
        try {
            val users = withContext(Dispatchers.IO) { userServices.getUsers() }
            state.value = State.SUCCESS(users)
        } catch (e: Exception) {
            state.value = State.FAILURE(e.localizedMessage)
        }
    }

}

sealed class State {
    object START : State()
    object LOADING : State()
    data class SUCCESS(val users: List<User>) : State()
    data class FAILURE(val message: String) : State()
}

Pretty straightforward.

In this post, we’re not going to cover any UI-related and business logic of networking calls. Now let’s start writing the unit tests for this ViewModel.

Add required dependencies for testing.

testImplementation 'junit:junit:4.13.2'
testImplementation 'org.mockito.kotlin:mockito-kotlin:4.0.0'
//For runBlockingTest, CoroutineDispatcher etc.
testImplementation 'org.jetbrains.kotlinx:kotlinx-coroutines-test:1.5.2'

Create Test class for MainViewModel

class MainViewModelTest {

    private val userServices = mock<UserServices>()

    private lateinit var viewModel: MainViewModel

    @Before
    fun setup(){
        viewModel = MainViewModel(userServices)
    }

}

Here, we simply mock UserService and create an instance of ViewModel with the required dependency.

Now, Let’s write the first test to verify that the value of our mutable state should be LOADING on loadUser method call whenever ViewModel gets initialized.

class MainViewModelTest {
   
    private val userServices = mock<UserServices>()
   
    private lateinit var viewModel: MainViewModel

    @Test
    fun `Loading state works`() = runBlocking {
        whenever(userServices.getUsers()).thenReturn(emptyList())
        viewModel = MainViewModel(userServices)
        Assert.assertEquals(State.LOADING, viewModel.state.value)
    }

}

When you run this test, the test fails and the error would be something like

This happened because we have tried to run the test in the Dispatcher.Main by using viewModelScope.launch{} which runs on the main thread. And we can’t use the main looper in the unit test.

So now, Let’s create a rule to replace our Dispatcher.Main with TestCoroutineDispatcher during the execution of a test.

@ExperimentalCoroutinesApi
class MainCoroutineRule : TestWatcher(),
                          TestCoroutineScope by TestCoroutineScope() {

    override fun starting(description: Description) {
        super.starting(description)
        Dispatchers.setMain(
            this.coroutineContext[ContinuationInterceptor] as CoroutineDispatcher)
    }

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

We simply set the TestCoroutine dispatcher to our main Coroutine dispatcher. The TestCoroutineScope controls the execution of coroutine within testes.

Let’s add this rule in our ViewModel test.

@get:Rule
val mainCoroutineRule = MainCoroutineRule()

Tada!!. now our test is working fine. Let’s add second test for success response from API.

class MainViewModelTest {

    @get:Rule
    val mainCoroutineRule = MainCoroutineRule()

    private val userServices = mock<UserServices>()

    private lateinit var viewModel: MainViewModel

    @Test
    fun `Success state works`() = runBlocking {
        whenever(userServices.getUsers()).thenReturn(emptyList())
        viewModel = MainViewModel(userServices)
        Assert.assertEquals(State.SUCCESS(emptyList()), viewModel.state.value)
    }
}

Here, getUsers() function is a suspend function, we have to call it from another suspend function or launch it from a coroutine. To mitigate this we have wrapped the call in runBlocking{} This function runs a new Coroutine and blocks the current thread until its completion. runBlocking should only be used for the unit tests not for production code as it will block the current running thread.

Run the test. Opps!! test fails with assertion error something like this

This happens because our test and ViewModel both execute in different thread. To solve this we have to inject our dispatcher to make execution of test and ViewModel in same thread which is recommended as per official doc.

Don’t hardcode Dispatchers when creating new coroutines or calling withContext.

Now Let’s create AppDispatcher, and inject in ViewModel by using dependency injection framework such as hilt.

data class AppDispatchers(
    val IO: CoroutineDispatcher = Dispatchers.IO
)

Let’s modify our viewmodel to use our AppDispatchers

@HiltViewModel
class MainViewModel
@Inject constructor(
    private val userServices: UserServices,
    private val appDispatchers: AppDispatchers,
) : ViewModel() {

    val state = MutableStateFlow<State>(State.START)

    init {
        loadUser()
    }

    private fun loadUser() = viewModelScope.launch {
        state.value = State.LOADING
        try {
            val users = withContext(appDispatchers.IO) { 
                     userServices.getUsers()
            }
            state.value = State.SUCCESS(users)
        } catch (e: Exception) {
            state.value = State.FAILURE(e.localizedMessage)
        }
    }

}

Here, we have replaced withContext(Dispatchers.IO) with our appDispatchers.IO Simple!! isn’t it?

Now, let’s add a test dispatcher for our unit test.

private val testDispatcher = AppDispatchers(
        IO = TestCoroutineDispatcher()
    )

And provide it in ViewModel constructor

class MainViewModelTest {

    @get:Rule
    val mainCoroutineRule = MainCoroutineRule()

    private val userServices = mock<UserServices>()

    private lateinit var viewModel: MainViewModel

    private val testDispatcher = AppDispatchers(
        IO = TestCoroutineDispatcher()
    )
    
    @Test
    fun `Success state works`() = runBlocking {
        whenever(userServices.getUsers()).thenReturn(emptyList())
        viewModel = MainViewModel(userServices, testDispatcher)
        Assert.assertEquals(State.SUCCESS(emptyList()), viewModel.state.value)
    }
}

Run test. It’s should be working now. Let’s add a test for the failure case.

class MainViewModelTest {

    @get:Rule
    val mainCoroutineRule = MainCoroutineRule()

    private val userServices = mock<UserServices>()

    private lateinit var viewModel: MainViewModel

    private val testDispatcher = AppDispatchers(
        IO = TestCoroutineDispatcher(),
        MAIN = Dispatchers.Unconfined
    )

    @Test
    fun `Failure state works`() = runBlocking {
        whenever(userServices.getUsers()).thenThrow(RuntimeException("Error"))
        viewModel = MainViewModel(userServices, testDispatcher)
        Assert.assertEquals(State.FAILURE("Error"), viewModel.state.value)
    }
}

Here, we have thrown RuntimeException whenever getUsers() called.

But, now when you run all tests together, you may find that our Loading state works test is not working. This happens because now our test and ViewModel execute in the same thread so on success response we have updated our mutable state variable with State.Success so we’re not able to verify LOADING state as we stub our API response and it’s executed without delay. So to fix this we have to add some delay in getUser() API call.

whenever(userServices.getUsers()).doSuspendableAnswer {
    withContext(Dispatchers.IO) { delay(5000) }
    emptyList()
}

Here we have added delay of 5000MS. So during this period, we can verify LOADING state. One more thing that needs to notice here is we have set this delay in a separate coroutine thread than our test execution thread by using withContext(Dispatchers.IO) so it’ll not block our current running thread.

Here’s our full Loading test verification test

class MainViewModelTest {

    @get:Rule
    val mainCoroutineRule = MainCoroutineRule()

    private val userServices = mock<UserServices>()

    private lateinit var viewModel: MainViewModel

    private val testDispatcher = AppDispatchers(
        IO = TestCoroutineDispatcher(),
        MAIN = Dispatchers.Unconfined
    )

    @Test
    fun `Loading state works`() = runBlocking {
        whenever(userServices.getUsers()).doSuspendableAnswer {
            withContext(Dispatchers.IO) { delay(5000) }
            emptyList()
        }

        viewModel = MainViewModel(userServices, testDispatcher)
        Assert.assertEquals(State.LOADING, viewModel.state.value)
    }
}

That’s it. Now run all tests again and verify that they all pass.

Related Useful Articles