Improve device list error handling

2 files changed, 50 insertions, 6 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 38fa9117f7..e71dcb8416 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
@@ -41,6 +41,7 @@ val uiModule = module {
     }
 
     single { ServiceConnectionManager(androidContext()) }
+    single { androidContext().resources }
 
     single { AccountExpiryNotification(get()) }
     single { TunnelStateNotification(get()) }
@@ -50,7 +51,7 @@ val uiModule = module {
     single { DeviceRepository(get()) }
     viewModel { LoginViewModel(get(), get()) }
     viewModel { DeviceRevokedViewModel(get(), get()) }
-    viewModel { DeviceListViewModel(get()) }
+    viewModel { DeviceListViewModel(get(), get()) }
 }
 
 const val APPS_SCOPE = "APPS_SCOPE"
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceListViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceListViewModel.kt
index a36727a8dc..0057457c5e 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceListViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceListViewModel.kt
@@ -1,24 +1,36 @@
 package net.mullvad.mullvadvpn.viewmodel
 
+import android.content.res.Resources
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.asSharedFlow
 import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.filter
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.onSubscription
 import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
+import kotlinx.coroutines.withTimeoutOrNull
+import net.mullvad.mullvadvpn.R
 import net.mullvad.mullvadvpn.compose.state.DeviceListItemUiState
 import net.mullvad.mullvadvpn.compose.state.DeviceListUiState
 import net.mullvad.mullvadvpn.model.Device
 import net.mullvad.mullvadvpn.model.DeviceList
+import net.mullvad.mullvadvpn.model.RemoveDeviceResult
 import net.mullvad.mullvadvpn.ui.serviceconnection.DeviceRepository
-import net.mullvad.mullvadvpn.util.safeLet
 
 typealias DeviceId = String
 
 class DeviceListViewModel(
-    private val deviceRepository: DeviceRepository
+    private val deviceRepository: DeviceRepository,
+    private val resources: Resources,
+    private val dispatcher: CoroutineDispatcher = Dispatchers.Default
 ) : ViewModel() {
     private val _stagedDeviceId = MutableStateFlow<DeviceId?>(null)
 
@@ -59,9 +71,36 @@ class DeviceListViewModel(
     }
 
     fun confirmRemovalOfStagedDevice() {
-        safeLet(accountToken, _stagedDeviceId.value) { token, deviceId ->
-            deviceRepository.removeDevice(token, deviceId)
-            _stagedDeviceId.value = null
+        val token = accountToken
+        val stagedDeviceId = _stagedDeviceId.value
+
+        if (token != null && stagedDeviceId != null) {
+            viewModelScope.launch {
+                withContext(dispatcher) {
+                    val result = withTimeoutOrNull(DEVICE_REMOVAL_TIMEOUT_MILLIS) {
+                        deviceRepository.deviceRemovalEvent
+                            .onSubscription {
+                                clearStagedDevice()
+                                deviceRepository.removeDevice(token, stagedDeviceId)
+                            }
+                            .filter { (deviceId, result) ->
+                                deviceId == stagedDeviceId && result == RemoveDeviceResult.Ok
+                            }
+                            .first()
+                    }
+
+                    if (result == null) {
+                        _toastMessages.tryEmit(
+                            resources.getString(R.string.failed_to_remove_device)
+                        )
+                        refreshDeviceList()
+                    }
+                }
+            }
+        } else {
+            _toastMessages.tryEmit(resources.getString(R.string.error_occurred))
+            clearStagedDevice()
+            refreshDeviceList()
         }
     }
 
@@ -70,4 +109,8 @@ class DeviceListViewModel(
     fun refreshDeviceList() = accountToken?.let { token ->
         deviceRepository.refreshDeviceList(token)
     }
+
+    companion object {
+        private const val DEVICE_REMOVAL_TIMEOUT_MILLIS = 5000L
+    }
 }