State Management in Jetpack Compose

Photo by Kelly Sikkema on Unsplash

State management is a crucial part of building modern Android applications using Jetpack Compose. Unlike traditional Android Views, which required manual UI updates, Compose follows a declarative UI model, meaning:

• The UI automatically updates when state changes.
• Data drives UI, rather than being modified manually.
• A single source of truth ensures predictable behavior.

However, improper state management can lead to unnecessary recompositions, performance issues, and maintainability problems. As Compose evolves, managing state effectively ensures your app remains responsive, maintainable, and scalable.

This guide covers every concept, pattern, and design principle necessary to master state management. You’ll also learn about Kotlin Flows for reactive state handling, with practical examples using a Weather App context. Additionally, we’ll explore how dependency injection with Koin can enhance state management in Jetpack Compose.

1. Understanding State in Jetpack Compose

What is State?

State refers to any data that can change over time and affect the UI. Examples include:

• The current temperature in a weather app.
• The selected city from a dropdown menu.
• A list of hourly forecasts retrieved from a network request.

Jetpack Compose follows a declarative UI model, meaning the UI is a function of the state. When the state changes, Compose automatically recomposes the UI to reflect the new state.

Key State Holders in Compose

+----------------------+----------------------------------------------------------+----------------------------+
|      State Type      |                      Description                         |  Survives Config Changes?  |
+----------------------+----------------------------------------------------------+----------------------------+
| remember             | Holds UI state but resets on config change               | No                         |
| rememberSaveable     | Stores UI state that survives config changes             | Yes                        |
| StateFlow            | Emits state changes, used in ViewModel                   | Yes                        |
| MutableStateFlow     | A mutable version of StateFlow                           | Yes                        |
| SharedFlow           | Emits state to multiple collectors, good for events      | Yes                        |
| SnapshotStateList    | Optimized state management for lists in Compose          | Yes                        |
| Repository (Flow)    | Handles data fetching, caching, and updates              | Yes                        |
| produceState         | Creates state that updates based on side effects         | Yes                        |
+----------------------+----------------------------------------------------------+----------------------------+

2. Weather App Example: ViewModel & StateFlow

A Weather App typically requires fetching data from an API and displaying the latest weather conditions. Here’s how you can structure state management in Jetpack Compose using Koin for dependency injection.

ViewModel with StateFlow (Koin Example)

class WeatherViewModel(private val repository: WeatherRepository) : ViewModel() {
    // MutableStateFlow to hold the temperature value
    private val _temperature = MutableStateFlow(0)
    // Expose the temperature as an immutable StateFlow
    val temperature: StateFlow<Int> get() = _temperature

    // Fetch weather data when the ViewModel is created
    init { fetchWeather() }

    private fun fetchWeather() {
        viewModelScope.launch {
            // Collect temperature data from the repository
            repository.getTemperature().collect { temp ->
                _temperature.value = temp // Update the state
            }
        }
    }
}

// Koin Module for Dependency Injection
val appModule = module {
    viewModel { WeatherViewModel(get()) } // Provide ViewModel
    single<WeatherRepository> { WeatherRepositoryImpl() } // Provide Repository
}

Explanation:

• MutableStateFlow: Used to hold mutable state that can be observed by the UI.
• StateFlow: Exposes the state as an immutable flow to the UI.
• viewModelScope.launch: Ensures the coroutine is tied to the ViewModel's lifecycle.
• Koin Module: Provides the WeatherViewModel and WeatherRepository dependencies.

3. WeatherScreen Composable

@Composable
fun WeatherScreen(viewModel: WeatherViewModel = getViewModel()) {
    // Collect temperature state from the ViewModel
    val temperature by viewModel.temperature.collectAsState()

    // UI Layout
    Column(horizontalAlignment = Alignment.CenterHorizontally) {
        Text("Current Temperature: $temperature°C", fontSize = 20.sp)
        Button(onClick = { viewModel.fetchWeather() }) {
            Text("Refresh Weather")
        }
    }
}

Explanation:

• collectAsState: Converts the StateFlow into a Compose state that triggers recomposition when the state changes.
• getViewModel(): Retrieves the WeatherViewModel instance using Koin's dependency injection.
• UI Components: Display the temperature and provide a button to refresh the data.

4. Error Handling in State Management

To handle errors gracefully, we can use a sealed class to represent different UI states: Loading, Success, and Error.

sealed class WeatherState {
    object Loading : WeatherState() // Represents a loading state
    data class Success(val temperature: Int) : WeatherState() // Represents a successful state with data
    data class Error(val message: String) : WeatherState() // Represents an error state with a message
}

class WeatherViewModel(private val repository: WeatherRepository) : ViewModel() {
    // MutableStateFlow to hold the weather state
    private val _weatherState = MutableStateFlow<WeatherState>(WeatherState.Loading)
    // Expose the weather state as an immutable StateFlow
    val weatherState: StateFlow<WeatherState> get() = _weatherState

    // Fetch weather data when the ViewModel is created
    init { fetchWeather() }

    private fun fetchWeather() {
        viewModelScope.launch {
            _weatherState.value = WeatherState.Loading // Set state to Loading
            try {
                // Collect temperature data from the repository
                repository.getTemperature().collect { temp ->
                    _weatherState.value = WeatherState.Success(temp) // Set state to Success
                }
            } catch (e: Exception) {
                // Set state to Error if an exception occurs
                _weatherState.value = WeatherState.Error("Failed to fetch weather data")
            }
        }
    }
}

Updated WeatherScreen Composable

@Composable
fun WeatherScreen(viewModel: WeatherViewModel = getViewModel()) {
    // Collect weather state from the ViewModel
    val weatherState by viewModel.weatherState.collectAsState()

    // UI Layout
    Column(horizontalAlignment = Alignment.CenterHorizontally) {
        when (val state = weatherState) {
            is WeatherState.Loading -> Text("Loading...") // Show loading state
            is WeatherState.Success -> Text("Current Temperature: ${state.temperature}°C", fontSize = 20.sp) // Show success state
            is WeatherState.Error -> Text("Error: ${state.message}", color = Color.Red) // Show error state
        }
        Button(onClick = { viewModel.fetchWeather() }) {
            Text("Refresh Weather")
        }
    }
}

Explanation:

• sealed class: Represents different UI states in a type-safe way.
• Error Handling: Catches exceptions and updates the state to Error if something goes wrong.
• UI Updates: Displays the appropriate UI based on the current state.

5. Using produceState for Side-Effect-Driven State

produceState is useful for creating state that updates based on side effects, such as fetching data from an API.

@Composable
fun WeatherScreen(viewModel: WeatherViewModel = getViewModel()) {
    // Create state using produceState
    val weatherState by produceState<WeatherState>(initialValue = WeatherState.Loading) {
        viewModel.fetchWeather() // Trigger data fetch
        viewModel.weatherState.collect { value = it } // Update state based on ViewModel's state
    }

    // UI Layout
    Column(horizontalAlignment = Alignment.CenterHorizontally) {
        when (val state = weatherState) {
            is WeatherState.Loading -> Text("Loading...")
            is WeatherState.Success -> Text("Current Temperature: ${state.temperature}°C", fontSize = 20.sp)
            is WeatherState.Error -> Text("Error: ${state.message}", color = Color.Red)
        }
        Button(onClick = { viewModel.fetchWeather() }) {
            Text("Refresh Weather")
        }
    }
}

Explanation:

• produceState: Creates a state that updates based on side effects (e.g., fetching data).
• Initial Value: Starts with WeatherState.Loading.
• State Updates: Collects state changes from the ViewModel and updates the UI.

6. Optimizing Recomposition with derivedStateOf

Use derivedStateOf to optimize recompositions when deriving state from other state variables.

@Composable
fun WeatherScreen(viewModel: WeatherViewModel = getViewModel()) {
    // Collect weather state from the ViewModel
    val weatherState by viewModel.weatherState.collectAsState()
    // Derive a boolean state indicating if it's hot outside
    val isHot = derivedStateOf {
        (weatherState as? WeatherState.Success)?.temperature ?: 0 > 30
    }

    // UI Layout
    Column(horizontalAlignment = Alignment.CenterHorizontally) {
        when (val state = weatherState) {
            is WeatherState.Loading -> Text("Loading...")
            is WeatherState.Success -> Text("Current Temperature: ${state.temperature}°C", fontSize = 20.sp)
            is WeatherState.Error -> Text("Error: ${state.message}", color = Color.Red)
        }
        if (isHot.value) {
            Text("It's hot outside!", color = Color.Red)
        }
        Button(onClick = { viewModel.fetchWeather() }) {
            Text("Refresh Weather")
        }
    }
}

Explanation:

• derivedStateOf: Creates a derived state that only triggers recomposition when the derived value changes.
• Optimization: Prevents unnecessary recompositions when unrelated state changes occur.

7. Unit Testing State Management

Testing WeatherViewModel

@RunWith(JUnit4::class)
class WeatherViewModelTest {
    private lateinit var viewModel: WeatherViewModel
    private val repository = FakeWeatherRepository()

    @Before
    fun setUp() {
        viewModel = WeatherViewModel(repository)
    }

    @Test
    fun `fetchWeather updates temperature`() = runTest {
        viewModel.fetchWeather()
        assertThat((viewModel.weatherState.value as WeatherState.Success).temperature).isGreaterThan(0)
    }

    @Test
    fun `fetchWeather handles error`() = runTest {
        repository.shouldThrowError = true
        viewModel.fetchWeather()
        assertThat(viewModel.weatherState.value).isInstanceOf(WeatherState.Error::class.java)
    }
}

Explanation:

• Fake Repository: Simulates network requests for testing.
• State Assertions: Verifies that the ViewModel updates the state correctly.

8. Best Practices

• Leverage ViewModel + StateFlow for managing asynchronous data and ensuring state persistence across configuration changes.
• Utilize Dependency Injection (Koin) for cleaner architecture and improved testability.
• Handle UI State Effectively:
  - Use sealed classes or Result wrappers to represent loading, success, and error states.
  - Optimize recompositions with remember, derivedStateOf, and produceState.
• Write Unit Tests for both ViewModel state and UI behavior to ensure reliability.
• Monitor Performance using tools like Compose Layout Inspector and Recomposition Counts.
• Plan for State Restoration and Deep Linking to ensure a seamless user experience.

By applying these modern Jetpack Compose state management principles, your apps will be efficient, scalable, and maintainable. For further insights into structuring your app for scalability and maintainability, check out this article on Modularization and Project Optimization in Android Development.