Tests for Kotlin Flow ViewModel tips - https://hrach.dev/posts/kotlin-flow-tips-for-viewmodel/

BaseViewModel.kt

package dev.hrach.examples

import androidx.lifecycle.SavedStateHandle
import androidx.lifecycle.ViewModel
import androidx.lifecycle.serialization.saved
import androidx.lifecycle.viewModelScope
import androidx.savedstate.serialization.serializers.MutableStateFlowSerializer
import kotlin.time.Duration.Companion.seconds
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.SharingStarted
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.WhileSubscribed
import kotlinx.coroutines.flow.channelFlow
import kotlinx.coroutines.flow.collectLatest
import kotlinx.coroutines.flow.onStart
import kotlinx.coroutines.flow.stateIn
import kotlinx.serialization.KSerializer
import kotlinx.serialization.serializer

abstract class BaseViewModel : ViewModel() {
	protected fun <T> Flow<T>.stateIn(initialValue: T): StateFlow<T> =
		stateIn(viewModelScope, SharingStarted.WhileSubscribed(5.seconds), initialValue)

	protected fun <T> Flow<T>.refreshOn(refreshSignal: Flow<Unit>): Flow<T> =
		channelFlow {
			refreshSignal
				.onStart { emit(Unit) }
				.collectLatest {
					this@refreshOn.collect { send(it) }
				}
		}

	protected inline fun <reified T> SavedStateHandle.savedStateFlow(
		serializer: KSerializer<T> = serializer<T>(),
		key: String? = null,
		crossinline init: () -> T,
	) = saved(
		serializer = MutableStateFlowSerializer(serializer),
		key = key,
		init = { MutableStateFlow(init()) },
	)
}

BaseViewModelTest.kt

package dev.hrach.examples

import androidx.lifecycle.SavedStateHandle
import androidx.lifecycle.createSavedStateHandle
import androidx.lifecycle.viewmodel.testing.viewModelScenario
import app.cash.turbine.test
import dev.hrach.examples.BaseViewModel
import kotlin.test.Test
import kotlin.test.assertEquals
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.combine
import kotlinx.coroutines.test.runTest
import org.junit.runner.RunWith
import org.robolectric.RobolectricTestRunner

@RunWith(RobolectricTestRunner::class)
class BaseViewModelTest {
	enum class Gender { Male, Female }

	class TestVM(
		state: SavedStateHandle,
	) : BaseViewModel() {
		val data1: MutableStateFlow<Int?> by state.savedStateFlow { 41 }
		val data2: MutableStateFlow<Int?> by state.savedStateFlow { null }
		val data3: MutableStateFlow<Gender> by state.savedStateFlow { Gender.Female }

		var c = 0
		val refresh = MutableSharedFlow<Unit>(extraBufferCapacity = 1)
		val data = combine(data1, data2) { a, b -> (a ?: 0) + (b ?: 0) + (++c) }
			.refreshOn(refresh)
	}

	@Test
	fun testReftech() = runTest {
		val vm = TestVM(SavedStateHandle())
		vm.data.test {
			assertEquals(42, awaitItem())
			expectNoEvents()

			vm.data2.value = 1
			assertEquals(44, awaitItem())
			expectNoEvents()

			vm.refresh.tryEmit(Unit)
			assertEquals(45, awaitItem())
			expectNoEvents()

			vm.refresh.emit(Unit)
			assertEquals(46, awaitItem())
			expectNoEvents()

			vm.data2.value = 2
			assertEquals(48, awaitItem())
			expectNoEvents()
		}
	}

	@Test
	fun fullRestoreStateTest() {
		val scenario = viewModelScenario { TestVM(createSavedStateHandle()) }
		assertEquals(41, scenario.viewModel.data1.value)
		assertEquals(null, scenario.viewModel.data2.value)
		assertEquals(Gender.Female, scenario.viewModel.data3.value)

		scenario.viewModel.data1.value = null
		scenario.viewModel.data2.value = 10
		scenario.viewModel.data3.value = Gender.Male

		assertEquals(null, scenario.viewModel.data1.value)
		assertEquals(10, scenario.viewModel.data2.value)
		assertEquals(Gender.Male, scenario.viewModel.data3.value)

		scenario.recreate()

		assertEquals(null, scenario.viewModel.data1.value)
		assertEquals(10, scenario.viewModel.data2.value)
		assertEquals(Gender.Male, scenario.viewModel.data3.value)
	}
}