UI state class based on this answer https://stackoverflow.com/a/63824909/6301831

SomeViewModel.kt

viewModel.viewState.observe(viewLifecycleOwner) { state ->
    state takeIfSuccess {
        // Here's the success state
    } takeIfError {
        // Here's the error state
    }
}

UIState.kt

sealed class UIState<out T> where T : Any? {

    object Loading : UIState<Nothing>()

    data class Success<T>(val data: T) : UIState<T>()

    data class Failure(val errorMessage: String, val messageType: UIComponentType) : UIState<Nothing>()
}
infix fun <T> UIState<T>.takeIfSuccess(onSuccess: UIState.Success<T>.() -> Unit): UIState<T> {
    return when (this) {
        is UIState.Success -> {
            onSuccess(this)
            this
        }
        else -> {
            this
        }
    }
}

infix fun <T> UIState<T>.takeIfError(onError: UIState.Failure.() -> Unit): UIState<T> {
    return when (this) {
        is UIState.Failure -> {
            onError(this)
            this
        }
        else -> {
            this
        }
    }
}