Transitioning to Unified State Management: Refactoring a ViewModel and Screen

In modern app development, managing UI state in a clean and scalable way is critical. Many developers start with a non-unified approach where individual fields (like posts, comments, isLiked) are stored as separate LiveData properties in the ViewModel. While this works for simple scenarios, it often leads to duplication, bugs, and difficulty testing as the application grows.

In this article, we’ll refactor a non-unified stateful design into a unified state design using the ViewModel and its corresponding screen as an example. We'll focus on cleaning up the ViewModel, structuring the UI state, and simplifying the screen’s logic.

Why Unified State Design?

Unified state design organizes all the state variables into a single object, often called UiState. The benefits include:

1. Single Source of Truth: All UI-related data is stored in one place.
2. Simpler State Management: Avoids inconsistent states or redundant LiveData updates.
3. Easier Testing: You can test the ViewModel by simply observing the UiState.
4. Improved Readability: Reduces the mental overhead of tracking multiple LiveData properties.

Step 1: The Initial Non-Unified ViewModel

Here’s what the original ViewDaynoteViewModel might look like:

class ViewDaynoteViewModel : ViewModel() {
    private val _daynote = MutableLiveData<Daynote?>()
    val daynote: LiveData<Daynote?> = _daynote

    private val _userPreviewData = MutableLiveData<UserPreview?>()
    val userPreviewData: LiveData<UserPreview?> = _userPreviewData

    private val _isDaynoteLiked = MutableLiveData<Boolean>()
    val isDaynoteLiked: LiveData<Boolean> = _isDaynoteLiked

    private val _commentList = MutableLiveData<List<Comment>>()
    val commentList: LiveData<List<Comment>> = _commentList

    private val _commentText = MutableLiveData<String>()
    val commentText: LiveData<String> = _commentText

    private val _isCommentSectionVisible = MutableLiveData<Boolean>()
    val isCommentSectionVisible: LiveData<Boolean> = _isCommentSectionVisible

    private val _toastMessage = MutableSharedFlow<ToastMessage>()
    val toastMessage: SharedFlow<ToastMessage> = _toastMessage

    // Functions like loadDaynote(), toggleCommentSection(), etc...
}

Problems with This Approach:

1. Scattered State: The UI has to observe multiple LiveData properties, leading to redundant updates and inconsistent states.
2. Boilerplate: For every new state, we need a new LiveData property and corresponding logic.
3. Hard to Test: Testing multiple LiveData properties individually increases complexity.

Step 2: Designing the Unified UiState

First, we define a single UiState class to encapsulate all the state in the ViewModel:

sealed class ViewDaynoteUiState {
    object Loading : ViewDaynoteUiState()
    data class Success(
        val daynote: Daynote,
        val userPreview: UserPreview,
        val isLiked: Boolean,
        val isCurrentUser: Boolean,
        val isCommentSectionVisible: Boolean = false,
        val comments: List<Comment> = emptyList(),
    ) : ViewDaynoteUiState()
    object Error : ViewDaynoteUiState()
    object Deleted : ViewDaynoteUiState()
}

This structure:

• Consolidates all state into a single object.
• Adds meaningful states like Loading, Error, and Deleted.
• Simplifies state transitions during operations like deleting a daynote.

Step 3: Refactoring the ViewModel

Next, we replace the scattered LiveData properties with a single StateFlow for the UI state:

class ViewDaynoteViewModel : ViewModel() {
    private val _uiState = MutableStateFlow<ViewDaynoteUiState>(ViewDaynoteUiState.Loading)
    val uiState: StateFlow<ViewDaynoteUiState> = _uiState

    var commentText = MutableStateFlow("")

    fun loadDaynote(daynoteId: String) {
        viewModelScope.launch {
            _uiState.value = ViewDaynoteUiState.Loading
            val result = daynoteRepository.getDaynoteById(daynoteId).flatMap { daynote ->
                userRepository.getUserPreviewById(daynote.userId).map { userPreview ->
                    Triple(daynote, userPreview, daynote.userId == currentUserId)
                }
            }

            result.onSuccess { (daynote, userPreview, isCurrentUser) ->
                val isLiked = likeRepository.isDaynoteLikedByUser(daynoteId, currentUserId!!)
                _uiState.value = ViewDaynoteUiState.Success(
                    daynote = daynote,
                    userPreview = userPreview,
                    isLiked = isLiked,
                    isCurrentUser = isCurrentUser,
                )
                loadComments(daynoteId)
            }.onFailure {
                _uiState.value = ViewDaynoteUiState.Error
            }
        }
    }

    fun toggleCommentSection() {
        val currentState = _uiState.value
        if (currentState is ViewDaynoteUiState.Success) {
            _uiState.value = currentState.copy(
                isCommentSectionVisible = !currentState.isCommentSectionVisible
            )
        }
    }

    fun deleteDaynote(daynoteId: String) {
        viewModelScope.launch {
            val result = daynoteRepository.deleteDaynote(daynoteId)
            result.onSuccess {
                _uiState.value = ViewDaynoteUiState.Deleted
            }.onFailure {
                _toastMessage.emit(ToastMessage.Error("Failed to delete daynote"))
            }
        }
    }

    // Additional logic for comments, likes, etc.
}

Key Changes:

1. StateFlow for Unified State: All UI state is derived from _uiState.
2. Simpler State Management:

• Operations like toggling the comment section update the state with copy().
• The UI automatically updates because it observes uiState.

Step 4: Refactoring the Screen

Finally, we update the ViewDaynoteScreen to consume the new unified uiState:

@Composable
fun ViewDaynoteScreen(
    navigateBack: () -> Unit,
    daynoteId: String,
    onTitleChange: (String) -> Unit,
    viewModel: ViewDaynoteViewModel = hiltViewModel()
) {
    val uiState by viewModel.uiState.collectAsState()
    val commentText by viewModel.commentText.collectAsState()

    when (val state = uiState) {
        is ViewDaynoteUiState.Loading -> {
            Box(
                modifier = Modifier.fillMaxSize(),
                contentAlignment = Alignment.Center
            ) {
                CircularProgressIndicator()
            }
        }

        is ViewDaynoteUiState.Success -> {
            LazyColumn {
                item {
                    DaynoteCard(
                        note = state.daynote,
                        userInfo = state.userPreview,
                        isDaynoteLiked = state.isLiked,
                        isCurrentUser = state.isCurrentUser,
                        onCommentSectionVisible = { viewModel.toggleCommentSection() },
                        onLikeDaynote = viewModel::toggleDaynoteLike,
                        onDeleteDaynote = { /* Show delete dialog */ }
                    )
                }

                if (state.isCommentSectionVisible) {
                    item {
                        CommentEntrySection(
                            commentText = commentText,
                            onCommentTextChanged = { viewModel.commentText.value = it },
                            onSendComment = viewModel::sendComment
                        )
                    }
                }

                items(state.comments) { comment ->
                    CommentCard(comment)
                }
            }
        }

        is ViewDaynoteUiState.Error -> {
            Box(
                modifier = Modifier.fillMaxSize(),
                contentAlignment = Alignment.Center
            ) {
                Text("Failed to load daynote", color = Color.Red)
            }
        }

        is ViewDaynoteUiState.Deleted -> {
            navigateBack()
        }
    }
}

Why the New Approach Works

1. Cleaner ViewModel:

• All logic is tied to uiState.
• There’s less duplication and easier state transitions.

2. Simpler Screen:

• The screen renders the UI based on the uiState.
• Each state (Loading, Success, Error) maps directly to a specific UI.

3. Easier Debugging:

• Observing a single StateFlow simplifies debugging and testing.

Conclusion

Refactoring to a unified state design can seem daunting initially, but it greatly simplifies your code in the long run. By consolidating all UI-related data into a single UiState, you make your ViewModel more maintainable and your screens easier to read.

This approach also prepares your app for scalability, enabling you to handle more complex features or states with minimal effort.