Android Architecture Components (AAC) provides several libraries that help you design robust, testable, and maintainable apps. One of the most important components of AAC is the ViewModel, which provides a lifecycle-aware container for UI-related data that survives configuration changes. The ViewModel can also be responsible for managing long-running tasks such as network requests, database queries, and other I/O operations.

One of the challenges of working with coroutines in Android is how to properly manage their lifecycle in a testable way. One popular approach is to use the ViewModelScope, which is a lifecycle-aware scope that automatically cancels all coroutines when the associated ViewModel is destroyed. However, using ViewModelScope can make it difficult to test coroutine code in isolation because the coroutines are tightly coupled to the lifecycle of the ViewModel and the ViewModelScope isn't aware about the structured concurrency within your runTest block (from the coroutines-test package).

In this blog post, we'll discuss why replacing ViewModelScope within AAC ViewModels with a custom solution may be a good idea, and how you can implement a better solution that allows for better testing of coroutine code (which also allows us to move away from having to set Dispatchers.setMain().

The Problem with ViewModelScope

When you use ViewModelScope to launch a coroutine, the coroutine is automatically tied to the lifecycle of the ViewModel. This means that if the ViewModel is destroyed, all running coroutines are canceled. While this behavior is useful in many cases, it can also make it difficult to test coroutine code in isolation. Additionally, the ViewModelScope isn't aware of the structured concurrency that exists within your runTest block, which can make it difficult to define the behavior we want for testing coroutines.

A Custom Solution

The solution we will be striving towards today involves using the ViewModelScope via Hilt in production code, and using the backgroundScope provided to us via TestScope.

The first thing we need to do is create the CoroutineScope we will be providing via Hilt to use in our ViewModel. This CloseableCoroutineScope's Job still needs to be canceled when the ViewModelScope is cleared, which is why we cancel it in onCleared()

class CloseableCoroutineScope(context: CoroutineContext) : CoroutineScope, RetainedLifecycle.OnClearedListener {
    override val coroutineContext: CoroutineContext = context

    override fun onCleared() {
        coroutineContext.cancel()
    }
}

Then, via our Hilt module, we can provide this scope in the ViewModelComponent. It is important to note, that we are also calling addOnClearedListener here, so the scope is canceled within the ViewModelLifecycle

@Module
@InstallIn(ViewModelComponent::class)
internal object ViewModelScopeModule {
    @Provides
    @ViewModelScoped
    fun provideViewModelCoroutineScope(lifecycle: ViewModelLifecycle): CoroutineScope {
        return CloseableCoroutineScope(SupervisorJob()).also { closeableCoroutineScope ->
            lifecycle.addOnClearedListener(closeableCoroutineScope)
        }
    }
}

Testing

Now that we have the basic setup in place, we can finally use it! First, we will create a ViewModel that takes in our new scope.

@HiltViewModel
class MainViewModel @Inject constructor(
    private val scope: CoroutineScope,
) : ViewModel()

Then, when we want to use that scope, we can do so via:

@HiltViewModel
class MainViewModel @Inject constructor(
    private val scope: CoroutineScope,
    private val apiService: ApiService,
) : ViewModel() {

    private val _result = MutableSharedFlow<Int>()
    val result = _result.asSharedFlow()

    fun getResult() {
        scope.launch {
            _result.emit(apiService.getResult())
        }
    }
}

The testing of this is really the impetus behind why we want to avoid ViewModelScope. Now that we are no longer using ViewModelScope, our testing setup becomes simpler. In this example, I am using Turbine for testing.

class MainViewModelTest {

    private val fakeApiService = object : ApiService {
        override suspend fun getResult() = 42
    }

    @Test
    fun `result is emitted after fetching`() {
        runTest {
            val viewModel = MainViewModel(
                scope = this.backgroundScope,
                apiService = fakeApiService,
            )

            viewModel.result.test {
                viewModel.getResult()
                assertEquals(42, awaitItem())
            }
        }
    }
}

From this sample code, we can see that we didn't have to set the main dispatcher anywhere, and we were able to use the backgroundScope provided via TestScope in our ViewModel.

Let's now compare this test code to what we would have to write if we instead used viewModelScope to launch the coroutine in the ViewModel.

fun getResult() {
        viewModelScope.launch {
            _result.emit(apiService.getResult())
        }
    }

Unfortunately, when we run our test now, we get an error. Referring to the docs, we can see "If your code under test references the main thread, it’ll throw an exception during unit tests.". To get around this, we can create a JUnit Test Rule to automatically set the main dispatcher.

class MainDispatcherRule(
    private val testDispatcher: TestDispatcher = UnconfinedTestDispatcher(),
) : TestWatcher() {
    override fun starting(description: Description) {
        Dispatchers.setMain(testDispatcher)
    }

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

This test rule can now be used to get our tests to pass as expected

class MainViewModelTest {

    @get:Rule
    val mainDispatcherRule = MainDispatcherRule()

    private val fakeApiService = object : ApiService {
        override suspend fun getResult() = 42
    }

    @Test
    fun `result is emitted after fetching`() {
        runTest {
            val viewModel = MainViewModel(
                scope = this.backgroundScope,
                apiService = fakeApiService,
            )

            viewModel.result.test {
                viewModel.getResult()
                assertEquals(42, awaitItem())
            }
        }
    }
}

class MainDispatcherRule(
    private val testDispatcher: TestDispatcher = UnconfinedTestDispatcher(),
) : TestWatcher() {
    override fun starting(description: Description) {
        Dispatchers.setMain(testDispatcher)
    }

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

While we were still able to get a working solution by setting the main dispatcher, that is ultimately an undesirable workaround because for everything to play together nicely, we'd have to pass the TestDispatcher around for anything that wanted to know about it so that the structured concurrency in our runTest block will play nicely.

Conclusion

Using ViewModelScope to launch coroutines can be convenient, but it can also make it difficult to test coroutine code in isolation. By creating a custom scope that's decoupled from the lifecycle of the `ViewModel`, as well as that's managed in isolation in your test, you can launch coroutines in a way that's more testable and easier to manage.

In this blog post, we discussed how to create a custom scope and use it in a ViewModel to launch coroutines. We also showed how to write a unit test for coroutine code that uses a custom scope.