Reset purchase state when closing the add time bottom sheet

Also immediately update the expiry date after a successful purchase

3 files changed, 35 insertions, 35 deletions

diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/AddTimeBottomSheet.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/AddTimeBottomSheet.kt
index ef92d131f5..edb28e0e69 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/AddTimeBottomSheet.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/AddTimeBottomSheet.kt
@@ -132,9 +132,9 @@ fun AddTimeBottomSheet(
             onSitePaymentClick = viewModel::onManageAccountClick,
             onRetryFetchProducts = viewModel::fetchPaymentAvailability,
             onRedeemVoucherClick = onRedeemVoucherClick,
-            resetPurchaseState = { viewModel.onClosePurchaseResultDialog(false) },
+            resetPurchaseState = { viewModel.resetPurchaseResult() },
             closeSheetAndResetPurchaseState = {
-                viewModel.onClosePurchaseResultDialog(it)
+                viewModel.resetPurchaseResult()
                 onCloseBottomSheet(true)
             },
             closeBottomSheet = onCloseBottomSheet,
@@ -162,7 +162,10 @@ fun AddTimeBottomSheetContent(
         sheetState = sheetState,
         backgroundColor = backgroundColor,
         onBackgroundColor = onBackgroundColor,
-        onDismissRequest = { closeBottomSheet(false) },
+        onDismissRequest = {
+            resetPurchaseState()
+            closeBottomSheet(false)
+        },
     ) {
         when (state) {
             is Lc.Loading ->
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/AddTimeViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/AddTimeViewModel.kt
index c865207353..141e854b81 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/AddTimeViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/AddTimeViewModel.kt
@@ -7,6 +7,7 @@ import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.receiveAsFlow
 import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
@@ -26,7 +27,7 @@ import net.mullvad.mullvadvpn.viewmodel.AddMoreTimeSideEffect.OpenAccountManagem
 class AddTimeViewModel(
     private val paymentUseCase: PaymentUseCase,
     private val accountRepository: AccountRepository,
-    private val connectionProxy: ConnectionProxy,
+    connectionProxy: ConnectionProxy,
     private val isPlayBuild: Boolean,
 ) : ViewModel() {
     private val _uiSideEffect = Channel<AddMoreTimeSideEffect>()
@@ -56,6 +57,7 @@ class AddTimeViewModel(
     init {
         verifyPurchases()
         fetchPaymentAvailability()
+        handlePurchaseResultTerminatingState()
     }
 
     fun onManageAccountClick() {
@@ -73,20 +75,8 @@ class AddTimeViewModel(
         viewModelScope.launch { paymentUseCase.purchaseProduct(productId, activityProvider) }
     }
 
-    fun onClosePurchaseResultDialog(success: Boolean) {
-        // We are closing the dialog without any action, this can happen either if an error occurred
-        // during the purchase or the purchase ended successfully.
-        // If the payment was successful we want to update the account expiry. If not successful we
-        // should check payment availability and verify any purchases to handle potential errors.
-        if (success) {
-            updateAccountExpiry()
-        } else {
-            fetchPaymentAvailability()
-            verifyPurchases() // Attempt to verify again
-        }
-        viewModelScope.launch {
-            paymentUseCase.resetPurchaseResult() // So that we do not show the dialog again.
-        }
+    fun resetPurchaseResult() {
+        viewModelScope.launch { paymentUseCase.resetPurchaseResult() }
     }
 
     private fun verifyPurchases() {
@@ -97,6 +87,22 @@ class AddTimeViewModel(
         }
     }
 
+    private fun handlePurchaseResultTerminatingState() {
+        viewModelScope.launch {
+            paymentUseCase.purchaseResult
+                .filter { it?.isTerminatingState() == true }
+                .collect {
+                    // Terminating states are either errors or completed purchases.
+                    if (it is PurchaseResult.Completed) {
+                        updateAccountExpiry()
+                    } else {
+                        fetchPaymentAvailability()
+                        verifyPurchases() // Attempt to verify again
+                    }
+                }
+        }
+    }
+
     private fun updateAccountExpiry() {
         viewModelScope.launch { accountRepository.getAccountData() }
     }
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/AddTimeViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/AddTimeViewModelTest.kt
index 3d4eb07905..6ee7bd2f81 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/AddTimeViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/AddTimeViewModelTest.kt
@@ -139,44 +139,35 @@ class AddTimeViewModelTest {
     }
 
     @Test
-    fun `onClosePurchaseResultDialog with success should invoke fetchAccountExpiry on AccountRepository`() {
+    fun `purchaseState success should invoke getAccountData on AccountRepository`() {
         // Arrange
+        val purchaseResultData = PurchaseResult.Completed.Success(ProductId("one_month"))
 
         // Act
-        viewModel.onClosePurchaseResultDialog(success = true)
+        purchaseResult.value = purchaseResultData
 
         // Assert
         coVerify { mockAccountRepository.getAccountData() }
     }
 
     @Test
-    fun `onClosePurchaseResultDialog with success should invoke resetPurchaseResult on PaymentUseCase`() {
-        // Arrange
-
-        // Act
-        viewModel.onClosePurchaseResultDialog(success = true)
-
-        // Assert
-        coVerify { mockPaymentUseCase.resetPurchaseResult() }
-    }
-
-    @Test
-    fun `onClosePurchaseResultDialog with success false should invoke queryPaymentAvailability on PaymentUseCase`() {
+    fun `purchaseState error should invoke queryPaymentAvailability on PaymentUseCase`() {
         // Arrange
+        val purchaseResultData = PurchaseResult.Error.VerificationError(Throwable())
 
         // Act
-        viewModel.onClosePurchaseResultDialog(success = false)
+        purchaseResult.value = purchaseResultData
 
         // Assert
         coVerify { mockPaymentUseCase.queryPaymentAvailability() }
     }
 
     @Test
-    fun `onClosePurchaseResultDialog with success false should invoke resetPurchaseResult on PaymentUseCase`() {
+    fun `resetPurchaseResult with success should invoke resetPurchaseResult on PaymentUseCase`() {
         // Arrange
 
         // Act
-        viewModel.onClosePurchaseResultDialog(success = false)
+        viewModel.resetPurchaseResult()
 
         // Assert
         coVerify { mockPaymentUseCase.resetPurchaseResult() }