Update payment flow within the app

Also add support for 3 months in-app purchases

5 files changed, 466 insertions, 539 deletions

diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/component/AddTimeBottomSheetTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/component/AddTimeBottomSheetTest.kt
new file mode 100644
index 0000000000..095c15dafc
--- /dev/null
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/component/AddTimeBottomSheetTest.kt
@@ -0,0 +1,372 @@
+package net.mullvad.mullvadvpn.compose.component
+
+import androidx.compose.material3.ExperimentalMaterial3Api
+import androidx.compose.material3.SheetState
+import androidx.compose.material3.SheetValue
+import androidx.compose.ui.test.ExperimentalTestApi
+import androidx.compose.ui.test.onNodeWithTag
+import androidx.compose.ui.test.onNodeWithText
+import androidx.compose.ui.test.performClick
+import androidx.compose.ui.unit.Density
+import de.mannodermaus.junit5.compose.ComposeContext
+import io.mockk.every
+import io.mockk.mockk
+import io.mockk.verify
+import kotlin.Unit
+import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
+import net.mullvad.mullvadvpn.compose.setContentWithTheme
+import net.mullvad.mullvadvpn.compose.state.AddTimeUiState
+import net.mullvad.mullvadvpn.compose.state.PaymentState
+import net.mullvad.mullvadvpn.compose.state.PurchaseState
+import net.mullvad.mullvadvpn.lib.payment.ProductIds
+import net.mullvad.mullvadvpn.lib.payment.model.PaymentProduct
+import net.mullvad.mullvadvpn.lib.payment.model.PaymentStatus
+import net.mullvad.mullvadvpn.lib.payment.model.ProductId
+import net.mullvad.mullvadvpn.lib.payment.model.ProductPrice
+import net.mullvad.mullvadvpn.lib.ui.tag.PLAY_PAYMENT_INFO_ICON_TEST_TAG
+import net.mullvad.mullvadvpn.util.Lc
+import net.mullvad.mullvadvpn.util.toLc
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.RegisterExtension
+
+@OptIn(ExperimentalMaterial3Api::class)
+class AddTimeBottomSheetTest {
+    @OptIn(ExperimentalTestApi::class)
+    @JvmField
+    @RegisterExtension
+    val composeExtension = createEdgeToEdgeComposeExtension()
+
+    private fun ComposeContext.initBottomSheet(
+        state: Lc<Unit, AddTimeUiState> = Lc.Loading(Unit),
+        sheetState: SheetState =
+            SheetState(
+                skipPartiallyExpanded = true,
+                density = Density(1f),
+                initialValue = SheetValue.Expanded,
+            ),
+        onPurchaseBillingProductClick: (ProductId) -> Unit = {},
+        onPlayPaymentInfoClick: () -> Unit = {},
+        onSitePaymentClick: () -> Unit = {},
+        onRedeemVoucherClick: () -> Unit = {},
+        onRetryFetchProducts: () -> Unit = {},
+        resetPurchaseState: () -> Unit = {},
+        closeSheetAndResetPurchaseState: (Boolean) -> Unit = {},
+        closeBottomSheet: (animate: Boolean) -> Unit = {},
+    ) {
+        setContentWithTheme {
+            AddTimeBottomSheetContent(
+                state = state,
+                sheetState = sheetState,
+                onPurchaseBillingProductClick = onPurchaseBillingProductClick,
+                onPlayPaymentInfoClick = onPlayPaymentInfoClick,
+                onSitePaymentClick = onSitePaymentClick,
+                onRedeemVoucherClick = onRedeemVoucherClick,
+                onRetryFetchProducts = onRetryFetchProducts,
+                resetPurchaseState = resetPurchaseState,
+                closeSheetAndResetPurchaseState = closeSheetAndResetPurchaseState,
+                closeBottomSheet = closeBottomSheet,
+            )
+        }
+    }
+
+    @Test
+    fun testBuyCreditClick() =
+        composeExtension.use {
+            // Arrange
+            val mockedClickHandler: () -> Unit = mockk(relaxed = true)
+            initBottomSheet(
+                state =
+                    AddTimeUiState(
+                            purchaseState = null,
+                            billingPaymentState = null,
+                            showSitePayment = true,
+                            tunnelStateBlocked = false,
+                        )
+                        .toLc(),
+                onSitePaymentClick = mockedClickHandler,
+            )
+
+            // Act
+            onNodeWithText(BUY_CREDIT_TEXT).performClick()
+
+            // Assert
+            verify(exactly = 1) { mockedClickHandler.invoke() }
+        }
+
+    @Test
+    fun testRedeemVoucherClick() =
+        composeExtension.use {
+            // Arrange
+            val mockedClickHandler: () -> Unit = mockk(relaxed = true)
+            initBottomSheet(
+                state =
+                    AddTimeUiState(
+                            purchaseState = null,
+                            billingPaymentState = null,
+                            tunnelStateBlocked = false,
+                            showSitePayment = true,
+                        )
+                        .toLc(),
+                onRedeemVoucherClick = mockedClickHandler,
+            )
+
+            // Act
+            onNodeWithText("Redeem voucher").performClick()
+
+            // Assert
+            verify { mockedClickHandler.invoke() }
+        }
+
+    @Test
+    fun testShowBillingErrorPaymentButton() =
+        composeExtension.use {
+            // Arrange
+            initBottomSheet(
+                state =
+                    AddTimeUiState(
+                            purchaseState = null,
+                            billingPaymentState = PaymentState.Error.Generic,
+                            tunnelStateBlocked = false,
+                            showSitePayment = false,
+                        )
+                        .toLc()
+            )
+
+            // Assert
+            onNodeWithText("Failed to load products, please try again").assertExists()
+        }
+
+    @Test
+    fun testShowBillingPaymentAvailable() =
+        composeExtension.use {
+            // Arrange
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.productId } returns ProductId(ProductIds.OneMonth)
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.status } returns null
+            initBottomSheet(
+                state =
+                    AddTimeUiState(
+                            purchaseState = null,
+                            billingPaymentState =
+                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
+                            tunnelStateBlocked = false,
+                            showSitePayment = false,
+                        )
+                        .toLc()
+            )
+
+            // Assert
+            onNodeWithText("Add 30 days time ($10)").assertExists()
+        }
+
+    @Test
+    fun testShowPendingPayment() =
+        composeExtension.use {
+            // Arrange
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.status } returns PaymentStatus.PENDING
+            every { mockPaymentProduct.productId } returns ProductId(ProductIds.OneMonth)
+            initBottomSheet(
+                state =
+                    AddTimeUiState(
+                            purchaseState = null,
+                            billingPaymentState =
+                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
+                            tunnelStateBlocked = false,
+                            showSitePayment = false,
+                        )
+                        .toLc()
+            )
+
+            // Assert
+            onNodeWithText("Google Play payment pending, this might take some time").assertExists()
+        }
+
+    @Test
+    fun testShowPendingPaymentInfoDialog() =
+        composeExtension.use {
+            // Arrange
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.status } returns PaymentStatus.PENDING
+            every { mockPaymentProduct.productId } returns ProductId(ProductIds.OneMonth)
+            val mockNavigateToVerificationPending: () -> Unit = mockk(relaxed = true)
+            PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
+            initBottomSheet(
+                state =
+                    AddTimeUiState(
+                            purchaseState = null,
+                            billingPaymentState =
+                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
+                            tunnelStateBlocked = false,
+                            showSitePayment = false,
+                        )
+                        .toLc(),
+                onPlayPaymentInfoClick = mockNavigateToVerificationPending,
+            )
+
+            // Act
+            onNodeWithTag(PLAY_PAYMENT_INFO_ICON_TEST_TAG).performClick()
+
+            // Assert
+            verify(exactly = 1) { mockNavigateToVerificationPending.invoke() }
+        }
+
+    @Test
+    fun testShowVerificationInProgress() =
+        composeExtension.use {
+            // Arrange
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.status } returns PaymentStatus.VERIFICATION_IN_PROGRESS
+            every { mockPaymentProduct.productId } returns ProductId(ProductIds.ThreeMonths)
+            initBottomSheet(
+                state =
+                    AddTimeUiState(
+                            purchaseState = null,
+                            billingPaymentState =
+                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
+                            tunnelStateBlocked = false,
+                            showSitePayment = false,
+                        )
+                        .toLc()
+            )
+
+            // Assert
+            onNodeWithText("Verifying purchase").assertExists()
+        }
+
+    @Test
+    fun testOnPurchaseBillingProductClick() =
+        composeExtension.use {
+            // Arrange
+            val clickHandler: (ProductId) -> Unit = mockk(relaxed = true)
+            val mockPaymentProduct: PaymentProduct = mockk()
+            every { mockPaymentProduct.price } returns ProductPrice("$10")
+            every { mockPaymentProduct.productId } returns ProductId(ProductIds.OneMonth)
+            every { mockPaymentProduct.status } returns null
+            initBottomSheet(
+                state =
+                    AddTimeUiState(
+                            purchaseState = null,
+                            billingPaymentState =
+                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
+                            tunnelStateBlocked = false,
+                            showSitePayment = false,
+                        )
+                        .toLc(),
+                onPurchaseBillingProductClick = clickHandler,
+            )
+
+            // Act
+            onNodeWithText("Add 30 days time ($10)").performClick()
+
+            // Assert
+            verify { clickHandler.invoke(ProductId(ProductIds.OneMonth)) }
+        }
+
+    @Test
+    fun testShowPurchaseCompleteDialog() =
+        composeExtension.use {
+            // Arrange
+            initBottomSheet(
+                state =
+                    AddTimeUiState(
+                            purchaseState =
+                                PurchaseState.Success(ProductId(ProductIds.ThreeMonths)),
+                            billingPaymentState = null,
+                            tunnelStateBlocked = false,
+                            showSitePayment = false,
+                        )
+                        .toLc()
+            )
+
+            // Assert
+            onNodeWithText("Time added").assertExists()
+            onNodeWithText("90 days was added to your account.").assertExists()
+        }
+
+    @Test
+    fun testShowVerificationErrorDialog() =
+        composeExtension.use {
+            // Arrange
+            initBottomSheet(
+                AddTimeUiState(
+                        purchaseState = PurchaseState.VerifyingPurchase,
+                        billingPaymentState = null,
+                        tunnelStateBlocked = false,
+                        showSitePayment = false,
+                    )
+                    .toLc()
+            )
+
+            // Assert
+            onNodeWithText("Verifying purchase").assertExists()
+        }
+
+    @Test
+    fun testShowFetchProductsErrorDialog() =
+        composeExtension.use {
+            // Arrange
+            initBottomSheet(
+                state =
+                    AddTimeUiState(
+                            purchaseState = PurchaseState.Error.OtherError(ProductId("ProductId")),
+                            billingPaymentState = null,
+                            tunnelStateBlocked = false,
+                            showSitePayment = false,
+                        )
+                        .toLc()
+            )
+
+            // Assert
+            onNodeWithText(
+                    "We were unable to start the payment process, please make sure you have the latest version of Google Play."
+                )
+                .assertExists()
+        }
+
+    @Test
+    fun testDisableSitePayment() =
+        composeExtension.use {
+            // Arrange
+            initBottomSheet(
+                state =
+                    AddTimeUiState(
+                            purchaseState = null,
+                            billingPaymentState = null,
+                            tunnelStateBlocked = false,
+                            showSitePayment = false,
+                        )
+                        .toLc()
+            )
+
+            // Assert
+            onNodeWithText(BUY_CREDIT_TEXT).assertDoesNotExist()
+        }
+
+    @Test
+    fun testShowInternetBlocked() =
+        composeExtension.use {
+            // Arrange
+            initBottomSheet(
+                state =
+                    AddTimeUiState(
+                            purchaseState = null,
+                            billingPaymentState = null,
+                            tunnelStateBlocked = true,
+                            showSitePayment = true,
+                        )
+                        .toLc()
+            )
+
+            // Assert
+            onNodeWithText("The app is blocking internet, please disconnect first").assertExists()
+        }
+
+    companion object {
+        private const val BUY_CREDIT_TEXT = "Buy credit"
+    }
+}
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/PaymentDialogTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/PaymentDialogTest.kt
deleted file mode 100644
index 09a5e9dd72..0000000000
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/dialog/PaymentDialogTest.kt
+++ /dev/null
@@ -1,72 +0,0 @@
-package net.mullvad.mullvadvpn.compose.dialog
-
-import androidx.compose.ui.test.ExperimentalTestApi
-import androidx.compose.ui.test.onNodeWithText
-import de.mannodermaus.junit5.compose.ComposeContext
-import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
-import net.mullvad.mullvadvpn.compose.dialog.payment.PaymentDialog
-import net.mullvad.mullvadvpn.compose.dialog.payment.PaymentDialogData
-import net.mullvad.mullvadvpn.compose.setContentWithTheme
-import net.mullvad.mullvadvpn.lib.payment.model.ProductId
-import net.mullvad.mullvadvpn.lib.payment.model.PurchaseResult
-import net.mullvad.mullvadvpn.util.toPaymentDialogData
-import org.junit.jupiter.api.Test
-import org.junit.jupiter.api.extension.RegisterExtension
-
-class PaymentDialogTest {
-    @OptIn(ExperimentalTestApi::class)
-    @JvmField
-    @RegisterExtension
-    val composeExtension = createEdgeToEdgeComposeExtension()
-
-    private fun ComposeContext.initDialog(
-        paymentDialogData: PaymentDialogData,
-        retryPurchase: (ProductId) -> Unit = {},
-        onCloseDialog: (isPaymentSuccessful: Boolean) -> Unit = {},
-    ) {
-        setContentWithTheme {
-            PaymentDialog(
-                paymentDialogData = paymentDialogData,
-                retryPurchase = retryPurchase,
-                onCloseDialog = onCloseDialog,
-            )
-        }
-    }
-
-    @Test
-    fun testShowPurchaseCompleteDialog() =
-        composeExtension.use {
-            // Arrange
-            initDialog(paymentDialogData = PurchaseResult.Completed.Success.toPaymentDialogData()!!)
-
-            // Assert
-            onNodeWithText("Time was successfully added").assertExists()
-        }
-
-    @Test
-    fun testShowVerificationErrorDialog() =
-        composeExtension.use {
-            // Arrange
-            initDialog(
-                paymentDialogData =
-                    PurchaseResult.Error.VerificationError(null).toPaymentDialogData()!!
-            )
-
-            // Assert
-            onNodeWithText("Verifying purchase").assertExists()
-        }
-
-    @Test
-    fun testShowFetchProductsErrorDialog() =
-        composeExtension.use {
-            // Arrange
-            initDialog(
-                paymentDialogData =
-                    PurchaseResult.Error.FetchProductsError(ProductId(""), null)
-                        .toPaymentDialogData()!!
-            )
-
-            // Assert
-            onNodeWithText("Google Play unavailable").assertExists()
-        }
-}
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 cc8c8e9943..9794fefa1e 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
@@ -2,7 +2,6 @@ package net.mullvad.mullvadvpn.compose.screen
 
 import androidx.compose.material3.ExperimentalMaterial3Api
 import androidx.compose.ui.test.ExperimentalTestApi
-import androidx.compose.ui.test.onNodeWithTag
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
 import de.mannodermaus.junit5.compose.ComposeContext
@@ -10,38 +9,42 @@ import io.mockk.MockKAnnotations
 import io.mockk.every
 import io.mockk.mockk
 import io.mockk.verify
+import kotlinx.coroutines.flow.MutableStateFlow
 import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
-import net.mullvad.mullvadvpn.compose.state.PaymentState
+import net.mullvad.mullvadvpn.compose.state.AddTimeUiState
 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
-import net.mullvad.mullvadvpn.lib.payment.model.ProductPrice
-import net.mullvad.mullvadvpn.lib.ui.tag.PLAY_PAYMENT_INFO_ICON_TEST_TAG
+import net.mullvad.mullvadvpn.util.Lc
 import net.mullvad.mullvadvpn.viewmodel.AccountUiState
+import net.mullvad.mullvadvpn.viewmodel.AddTimeViewModel
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.extension.RegisterExtension
+import org.koin.core.context.loadKoinModules
+import org.koin.core.module.dsl.viewModel
+import org.koin.dsl.module
 
 @ExperimentalTestApi
 @OptIn(ExperimentalMaterial3Api::class)
 class AccountScreenTest {
     @JvmField @RegisterExtension val composeExtension = createEdgeToEdgeComposeExtension()
 
+    private val addTimeViewModel: AddTimeViewModel = mockk(relaxed = true)
+
     @BeforeEach
     fun setup() {
         MockKAnnotations.init(this)
+        loadKoinModules(module { viewModel { addTimeViewModel } })
+        every { addTimeViewModel.uiState } returns
+            MutableStateFlow<Lc<Unit, AddTimeUiState>>(Lc.Loading(Unit))
     }
 
     private fun ComposeContext.initScreen(
-        state: AccountUiState = AccountUiState.default(),
+        state: AccountUiState? = null,
         onCopyAccountNumber: (String) -> Unit = {},
         onRedeemVoucherClick: () -> Unit = {},
-        onManageAccountClick: () -> Unit = {},
         onLogoutClick: () -> Unit = {},
-        onPurchaseBillingProductClick: (productId: ProductId) -> Unit = {},
-        navigateToVerificationPendingDialog: () -> Unit = {},
+        onPlayPaymentInfoClick: () -> Unit = {},
         onBackClick: () -> Unit = {},
         onManageDevicesClick: () -> Unit = {},
     ) {
@@ -49,13 +52,11 @@ class AccountScreenTest {
             AccountScreen(
                 state = state,
                 onCopyAccountNumber = onCopyAccountNumber,
-                onRedeemVoucherClick = onRedeemVoucherClick,
-                onManageAccountClick = onManageAccountClick,
+                onManageDevicesClick = onManageDevicesClick,
                 onLogoutClick = onLogoutClick,
-                onPurchaseBillingProductClick = onPurchaseBillingProductClick,
-                navigateToVerificationPendingDialog = navigateToVerificationPendingDialog,
+                onRedeemVoucherClick = onRedeemVoucherClick,
+                onPlayPaymentInfoClick = onPlayPaymentInfoClick,
                 onBackClick = onBackClick,
-                onManageDevicesClick = onManageDevicesClick,
             )
         }
     }
@@ -70,44 +71,17 @@ class AccountScreenTest {
                         deviceName = DUMMY_DEVICE_NAME,
                         accountNumber = DUMMY_ACCOUNT_NUMBER,
                         accountExpiry = null,
-                        showSitePayment = false,
                         showLogoutLoading = false,
-                        showManageAccountLoading = false,
+                        verificationPending = false,
                     )
             )
 
             // Assert
-            onNodeWithText("Redeem voucher").assertExists()
             onNodeWithText("Log out").assertExists()
         }
 
     @Test
-    fun testManageAccountClick() =
-        composeExtension.use {
-            // Arrange
-            val mockedClickHandler: () -> Unit = mockk(relaxed = true)
-            initScreen(
-                state =
-                    AccountUiState(
-                        deviceName = DUMMY_DEVICE_NAME,
-                        accountNumber = DUMMY_ACCOUNT_NUMBER,
-                        accountExpiry = null,
-                        showSitePayment = true,
-                        showLogoutLoading = false,
-                        showManageAccountLoading = false,
-                    ),
-                onManageAccountClick = mockedClickHandler,
-            )
-
-            // Act
-            onNodeWithText("Manage account").performClick()
-
-            // Assert
-            verify(exactly = 1) { mockedClickHandler.invoke() }
-        }
-
-    @Test
-    fun testRedeemVoucherClick() =
+    fun testLogoutClick() =
         composeExtension.use {
             // Arrange
             val mockedClickHandler: () -> Unit = mockk(relaxed = true)
@@ -117,169 +91,38 @@ class AccountScreenTest {
                         deviceName = DUMMY_DEVICE_NAME,
                         accountNumber = DUMMY_ACCOUNT_NUMBER,
                         accountExpiry = null,
-                        showSitePayment = false,
                         showLogoutLoading = false,
-                        showManageAccountLoading = false,
+                        verificationPending = false,
                     ),
-                onRedeemVoucherClick = mockedClickHandler,
+                onLogoutClick = mockedClickHandler,
             )
 
             // Act
-            onNodeWithText("Redeem voucher").performClick()
+            onNodeWithText("Log out").performClick()
 
             // Assert
             verify { mockedClickHandler.invoke() }
         }
 
     @Test
-    fun testLogoutClick() =
+    fun testShowVerificationInProgress() =
         composeExtension.use {
             // Arrange
-            val mockedClickHandler: () -> Unit = mockk(relaxed = true)
             initScreen(
                 state =
                     AccountUiState(
                         deviceName = DUMMY_DEVICE_NAME,
                         accountNumber = DUMMY_ACCOUNT_NUMBER,
                         accountExpiry = null,
-                        showSitePayment = false,
                         showLogoutLoading = false,
-                        showManageAccountLoading = false,
-                    ),
-                onLogoutClick = mockedClickHandler,
-            )
-
-            // Act
-            onNodeWithText("Log out").performClick()
-
-            // Assert
-            verify { mockedClickHandler.invoke() }
-        }
-
-    @Test
-    fun testShowBillingErrorPaymentButton() =
-        composeExtension.use {
-            // Arrange
-            initScreen(
-                state =
-                    AccountUiState.default().copy(billingPaymentState = PaymentState.Error.Billing)
-            )
-
-            // Assert
-            onNodeWithText("Add 30 days time").assertExists()
-        }
-
-    @Test
-    fun testShowBillingPaymentAvailable() =
-        composeExtension.use {
-            // Arrange
-            val mockPaymentProduct: PaymentProduct = mockk()
-            every { mockPaymentProduct.price } returns ProductPrice("$10")
-            every { mockPaymentProduct.status } returns null
-            initScreen(
-                state =
-                    AccountUiState.default()
-                        .copy(
-                            billingPaymentState =
-                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                        )
-            )
-
-            // Assert
-            onNodeWithText("Add 30 days time ($10)").assertExists()
-        }
-
-    @Test
-    fun testShowPendingPayment() =
-        composeExtension.use {
-            // Arrange
-            val mockPaymentProduct: PaymentProduct = mockk()
-            every { mockPaymentProduct.price } returns ProductPrice("$10")
-            every { mockPaymentProduct.status } returns PaymentStatus.PENDING
-            initScreen(
-                state =
-                    AccountUiState.default()
-                        .copy(
-                            billingPaymentState =
-                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                        )
+                        verificationPending = true,
+                    )
             )
 
             // Assert
             onNodeWithText("Google Play payment pending").assertExists()
         }
 
-    @Test
-    fun testShowPendingPaymentInfoDialog() =
-        composeExtension.use {
-            // Arrange
-            val mockPaymentProduct: PaymentProduct = mockk()
-            every { mockPaymentProduct.price } returns ProductPrice("$10")
-            every { mockPaymentProduct.status } returns PaymentStatus.PENDING
-            val mockNavigateToVerificationPending: () -> Unit = mockk(relaxed = true)
-            initScreen(
-                state =
-                    AccountUiState.default()
-                        .copy(
-                            billingPaymentState =
-                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                        ),
-                navigateToVerificationPendingDialog = mockNavigateToVerificationPending,
-            )
-
-            // Act
-            onNodeWithTag(PLAY_PAYMENT_INFO_ICON_TEST_TAG).performClick()
-
-            // Assert
-            verify(exactly = 1) { mockNavigateToVerificationPending.invoke() }
-        }
-
-    @Test
-    fun testShowVerificationInProgress() =
-        composeExtension.use {
-            // Arrange
-            val mockPaymentProduct: PaymentProduct = mockk()
-            every { mockPaymentProduct.price } returns ProductPrice("$10")
-            every { mockPaymentProduct.status } returns PaymentStatus.VERIFICATION_IN_PROGRESS
-            initScreen(
-                state =
-                    AccountUiState.default()
-                        .copy(
-                            billingPaymentState =
-                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                        )
-            )
-
-            // Assert
-            onNodeWithText("Verifying purchase").assertExists()
-        }
-
-    @Test
-    fun testOnPurchaseBillingProductClick() =
-        composeExtension.use {
-            // Arrange
-            val clickHandler: (ProductId) -> Unit = mockk(relaxed = true)
-            val mockPaymentProduct: PaymentProduct = mockk()
-            every { mockPaymentProduct.price } returns ProductPrice("$10")
-            every { mockPaymentProduct.productId } returns ProductId("PRODUCT_ID")
-            every { mockPaymentProduct.status } returns null
-            initScreen(
-                state =
-                    AccountUiState.default()
-                        .copy(
-                            billingPaymentState =
-                                PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                        ),
-                onPurchaseBillingProductClick = clickHandler,
-            )
-
-            // Act
-            onNodeWithText("Add 30 days time ($10)").performClick()
-
-            // Assert
-            verify { clickHandler.invoke(ProductId("PRODUCT_ID")) }
-        }
-
     companion object {
         private const val DUMMY_DEVICE_NAME = "fake_name"
         private val DUMMY_ACCOUNT_NUMBER = AccountNumber("1234123412341234")
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/OutOfTimeScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/OutOfTimeScreenTest.kt
index 192daa7199..451c02309f 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/OutOfTimeScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/OutOfTimeScreenTest.kt
@@ -10,50 +10,53 @@ import io.mockk.MockKAnnotations
 import io.mockk.every
 import io.mockk.mockk
 import io.mockk.verify
+import kotlinx.coroutines.flow.MutableStateFlow
 import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
+import net.mullvad.mullvadvpn.compose.state.AddTimeUiState
 import net.mullvad.mullvadvpn.compose.state.OutOfTimeUiState
-import net.mullvad.mullvadvpn.compose.state.PaymentState
 import net.mullvad.mullvadvpn.lib.model.TunnelState
-import net.mullvad.mullvadvpn.lib.payment.model.PaymentProduct
-import net.mullvad.mullvadvpn.lib.payment.model.PaymentStatus
-import net.mullvad.mullvadvpn.lib.payment.model.ProductId
-import net.mullvad.mullvadvpn.lib.payment.model.ProductPrice
 import net.mullvad.mullvadvpn.lib.ui.tag.PLAY_PAYMENT_INFO_ICON_TEST_TAG
+import net.mullvad.mullvadvpn.util.Lc
+import net.mullvad.mullvadvpn.viewmodel.AddTimeViewModel
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.extension.RegisterExtension
+import org.koin.core.context.loadKoinModules
+import org.koin.core.module.dsl.viewModel
+import org.koin.dsl.module
 
 @OptIn(ExperimentalTestApi::class)
 class OutOfTimeScreenTest {
     @JvmField @RegisterExtension val composeExtension = createEdgeToEdgeComposeExtension()
 
+    private val addTimeViewModel: AddTimeViewModel = mockk(relaxed = true)
+
     @BeforeEach
     fun setup() {
         MockKAnnotations.init(this)
+        loadKoinModules(module { viewModel { addTimeViewModel } })
+        every { addTimeViewModel.uiState } returns
+            MutableStateFlow<Lc<Unit, AddTimeUiState>>(Lc.Loading(Unit))
     }
 
     private fun ComposeContext.initScreen(
         state: OutOfTimeUiState = OutOfTimeUiState(),
         onDisconnectClick: () -> Unit = {},
-        onSitePaymentClick: () -> Unit = {},
         onRedeemVoucherClick: () -> Unit = {},
         onSettingsClick: () -> Unit = {},
         onAccountClick: () -> Unit = {},
-        onPurchaseBillingProductClick: (ProductId) -> Unit = {},
-        navigateToVerificationPendingDialog: () -> Unit = {},
+        onPlayPaymentInfoClick: () -> Unit = {},
     ) {
 
         setContentWithTheme {
             OutOfTimeScreen(
                 state = state,
                 onDisconnectClick = onDisconnectClick,
-                onSitePaymentClick = onSitePaymentClick,
                 onRedeemVoucherClick = onRedeemVoucherClick,
                 onSettingsClick = onSettingsClick,
                 onAccountClick = onAccountClick,
-                onPurchaseBillingProductClick = onPurchaseBillingProductClick,
-                navigateToVerificationPendingDialog = navigateToVerificationPendingDialog,
+                onPlayPaymentInfoClick = onPlayPaymentInfoClick,
             )
         }
     }
@@ -70,7 +73,6 @@ class OutOfTimeScreenTest {
                     substring = true,
                 )
                 .assertDoesNotExist()
-            onNodeWithText("Buy credit").assertDoesNotExist()
         }
 
     @Test
@@ -91,40 +93,6 @@ class OutOfTimeScreenTest {
         }
 
     @Test
-    fun testClickSitePaymentButton() =
-        composeExtension.use {
-            // Arrange
-            val mockClickListener: () -> Unit = mockk(relaxed = true)
-            initScreen(
-                state = OutOfTimeUiState(deviceName = "", showSitePayment = true),
-                onSitePaymentClick = mockClickListener,
-            )
-
-            // Act
-            onNodeWithText("Buy credit").performClick()
-
-            // Assert
-            verify(exactly = 1) { mockClickListener.invoke() }
-        }
-
-    @Test
-    fun testClickRedeemVoucher() =
-        composeExtension.use {
-            // Arrange
-            val mockClickListener: () -> Unit = mockk(relaxed = true)
-            initScreen(
-                state = OutOfTimeUiState(deviceName = "", showSitePayment = true),
-                onRedeemVoucherClick = mockClickListener,
-            )
-
-            // Act
-            onNodeWithText("Redeem voucher").performClick()
-
-            // Assert
-            verify(exactly = 1) { mockClickListener.invoke() }
-        }
-
-    @Test
     fun testClickDisconnect() =
         composeExtension.use {
             // Arrange
@@ -147,77 +115,13 @@ class OutOfTimeScreenTest {
         }
 
     @Test
-    fun testShowBillingErrorPaymentButton() =
-        composeExtension.use {
-            // Arrange
-            initScreen(
-                state =
-                    OutOfTimeUiState(
-                        showSitePayment = true,
-                        billingPaymentState = PaymentState.Error.Billing,
-                    )
-            )
-
-            // Assert
-            onNodeWithText("Add 30 days time").assertExists()
-        }
-
-    @Test
-    fun testShowBillingPaymentAvailable() =
-        composeExtension.use {
-            // Arrange
-            val mockPaymentProduct: PaymentProduct = mockk()
-            every { mockPaymentProduct.price } returns ProductPrice("$10")
-            every { mockPaymentProduct.status } returns null
-            initScreen(
-                state =
-                    OutOfTimeUiState(
-                        showSitePayment = true,
-                        billingPaymentState =
-                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
-                    )
-            )
-
-            // Assert
-            onNodeWithText("Add 30 days time ($10)").assertExists()
-        }
-
-    @Test
-    fun testShowPendingPayment() =
-        composeExtension.use {
-            // Arrange
-            val mockPaymentProduct: PaymentProduct = mockk()
-            every { mockPaymentProduct.price } returns ProductPrice("$10")
-            every { mockPaymentProduct.status } returns PaymentStatus.PENDING
-            initScreen(
-                state =
-                    OutOfTimeUiState(
-                        showSitePayment = true,
-                        billingPaymentState =
-                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
-                    )
-            )
-
-            // Assert
-            onNodeWithText("Google Play payment pending").assertExists()
-        }
-
-    @Test
     fun testShowPendingPaymentInfoDialog() =
         composeExtension.use {
             // Arrange
-            val mockPaymentProduct: PaymentProduct = mockk()
-            every { mockPaymentProduct.price } returns ProductPrice("$10")
-            every { mockPaymentProduct.status } returns PaymentStatus.PENDING
-            val mockNavigateToVerificationPending: () -> Unit = mockk(relaxed = true)
+            val mockOnPlayPaymentInfoClick: () -> Unit = mockk(relaxed = true)
             initScreen(
-                state =
-                    OutOfTimeUiState(
-                        showSitePayment = true,
-                        billingPaymentState =
-                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
-                    ),
-                navigateToVerificationPendingDialog = mockNavigateToVerificationPending,
+                state = OutOfTimeUiState(showSitePayment = true, verificationPending = true),
+                onPlayPaymentInfoClick = mockOnPlayPaymentInfoClick,
             )
 
             // Act
@@ -225,52 +129,16 @@ class OutOfTimeScreenTest {
             onNodeWithTag(PLAY_PAYMENT_INFO_ICON_TEST_TAG).assertExists()
 
             // Assert
-            verify(exactly = 1) { mockNavigateToVerificationPending.invoke() }
+            verify(exactly = 1) { mockOnPlayPaymentInfoClick.invoke() }
         }
 
     @Test
     fun testShowVerificationInProgress() =
         composeExtension.use {
             // Arrange
-            val mockPaymentProduct: PaymentProduct = mockk()
-            every { mockPaymentProduct.price } returns ProductPrice("$10")
-            every { mockPaymentProduct.status } returns PaymentStatus.VERIFICATION_IN_PROGRESS
-            initScreen(
-                state =
-                    OutOfTimeUiState(
-                        billingPaymentState =
-                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
-                        showSitePayment = true,
-                    )
-            )
-
-            // Assert
-            onNodeWithText("Verifying purchase").assertExists()
-        }
-
-    @Test
-    fun testOnPurchaseBillingProductClick() =
-        composeExtension.use {
-            // Arrange
-            val clickHandler: (ProductId) -> Unit = mockk(relaxed = true)
-            val mockPaymentProduct: PaymentProduct = mockk()
-            every { mockPaymentProduct.price } returns ProductPrice("$10")
-            every { mockPaymentProduct.productId } returns ProductId("PRODUCT_ID")
-            every { mockPaymentProduct.status } returns null
-            initScreen(
-                state =
-                    OutOfTimeUiState(
-                        billingPaymentState =
-                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct)),
-                        showSitePayment = true,
-                    ),
-                onPurchaseBillingProductClick = clickHandler,
-            )
-
-            // Act
-            onNodeWithText("Add 30 days time ($10)").performClick()
+            initScreen(state = OutOfTimeUiState(showSitePayment = true, verificationPending = true))
 
             // Assert
-            verify { clickHandler(ProductId("PRODUCT_ID")) }
+            onNodeWithText("Google Play payment pending").assertExists()
         }
 }
diff --git a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/WelcomeScreenTest.kt b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/WelcomeScreenTest.kt
index b3ac54fb73..cf25afee16 100644
--- a/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/WelcomeScreenTest.kt
+++ b/android/app/src/androidTest/kotlin/net/mullvad/mullvadvpn/compose/screen/WelcomeScreenTest.kt
@@ -9,52 +9,56 @@ import io.mockk.MockKAnnotations
 import io.mockk.every
 import io.mockk.mockk
 import io.mockk.verify
+import kotlinx.coroutines.flow.MutableStateFlow
 import net.mullvad.mullvadvpn.compose.createEdgeToEdgeComposeExtension
 import net.mullvad.mullvadvpn.compose.setContentWithTheme
-import net.mullvad.mullvadvpn.compose.state.PaymentState
+import net.mullvad.mullvadvpn.compose.state.AddTimeUiState
 import net.mullvad.mullvadvpn.compose.state.WelcomeUiState
 import net.mullvad.mullvadvpn.lib.model.AccountNumber
 import net.mullvad.mullvadvpn.lib.model.TunnelState
-import net.mullvad.mullvadvpn.lib.payment.model.PaymentProduct
-import net.mullvad.mullvadvpn.lib.payment.model.PaymentStatus
-import net.mullvad.mullvadvpn.lib.payment.model.ProductId
-import net.mullvad.mullvadvpn.lib.payment.model.ProductPrice
 import net.mullvad.mullvadvpn.lib.ui.tag.PLAY_PAYMENT_INFO_ICON_TEST_TAG
+import net.mullvad.mullvadvpn.util.Lc
+import net.mullvad.mullvadvpn.util.toLc
+import net.mullvad.mullvadvpn.viewmodel.AddTimeViewModel
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.extension.RegisterExtension
+import org.koin.core.context.loadKoinModules
+import org.koin.core.module.dsl.viewModel
+import org.koin.dsl.module
 
 @OptIn(ExperimentalTestApi::class)
 class WelcomeScreenTest {
     @JvmField @RegisterExtension val composeExtension = createEdgeToEdgeComposeExtension()
 
+    private val addTimeViewModel: AddTimeViewModel = mockk(relaxed = true)
+
     @BeforeEach
     fun setup() {
         MockKAnnotations.init(this)
+        loadKoinModules(module { viewModel { addTimeViewModel } })
+        every { addTimeViewModel.uiState } returns
+            MutableStateFlow<Lc<Unit, AddTimeUiState>>(Lc.Loading(Unit))
     }
 
     private fun ComposeContext.initScreen(
-        state: WelcomeUiState = WelcomeUiState(),
-        onSitePaymentClick: () -> Unit = {},
+        state: Lc<Unit, WelcomeUiState> = Lc.Loading(Unit),
         onRedeemVoucherClick: () -> Unit = {},
         onSettingsClick: () -> Unit = {},
         onAccountClick: () -> Unit = {},
-        onPurchaseBillingProductClick: (productId: ProductId) -> Unit = {},
         onDisconnectClick: () -> Unit = {},
         navigateToDeviceInfoDialog: () -> Unit = {},
-        navigateToVerificationPendingDialog: () -> Unit = {},
+        onPlayPaymentInfoClick: () -> Unit = {},
     ) {
         setContentWithTheme {
             WelcomeScreen(
                 state = state,
-                onSitePaymentClick = onSitePaymentClick,
                 onRedeemVoucherClick = onRedeemVoucherClick,
                 onSettingsClick = onSettingsClick,
                 onAccountClick = onAccountClick,
-                onPurchaseBillingProductClick = onPurchaseBillingProductClick,
                 navigateToDeviceInfoDialog = navigateToDeviceInfoDialog,
-                navigateToVerificationPendingDialog = navigateToVerificationPendingDialog,
                 onDisconnectClick = onDisconnectClick,
+                onPlayPaymentInfoClick = onPlayPaymentInfoClick,
             )
         }
     }
@@ -82,7 +86,6 @@ class WelcomeScreenTest {
                     substring = true,
                 )
                 .assertDoesNotExist()
-            onNodeWithText("Buy credit").assertDoesNotExist()
         }
 
     @Test
@@ -91,108 +94,38 @@ class WelcomeScreenTest {
             // Arrange
             val rawAccountNumber = AccountNumber("1111222233334444")
             val expectedAccountNumber = "1111 2222 3333 4444"
-            initScreen(state = WelcomeUiState(accountNumber = rawAccountNumber))
-
-            // Assert
-            onNodeWithText(expectedAccountNumber).assertExists()
-        }
-
-    @Test
-    fun testClickSitePaymentButton() =
-        composeExtension.use {
-            // Arrange
-            val mockClickListener: () -> Unit = mockk(relaxed = true)
-            initScreen(
-                state = WelcomeUiState(showSitePayment = true),
-                onSitePaymentClick = mockClickListener,
-            )
-
-            // Act
-            onNodeWithText("Buy credit").performClick()
-
-            // Assert
-            verify(exactly = 1) { mockClickListener.invoke() }
-        }
-
-    @Test
-    fun testClickRedeemVoucher() =
-        composeExtension.use {
-            // Arrange
-            val mockClickListener: () -> Unit = mockk(relaxed = true)
-            initScreen(state = WelcomeUiState(), onRedeemVoucherClick = mockClickListener)
-
-            // Act
-            onNodeWithText("Redeem voucher").performClick()
-
-            // Assert
-            verify(exactly = 1) { mockClickListener.invoke() }
-        }
-
-    @Test
-    fun testShowBillingErrorPaymentButton() =
-        composeExtension.use {
-            // Arrange
-            initScreen(
-                state = WelcomeUiState().copy(billingPaymentState = PaymentState.Error.Billing)
-            )
-
-            // Assert
-            onNodeWithText("Add 30 days time").assertExists()
-        }
-
-    @Test
-    fun testShowBillingPaymentAvailable() =
-        composeExtension.use {
-            // Arrange
-            val mockPaymentProduct: PaymentProduct = mockk()
-            every { mockPaymentProduct.price } returns ProductPrice("$10")
-            every { mockPaymentProduct.status } returns null
             initScreen(
                 state =
                     WelcomeUiState(
-                        billingPaymentState =
-                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                    )
+                            tunnelState = TunnelState.Disconnected(),
+                            accountNumber = rawAccountNumber,
+                            deviceName = null,
+                            showSitePayment = false,
+                            verificationPending = false,
+                        )
+                        .toLc()
             )
 
             // Assert
-            onNodeWithText("Add 30 days time ($10)").assertExists()
-        }
-
-    @Test
-    fun testShowPendingPayment() =
-        composeExtension.use {
-            // Arrange
-            val mockPaymentProduct: PaymentProduct = mockk()
-            every { mockPaymentProduct.price } returns ProductPrice("$10")
-            every { mockPaymentProduct.status } returns PaymentStatus.PENDING
-            initScreen(
-                state =
-                    WelcomeUiState(
-                        billingPaymentState =
-                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                    )
-            )
-
-            // Assert
-            onNodeWithText("Google Play payment pending").assertExists()
+            onNodeWithText(expectedAccountNumber).assertExists()
         }
 
     @Test
     fun testShowPendingPaymentInfoDialog() =
         composeExtension.use {
             // Arrange
-            val mockPaymentProduct: PaymentProduct = mockk()
-            every { mockPaymentProduct.price } returns ProductPrice("$10")
-            every { mockPaymentProduct.status } returns PaymentStatus.PENDING
             val mockShowPendingInfo = mockk<() -> Unit>(relaxed = true)
             initScreen(
                 state =
                     WelcomeUiState(
-                        billingPaymentState =
-                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                    ),
-                navigateToVerificationPendingDialog = mockShowPendingInfo,
+                            tunnelState = TunnelState.Disconnected(),
+                            accountNumber = null,
+                            deviceName = null,
+                            showSitePayment = false,
+                            verificationPending = true,
+                        )
+                        .toLc(),
+                onPlayPaymentInfoClick = mockShowPendingInfo,
             )
 
             // Act
@@ -206,45 +139,20 @@ class WelcomeScreenTest {
     fun testShowVerificationInProgress() =
         composeExtension.use {
             // Arrange
-            val mockPaymentProduct: PaymentProduct = mockk()
-            every { mockPaymentProduct.price } returns ProductPrice("$10")
-            every { mockPaymentProduct.status } returns PaymentStatus.VERIFICATION_IN_PROGRESS
-
             initScreen(
                 state =
                     WelcomeUiState(
-                        billingPaymentState =
-                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                    )
+                            tunnelState = TunnelState.Disconnected(),
+                            accountNumber = null,
+                            deviceName = null,
+                            showSitePayment = false,
+                            verificationPending = true,
+                        )
+                        .toLc()
             )
 
             // Assert
-            onNodeWithText("Verifying purchase").assertExists()
-        }
-
-    @Test
-    fun testOnPurchaseBillingProductClick() =
-        composeExtension.use {
-            // Arrange
-            val clickHandler: (ProductId) -> Unit = mockk(relaxed = true)
-            val mockPaymentProduct: PaymentProduct = mockk()
-            every { mockPaymentProduct.price } returns ProductPrice("$10")
-            every { mockPaymentProduct.productId } returns ProductId("PRODUCT_ID")
-            every { mockPaymentProduct.status } returns null
-            initScreen(
-                state =
-                    WelcomeUiState(
-                        billingPaymentState =
-                            PaymentState.PaymentAvailable(listOf(mockPaymentProduct))
-                    ),
-                onPurchaseBillingProductClick = clickHandler,
-            )
-
-            // Act
-            onNodeWithText("Add 30 days time ($10)").performClick()
-
-            // Assert
-            verify { clickHandler(ProductId("PRODUCT_ID")) }
+            onNodeWithText("Google Play payment pending").assertExists()
         }
 
     @Test
@@ -255,7 +163,15 @@ class WelcomeScreenTest {
             val tunnelState: TunnelState = mockk(relaxed = true)
             every { tunnelState.isSecured() } returns true
             initScreen(
-                state = WelcomeUiState(tunnelState = tunnelState),
+                state =
+                    WelcomeUiState(
+                            tunnelState = tunnelState,
+                            accountNumber = null,
+                            deviceName = null,
+                            showSitePayment = false,
+                            verificationPending = false,
+                        )
+                        .toLc(),
                 onDisconnectClick = clickHandler,
             )