UI State, StateFlow and Unit-Testing the ViewModel

In lesson 24 you moved Dessert Clicker's state into a ViewModel. That was the first half of the job. The second half is the payoff: because the logic now lives in a plain Kotlin class, you can test it automatically in milliseconds, with no emulator, no phone and no tapping.

In this lesson you will look more closely at UI state and StateFlow, learn what SavedStateHandle adds, and write real unit tests for a ViewModel using JUnit and the coroutines test library. You will finish the Unit 4 ViewModel work with the practice exercise of adding a ViewModel to an app that does not have one.

This lesson covers two items from Unit 4 of the official Android Basics with Compose course: Write unit tests to test the ViewModel and Practice: Add a ViewModel to Dessert Clicker.

By the end you will be able to:

• Design a UI state class and update it with MutableStateFlow.update.
• Explain what a StateFlow is: always has a value, replays the latest to new collectors, ignores equal values.
• Keep small essential state across process death with SavedStateHandle.
• Write local unit tests (src/test) with JUnit 4 and assertions.
• Test code that launches coroutines with runTest and a test dispatcher.
• Explain why a Dispatchers.Main replacement is needed in ViewModel tests, and install one with a JUnit rule.

Big idea: a test is an executable specification. Writing "after 5 taps the dessert is a Donut" as a test means you never have to re-tap by hand again.

• ViewModel and MutableStateFlow (lesson 24).
• Data classes and copy (lesson 11).
• Why test? In lesson 20 you wrote tests for a function and a composable. Today's tests target logic that has no UI at all.
• Local versus instrumented tests. Android Studio gives you two test folders:

ViewModel logic with no Android classes belongs in src/test.

StateFlow is a hot, state-holding flow:

• It always has a current value (.value). You must give it an initial value when you create it.
• A new collector immediately receives the latest value.
• It conflates: if you set many values quickly, a slow collector may only see the latest.
• It ignores a new value that is equal (==) to the current one. This is why immutable data classes work so well: set an equal state and nothing recomposes.

MutableStateFlow(initial)           StateFlow (read-only view)
  .value = new      ───────────▶     .value   (read the latest)
  .update { old -> new }             .collect { }  (observe changes)

Three ways to change a MutableStateFlow

_uiState.value = UiState() // 1: replace outright
_uiState.value = _uiState.value.copy(count = 1) // 2: read then write (NOT thread-safe)
_uiState.update { it.copy(count = it.count + 1) } // 3: atomic read-modify-write

Option 3 is safest. If two coroutines both run option 2 at the same moment they can read the same old value and one change gets lost. update retries until the swap succeeds. Use it whenever the new value depends on the old one.

Designing a UI state class

Ask: "what does the screen need to draw?" Put exactly that in the class. Good UI state is:

• Immutable (val properties).
• Complete: the screen can render purely from it.
• Free of impossible combinations where possible. For a screen that loads data, a sealed interface often beats several booleans:

sealed interface NewsUiState {
 data object Loading : NewsUiState
 data class Success(val headlines: List<String>) : NewsUiState
 data class Error(val message: String) : NewsUiState
}

With three booleans (isLoading, hasError, hasData) there are eight combinations and most are nonsense. With the sealed type there are exactly three, and Kotlin's when forces you to handle each one. You will use this pattern in lesson 29.

In lesson 24 you saw that a ViewModel does not survive process death. SavedStateHandle is a small key-value map that does, because it is backed by the same saved instance state as rememberSaveable. The ViewModel receives it through its constructor.

import androidx.lifecycle.SavedStateHandle
import androidx.lifecycle.ViewModel
import kotlinx.coroutines.flow.StateFlow
class CounterViewModel(private val savedStateHandle: SavedStateHandle) : ViewModel() {
 val count: StateFlow<Int> = savedStateHandle.getStateFlow("count", 0)
 fun increment() {
 savedStateHandle["count"] = count.value + 1
 }
}

• getStateFlow(key, initialValue) gives a StateFlow that is always in sync with the saved value.
• savedStateHandle[key] = value writes the value, and it is stored for process death.

With the plain viewModel() function in Compose, the default factory supplies the SavedStateHandle automatically for ViewModels whose constructor takes only a SavedStateHandle (or has no parameters). Anything more complex needs a factory, which lesson 30 covers.

Keep it to small things: the selected tab, a search query, an item ID. Big data belongs in storage. The same 1 MB Bundle limit from lesson 23 applies.

To make testing straightforward, we will use a slightly extended Dessert Clicker ViewModel. Start from lesson 24's code. For clarity, here is the full ViewModel once more, plus one new rule: the reset function.

package com.example.dessertclicker.ui
import androidx.lifecycle.ViewModel
import com.example.dessertclicker.desserts
import com.example.dessertclicker.determineDessertIndex
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.flow.update
class DessertViewModel : ViewModel() {
 private val _uiState = MutableStateFlow(DessertUiState())
 val uiState: StateFlow<DessertUiState> = _uiState.asStateFlow()
 fun onDessertClicked() {
 _uiState.update { current ->
 val sold = current.dessertsSold + 1
 val index = determineDessertIndex(sold)
 current.copy(
 revenue = current.revenue + desserts[current.currentDessertIndex].price,
 dessertsSold = sold,
 currentDessertIndex = index,
 currentDessertName = desserts[index].name
 )
 }
 }
 fun onReset() {
 _uiState.value = DessertUiState()
 }
}

Nothing in it knows about Android's UI. That is exactly why we can test it on the JVM.

Check the test dependencies

A new project already has JUnit 4 (testImplementation(libs.junit)) in app/build.gradle.kts. For coroutine tests you also need the coroutines test library. Add this line and sync (look up the latest stable version on the kotlinx.coroutines releases page, and use the same version as your coroutines library; the Compose and lifecycle libraries pull in coroutines for you):

dependencies {
 testImplementation(libs.junit)
 testImplementation("org.jetbrains.kotlinx:kotlinx-coroutines-test:<same version as coroutines>")
}

(A version-catalog entry is tidier; lesson 28 shows both styles.)

The test class

Right-click app/src/test/java/com/example/dessertclicker, choose New > Kotlin Class/File, and name it DessertViewModelTest:

package com.example.dessertclicker
import com.example.dessertclicker.ui.DessertUiState
import com.example.dessertclicker.ui.DessertViewModel
import org.junit.Assert.assertEquals
import org.junit.Before
import org.junit.Test
class DessertViewModelTest {
 private lateinit var viewModel: DessertViewModel
 @Before
 fun setUp() {
 viewModel = DessertViewModel()
 }
 @Test
 fun initialState_isCupcakeWithZeroSales() {
 val state = viewModel.uiState.value
 assertEquals(0, state.dessertsSold)
 assertEquals(0, state.revenue)
 assertEquals("Cupcake", state.currentDessertName)
 }
 @Test
 fun oneClick_sellsOneCupcake() {
 viewModel.onDessertClicked()
 val state = viewModel.uiState.value
 assertEquals(1, state.dessertsSold)
 assertEquals(5, state.revenue) // a cupcake costs 5
 assertEquals("Cupcake", state.currentDessertName)
 }
 @Test
 fun fiveClicks_unlockDonut() {
 repeat(5) { viewModel.onDessertClicked() }
 val state = viewModel.uiState.value
 assertEquals(5, state.dessertsSold)
 assertEquals(25, state.revenue) // five cupcakes at 5 each
 assertEquals("Donut", state.currentDessertName)
 }
 @Test
 fun reset_returnsToInitialState() {
 repeat(8) { viewModel.onDessertClicked() }
 viewModel.onReset()
 assertEquals(DessertUiState(), viewModel.uiState.value)
 }
}

Understanding the test

• @Test marks a function JUnit should run. Each test gets a fresh instance of the test class.
• @Before runs before every test, so each starts from a new ViewModel. Tests must never depend on each other or on order.
• assertEquals(expected, actual) fails the test if the values differ. Note the argument order: expected first.
• The test names describe behaviour (fiveClicks_unlockDonut). A good name tells you what broke without opening the file.
• assertEquals(DessertUiState(), viewModel.uiState.value) compares whole data classes. This works because data classes generate equals for you.

Check the arithmetic for fiveClicks_unlockDonut. Click 1 to 4 sell cupcakes (index 0, price 5) each, 20 in revenue. At click 5 the sold count reaches 5 after the click, but the price comes from the dessert shown before that click, which is still the cupcake, so revenue is 25. Then the index becomes 1 (Donut). The test asserts 25, matching that reasoning.

Running the tests

Click the green arrow beside the class name, and choose Run 'DessertViewModelTest'. A green tick means pass. Or use the terminal:

./gradlew test

Break something on purpose (change the cupcake price in desserts) and watch the test turn red with a message saying expected:<25> but was:<...>. A test you have never seen fail is a test you cannot fully trust.

Soon your ViewModels will launch coroutines with viewModelScope.launch (lesson 28). Consider a version that delays before applying a "bonus":

import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewModelScope
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.flow.update
import kotlinx.coroutines.launch
class BonusViewModel : ViewModel() {
 private val _bonus = MutableStateFlow(0)
 val bonus: StateFlow<Int> = _bonus.asStateFlow()
 fun claimBonus() {
 viewModelScope.launch {
 delay(2_000) // pretend a server call takes two seconds
 _bonus.update { it + 100 }
 }
 }
}

Testing this naively has two problems:

1. viewModelScope runs on Dispatchers.Main, which does not exist in a plain JVM unit test, because there is no Android main thread. The test crashes with Module with the Main dispatcher had failed to initialize.
2. A real two-second delay would make the test slow.

The solution is the coroutines test library:

• Dispatchers.setMain(testDispatcher) replaces the Main dispatcher for the test, and Dispatchers.resetMain() restores it afterwards.
• runTest { } runs a coroutine with a virtual clock. A delay(2_000) inside it completes instantly, because the test scheduler skips ahead.

A reusable JUnit rule installs and removes the dispatcher for every test:

package com.example.dessertclicker
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.test.StandardTestDispatcher
import kotlinx.coroutines.test.TestDispatcher
import kotlinx.coroutines.test.resetMain
import kotlinx.coroutines.test.setMain
import org.junit.rules.TestWatcher
import org.junit.runner.Description
@OptIn(ExperimentalCoroutinesApi::class)
class MainDispatcherRule(
 val testDispatcher: TestDispatcher = StandardTestDispatcher()
) : TestWatcher() {
 override fun starting(description: Description) {
 Dispatchers.setMain(testDispatcher)
 }
 override fun finished(description: Description) {
 Dispatchers.resetMain()
 }
}

setMain and resetMain are marked experimental in the library, hence the @OptIn. Now the test:

package com.example.dessertclicker
import kotlinx.coroutines.test.advanceUntilIdle
import kotlinx.coroutines.test.runTest
import org.junit.Assert.assertEquals
import org.junit.Rule
import org.junit.Test
class BonusViewModelTest {
 @get:Rule
 val mainDispatcherRule = MainDispatcherRule()
 @Test
 fun claimBonus_addsOneHundredAfterDelay() = runTest {
 val viewModel = BonusViewModel()
 viewModel.claimBonus()
 assertEquals(0, viewModel.bonus.value) // the two seconds have not passed yet
 advanceUntilIdle() // run pending coroutines, skipping the delay
 assertEquals(100, viewModel.bonus.value)
 }
}

• @get:Rule applies the rule to each test. (The get: is Kotlin syntax so JUnit sees the annotation on the getter.)
• StandardTestDispatcher does not run launched coroutines immediately. They are queued until you advance the scheduler, which lets you assert on the "before" state.
• advanceUntilIdle() runs all queued work. In Kotlin coroutines tests the runTest scheduler and the rule's dispatcher must share a scheduler. When you create the ViewModel inside runTest after setMain, runTest picks up the Main test dispatcher's scheduler automatically, so they cooperate.

For the more complete treatment of flows and testing them, see lesson 33.

StateFlow exposes .value, so most tests simply read it. To check the sequence of emissions you can collect into a list:

import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.flow.toList
import kotlinx.coroutines.launch
import kotlinx.coroutines.test.UnconfinedTestDispatcher
import kotlinx.coroutines.test.runTest
import org.junit.Assert.assertEquals
import org.junit.Test
@OptIn(ExperimentalCoroutinesApi::class)
@Test
fun uiState_emitsEveryChange() = runTest {
 val viewModel = DessertViewModelFake()
 val seen = mutableListOf<Int>()
 val job = launch(UnconfinedTestDispatcher(testScheduler)) {
 viewModel.uiState.toList(seen.let { mutableListOf() }.also { }) // placeholder
 }
 job.cancel()
}

That example is deliberately not what you should write, and it is a good reminder: do not make tests clever. A StateFlow conflates updates, so collecting a list can miss intermediate values and produce flaky tests. Prefer asserting on viewModel.uiState.value after each action. For genuine sequences, use first() or the library approach shown in lesson 33.

(That dead-end snippet is here so you learn to recognise over-engineering. Delete it from your own file.)

Mistake #1: Tests that depend on each other. What happens: test B only passes after test A ran. The fix: create a fresh ViewModel in @Before, and never share mutable state between tests.

Mistake #2: Forgetting to replace Dispatchers.Main. Symptom: Module with the Main dispatcher had failed to initialize. The fix: use a MainDispatcherRule as shown.

Mistake #3: Thread.sleep in tests. What happens: slow and flaky tests. The fix: runTest with the virtual clock and advanceUntilIdle().

Mistake #4: Swapping expected and actual in assertEquals. What happens: the failure message is misleading. The fix: expected first, always.

Mistake #5: Reading .value before the coroutine ran. What happens: the assertion sees the old state. The fix: advance the test scheduler before asserting.

Mistake #6: Testing implementation, not behaviour. The wrong test: checking that a private function was called. The fix: assert on the state the user would see.

Mistake #7: Putting Android classes in a local test. Symptom: Method ... not mocked. The fix: keep the ViewModel free of Android classes, or move that test to androidTest.

Mistake #8: Trusting an exercise instead of the arithmetic. What happens: you "fix" correct code to match a wrong expectation. The fix: calculate by hand, then decide which side is wrong.

Question 1. Why does MutableStateFlow require an initial value?

Question 2. Why use update { } instead of value = value.copy(...)?

Question 3. Where do local unit tests live and what do they run on?

Question 4. Why do ViewModel tests that use viewModelScope need Dispatchers.setMain?

Question 5. What does SavedStateHandle add to a ViewModel?

Go deeper: official resources

• Android Basics with Compose, Unit 4, Pathway 1: Write unit tests to test the ViewModel and the Dessert Clicker practice.
• UI layer: state production
• UI events
• Kotlin coroutines on Android