Prevent loading when logging out

4 files changed, 37 insertions, 17 deletions

diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/AccountScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/AccountScreenTest.kt
index a40cb14d65..48ce7590a2 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/AccountScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/AccountScreenTest.kt
@@ -15,6 +15,7 @@ import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
 import net.mullvad.mullvadvpn.compose.state.PaymentState
 import net.mullvad.mullvadvpn.compose.test.PLAY_PAYMENT_INFO_ICON_TEST_TAG
+import net.mullvad.mullvadvpn.lib.model.AccountNumber
 import net.mullvad.mullvadvpn.lib.payment.model.PaymentProduct
 import net.mullvad.mullvadvpn.lib.payment.model.PaymentStatus
 import net.mullvad.mullvadvpn.lib.payment.model.ProductId
@@ -290,6 +291,6 @@ class AccountScreenTest {
 
     companion object {
         private const val DUMMY_DEVICE_NAME = "fake_name"
-        private const val DUMMY_ACCOUNT_NUMBER = "fake_number"
+        private val DUMMY_ACCOUNT_NUMBER = AccountNumber("1234123412341234")
     }
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/AccountScreen.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/AccountScreen.kt
index 4d19095a6e..d78592775b 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/AccountScreen.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/AccountScreen.kt
@@ -55,6 +55,7 @@ import net.mullvad.mullvadvpn.compose.transitions.SlideInFromBottomTransition
 import net.mullvad.mullvadvpn.compose.util.LaunchedEffectCollect
 import net.mullvad.mullvadvpn.compose.util.SecureScreenWhileInView
 import net.mullvad.mullvadvpn.compose.util.createCopyToClipboardHandle
+import net.mullvad.mullvadvpn.lib.model.AccountNumber
 import net.mullvad.mullvadvpn.lib.payment.model.PaymentProduct
 import net.mullvad.mullvadvpn.lib.payment.model.PaymentStatus
 import net.mullvad.mullvadvpn.lib.payment.model.ProductId
@@ -76,7 +77,7 @@ private fun PreviewAccountScreen() {
             state =
                 AccountUiState(
                     deviceName = "Test Name",
-                    accountNumber = "1234123412341234",
+                    accountNumber = AccountNumber("1234123412341234"),
                     accountExpiry = null,
                     showSitePayment = true,
                     billingPaymentState =
@@ -201,7 +202,10 @@ fun AccountScreen(
                     onInfoClick = navigateToDeviceInfo
                 )
 
-                AccountNumberRow(accountNumber = state.accountNumber ?: "", onCopyAccountNumber)
+                AccountNumberRow(
+                    accountNumber = state.accountNumber?.value ?: "",
+                    onCopyAccountNumber
+                )
 
                 PaidUntilRow(accountExpiry = state.accountExpiry)
             }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/AccountViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/AccountViewModel.kt
index 0a497c22f6..a42003d6e2 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/AccountViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/AccountViewModel.kt
@@ -4,13 +4,21 @@ import android.app.Activity
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.filterIsInstance
+import kotlinx.coroutines.flow.filterNotNull
+import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.flow.receiveAsFlow
 import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.compose.state.PaymentState
+import net.mullvad.mullvadvpn.lib.model.AccountData
+import net.mullvad.mullvadvpn.lib.model.AccountNumber
+import net.mullvad.mullvadvpn.lib.model.DeviceState
 import net.mullvad.mullvadvpn.lib.model.WebsiteAuthToken
 import net.mullvad.mullvadvpn.lib.payment.model.ProductId
 import net.mullvad.mullvadvpn.lib.shared.AccountRepository
@@ -30,13 +38,13 @@ class AccountViewModel(
 
     val uiState: StateFlow<AccountUiState> =
         combine(
-                deviceRepository.deviceState,
-                accountRepository.accountData,
+                deviceRepository.deviceState.filterIsInstance<DeviceState.LoggedIn>(),
+                accountData(),
                 paymentUseCase.paymentAvailability
             ) { deviceState, accountData, paymentAvailability ->
                 AccountUiState(
-                    deviceName = deviceState?.displayName() ?: "",
-                    accountNumber = deviceState?.token()?.value ?: "",
+                    deviceName = deviceState.device.displayName(),
+                    accountNumber = deviceState.accountNumber,
                     accountExpiry = accountData?.expiryDate,
                     showSitePayment = !isPlayBuild,
                     billingPaymentState = paymentAvailability?.toPaymentState()
@@ -50,6 +58,13 @@ class AccountViewModel(
         fetchPaymentAvailability()
     }
 
+    private fun accountData(): Flow<AccountData?> =
+        // Ignore nulls expect first, to avoid loading when logging out.
+        accountRepository.accountData
+            .filterNotNull()
+            .onStart<AccountData?> { emit(accountRepository.accountData.value) }
+            .distinctUntilChanged()
+
     fun onManageAccountClick() {
         viewModelScope.launch {
             accountRepository.getWebsiteAuthToken()?.let { wwwAuthToken ->
@@ -115,7 +130,7 @@ class AccountViewModel(
 
 data class AccountUiState(
     val deviceName: String?,
-    val accountNumber: String?,
+    val accountNumber: AccountNumber?,
     val accountExpiry: DateTime?,
     val showSitePayment: Boolean,
     val billingPaymentState: PaymentState? = null,
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/AccountViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/AccountViewModelTest.kt
index fd1a363e11..76c176c519 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/AccountViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/AccountViewModelTest.kt
@@ -39,11 +39,6 @@ class AccountViewModelTest {
     private val mockDeviceRepository: DeviceRepository = mockk(relaxUnitFun = true)
     private val mockPaymentUseCase: PaymentUseCase = mockk(relaxed = true)
 
-    private val deviceState: MutableStateFlow<DeviceState?> = MutableStateFlow(null)
-    private val paymentAvailability = MutableStateFlow<PaymentAvailability?>(null)
-    private val purchaseResult = MutableStateFlow<PurchaseResult?>(null)
-    private val accountExpiryState = MutableStateFlow(null)
-
     private val dummyDevice =
         Device(id = DeviceId.fromString(UUID), name = "fake_name", creationDate = DateTime.now())
     private val dummyAccountNumber: AccountNumber =
@@ -51,6 +46,14 @@ class AccountViewModelTest {
             DUMMY_DEVICE_NAME,
         )
 
+    private val deviceState: MutableStateFlow<DeviceState?> =
+        MutableStateFlow(
+            DeviceState.LoggedIn(accountNumber = dummyAccountNumber, device = dummyDevice)
+        )
+    private val paymentAvailability = MutableStateFlow<PaymentAvailability?>(null)
+    private val purchaseResult = MutableStateFlow<PurchaseResult?>(null)
+    private val accountExpiryState = MutableStateFlow(null)
+
     private lateinit var viewModel: AccountViewModel
 
     @BeforeEach
@@ -80,11 +83,10 @@ class AccountViewModelTest {
     fun `given device state LoggedIn uiState should contain accountNumber`() = runTest {
         // Act, Assert
         viewModel.uiState.test {
-            awaitItem() // Default state
             deviceState.value =
                 DeviceState.LoggedIn(accountNumber = dummyAccountNumber, device = dummyDevice)
             val result = awaitItem()
-            assertEquals(DUMMY_DEVICE_NAME, result.accountNumber)
+            assertEquals(dummyAccountNumber, result.accountNumber)
         }
     }
 
@@ -100,8 +102,6 @@ class AccountViewModelTest {
     @Test
     fun `when paymentAvailability emits ProductsUnavailable uiState should be NoPayment`() =
         runTest {
-            // Arrange in setup
-
             // Act, Assert
             viewModel.uiState.test {
                 awaitItem() // Default state