Improve login vm service connection handling

4 files changed, 70 insertions, 48 deletions

diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
index 0910754a8d..f86a8d69e2 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
@@ -37,7 +37,7 @@ val uiModule = module {
 
     single { ServiceConnectionManager(androidContext()) }
     single { DeviceRepository(get()) }
-    viewModel { LoginViewModel() }
+    viewModel { LoginViewModel(get()) }
     viewModel { DeviceRevokedViewModel(get()) }
 }
 
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/model/AccountHistory.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/model/AccountHistory.kt
index 114463aaaa..11e9b20604 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/model/AccountHistory.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/model/AccountHistory.kt
@@ -9,4 +9,6 @@ sealed class AccountHistory : Parcelable {
 
     @Parcelize
     object Missing : AccountHistory()
+
+    fun accountToken() = (this as? Available)?.accountToken
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/LoginFragment.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/LoginFragment.kt
index dcd8844656..3b0b78d9a2 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/LoginFragment.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/LoginFragment.kt
@@ -15,8 +15,6 @@ import androidx.lifecycle.repeatOnLifecycle
 import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.R
-import net.mullvad.mullvadvpn.model.AccountHistory
-import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionContainer
 import net.mullvad.mullvadvpn.ui.widget.AccountLogin
 import net.mullvad.mullvadvpn.ui.widget.HeaderBar
 import net.mullvad.mullvadvpn.viewmodel.LoginViewModel
@@ -57,8 +55,6 @@ class LoginFragment :
         loggedInStatus = view.findViewById(R.id.logged_in_status)
         loginFailStatus = view.findViewById(R.id.login_fail_status)
 
-        loginViewModel.updateAccountCacheInstance(accountCache)
-
         accountLogin = view.findViewById<AccountLogin>(R.id.account_login).apply {
             onLogin = loginViewModel::login
             onClearHistory = loginViewModel::clearAccountHistory
@@ -78,16 +74,6 @@ class LoginFragment :
         return view
     }
 
-    override fun onNewServiceConnection(serviceConnectionContainer: ServiceConnectionContainer) {
-        super.onNewServiceConnection(serviceConnectionContainer)
-        loginViewModel.updateAccountCacheInstance(accountCache)
-    }
-
-    override fun onNoServiceConnection() {
-        super.onNoServiceConnection()
-        loginViewModel.updateAccountCacheInstance(null)
-    }
-
     override fun onSafelyStart() {
         parentActivity.backButtonHandler = {
             if (accountLogin.hasFocus) {
@@ -113,8 +99,7 @@ class LoginFragment :
             repeatOnLifecycle(Lifecycle.State.RESUMED) {
                 launch {
                     loginViewModel.accountHistory.collect { history ->
-                        accountLogin.accountHistory = history
-                            .let { it as? AccountHistory.Available }?.accountToken
+                        accountLogin.accountHistory = history.accountToken()
                     }
                 }
                 launch {
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt
index e9cb27fda6..6aac2d8177 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt
@@ -2,24 +2,50 @@ package net.mullvad.mullvadvpn.viewmodel
 
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.collect
+import kotlinx.coroutines.flow.emptyFlow
 import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.onStart
+import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.model.AccountCreationResult
 import net.mullvad.mullvadvpn.model.AccountHistory
 import net.mullvad.mullvadvpn.model.LoginResult
 import net.mullvad.mullvadvpn.ui.serviceconnection.AccountCache
+import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
+import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionState
 
-class LoginViewModel : ViewModel() {
+class LoginViewModel(
+    private val serviceConnectionManager: ServiceConnectionManager,
+    private val dispatcher: CoroutineDispatcher = Dispatchers.IO
+) : ViewModel() {
     private val _uiState = MutableStateFlow<LoginUiState>(LoginUiState.Default)
     val uiState: StateFlow<LoginUiState> = _uiState
 
-    private val _accountHistory = MutableStateFlow<AccountHistory>(AccountHistory.Missing)
-    val accountHistory: StateFlow<AccountHistory> = _accountHistory
+    private val accountCache: AccountCache?
+        get() {
+            return serviceConnectionManager.connectionState.value.readyContainer()?.accountCache
+        }
 
-    private var accountCache: AccountCache? = null
+    val accountHistory = serviceConnectionManager.connectionState
+        .flatMapLatest { state ->
+            if (state is ServiceConnectionState.ConnectedReady) {
+                state.container.accountCache.accountHistoryEvents
+                    .onStart {
+                        state.container.accountCache.fetchAccountHistory()
+                    }
+            } else {
+                emptyFlow()
+            }
+        }
+        .stateIn(CoroutineScope(dispatcher), SharingStarted.Lazily, AccountHistory.Missing)
 
     sealed class LoginUiState {
         object Default : LoginUiState()
@@ -36,45 +62,50 @@ class LoginViewModel : ViewModel() {
         data class OtherError(val errorMessage: String) : LoginUiState()
     }
 
-    // Ensures the view model has an up-to-date instance of account cache. This is an intermediate
-    // solution due to limitations in the current app architecture.
-    fun updateAccountCacheInstance(newAccountCache: AccountCache?) {
-        accountCache = newAccountCache?.apply {
-            viewModelScope.launch {
-                accountHistoryEvents.collect {
-                    _accountHistory.value = it
-                }
-            }
-
-            fetchAccountHistory()
-        }
-    }
-
     fun clearAccountHistory() {
-        accountCache?.clearAccountHistory()
+        accountCache.tryPerformAction(
+            errorMessageIfAccountCacheNotAvailable = SERVICE_NOT_CONNECTED_ERROR_MESSAGE
+        ) { cache ->
+            cache.clearAccountHistory()
+        }
     }
 
     fun createAccount() {
-        accountCache?.apply {
+        accountCache.tryPerformAction(
+            errorMessageIfAccountCacheNotAvailable = SERVICE_NOT_CONNECTED_ERROR_MESSAGE
+        ) { cache ->
             _uiState.value = LoginUiState.CreatingAccount
-
-            viewModelScope.launch {
-                _uiState.value = accountCreationEvents.first().mapToUiState()
+            viewModelScope.launch(dispatcher) {
+                _uiState.value = cache.accountCreationEvents
+                    .onStart { cache.createNewAccount() }
+                    .first()
+                    .mapToUiState()
             }
-
-            createNewAccount()
         }
     }
 
     fun login(accountToken: String) {
-        accountCache?.apply {
+        accountCache.tryPerformAction(
+            errorMessageIfAccountCacheNotAvailable = SERVICE_NOT_CONNECTED_ERROR_MESSAGE
+        ) { cache ->
             _uiState.value = LoginUiState.Loading
-
-            viewModelScope.launch {
-                _uiState.value = loginEvents.first().result.mapToUiState()
+            viewModelScope.launch(dispatcher) {
+                _uiState.value = cache.loginEvents
+                    .onStart { cache.login(accountToken) }
+                    .map { it.result.mapToUiState() }
+                    .first()
             }
+        }
+    }
 
-            login(accountToken)
+    private fun AccountCache?.tryPerformAction(
+        errorMessageIfAccountCacheNotAvailable: String,
+        action: (AccountCache) -> Unit
+    ) {
+        if (this != null) {
+            action(this)
+        } else {
+            _uiState.value = LoginUiState.OtherError(errorMessageIfAccountCacheNotAvailable)
         }
     }
 
@@ -94,4 +125,8 @@ class LoginViewModel : ViewModel() {
             else -> LoginUiState.OtherError(errorMessage = this.toString())
         }
     }
+
+    companion object {
+        private const val SERVICE_NOT_CONNECTED_ERROR_MESSAGE = "Not connected to service!"
+    }
 }