Merge branch 'migrate-to-compose-navigation-droid-173'

11 files changed, 366 insertions, 273 deletions

diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/OutOfTimeUseCaseTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/OutOfTimeUseCaseTest.kt
new file mode 100644
index 0000000000..74683813ae
--- /dev/null
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/usecase/OutOfTimeUseCaseTest.kt
@@ -0,0 +1,128 @@
+package net.mullvad.mullvadvpn.usecase
+
+import app.cash.turbine.test
+import io.mockk.every
+import io.mockk.mockk
+import kotlin.test.assertEquals
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.test.runTest
+import net.mullvad.mullvadvpn.lib.ipc.Event
+import net.mullvad.mullvadvpn.lib.ipc.MessageHandler
+import net.mullvad.mullvadvpn.lib.ipc.events
+import net.mullvad.mullvadvpn.model.AccountExpiry
+import net.mullvad.mullvadvpn.model.TunnelState
+import net.mullvad.mullvadvpn.repository.AccountRepository
+import net.mullvad.talpid.tunnel.ErrorState
+import net.mullvad.talpid.tunnel.ErrorStateCause
+import org.joda.time.DateTime
+import org.junit.Before
+import org.junit.Test
+
+class OutOfTimeUseCaseTest {
+    private val mockAccountRepository: AccountRepository = mockk()
+    private val mockMessageHandler: MessageHandler = mockk()
+
+    private val events = MutableSharedFlow<Event.TunnelStateChange>()
+    private val expiry = MutableStateFlow<AccountExpiry>(AccountExpiry.Missing)
+
+    lateinit var outOfTimeUseCase: OutOfTimeUseCase
+
+    @Before
+    fun setup() {
+        every { mockAccountRepository.accountExpiryState } returns expiry
+        every { mockMessageHandler.events<Event.TunnelStateChange>() } returns events
+        outOfTimeUseCase = OutOfTimeUseCase(mockAccountRepository, mockMessageHandler)
+    }
+
+    @Test
+    fun `No events should result in no expiry`() = runTest {
+        // Arrange
+        // Act, Assert
+        outOfTimeUseCase.isOutOfTime().test { assertEquals(null, awaitItem()) }
+    }
+
+    @Test
+    fun `Tunnel is blocking because out of time should emit true`() = runTest {
+        // Arrange
+        // Act, Assert
+        val errorStateCause = ErrorStateCause.AuthFailed("[EXPIRED_ACCOUNT]")
+        val tunnelStateError = TunnelState.Error(ErrorState(errorStateCause, true))
+        val errorChange = Event.TunnelStateChange(tunnelStateError)
+
+        outOfTimeUseCase.isOutOfTime().test {
+            assertEquals(null, awaitItem())
+            events.emit(errorChange)
+            assertEquals(true, awaitItem())
+        }
+    }
+
+    @Test
+    fun `Tunnel is connected should emit false`() = runTest {
+        // Arrange
+        val expiredAccountExpiry = AccountExpiry.Available(DateTime.now().plusDays(1))
+        val tunnelStateChanges =
+            listOf(
+                    TunnelState.Disconnected,
+                    TunnelState.Connected(mockk(), null),
+                    TunnelState.Connecting(null, null),
+                    TunnelState.Disconnecting(mockk()),
+                    TunnelState.Error(ErrorState(ErrorStateCause.StartTunnelError, false)),
+                )
+                .map(Event::TunnelStateChange)
+
+        // Act, Assert
+        outOfTimeUseCase.isOutOfTime().test {
+            assertEquals(null, awaitItem())
+            events.emit(tunnelStateChanges.first())
+            expiry.emit(expiredAccountExpiry)
+            assertEquals(false, awaitItem())
+
+            tunnelStateChanges.forEach { events.emit(it) }
+
+            // Should not emit again
+            expectNoEvents()
+        }
+    }
+
+    @Test
+    fun `Account expiry that has expired should emit true`() = runTest {
+        // Arrange
+        val expiredAccountExpiry = AccountExpiry.Available(DateTime.now().minusDays(1))
+        // Act, Assert
+        outOfTimeUseCase.isOutOfTime().test {
+            assertEquals(null, awaitItem())
+            expiry.emit(expiredAccountExpiry)
+            assertEquals(true, awaitItem())
+        }
+    }
+
+    @Test
+    fun `Account expiry that has not expired should emit false`() = runTest {
+        // Arrange
+        val expiredAccountExpiry = AccountExpiry.Available(DateTime.now().plusDays(1))
+
+        // Act, Assert
+        outOfTimeUseCase.isOutOfTime().test {
+            assertEquals(null, awaitItem())
+            expiry.emit(expiredAccountExpiry)
+            assertEquals(false, awaitItem())
+        }
+    }
+
+    @Test
+    fun `Account that expires without new expiry event`() = runTest {
+        // Arrange
+        val expiredAccountExpiry = AccountExpiry.Available(DateTime.now().plusSeconds(62))
+
+        // Act, Assert
+        outOfTimeUseCase.isOutOfTime().test {
+            // Initial event
+            assertEquals(null, awaitItem())
+
+            expiry.emit(expiredAccountExpiry)
+            assertEquals(false, awaitItem())
+            assertEquals(true, awaitItem())
+        }
+    }
+}
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 c02e755951..282d1d3a27 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
@@ -11,10 +11,8 @@ import io.mockk.unmockkAll
 import io.mockk.verify
 import kotlin.test.assertEquals
 import kotlin.test.assertIs
-import kotlin.test.assertNull
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.test.runTest
-import net.mullvad.mullvadvpn.compose.dialog.payment.PaymentDialogData
 import net.mullvad.mullvadvpn.compose.state.PaymentState
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
 import net.mullvad.mullvadvpn.lib.common.test.assertLists
@@ -32,7 +30,6 @@ import net.mullvad.mullvadvpn.ui.serviceconnection.AuthTokenCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
 import net.mullvad.mullvadvpn.ui.serviceconnection.authTokenCache
 import net.mullvad.mullvadvpn.usecase.PaymentUseCase
-import net.mullvad.mullvadvpn.util.toPaymentDialogData
 import org.junit.After
 import org.junit.Before
 import org.junit.Rule
@@ -161,29 +158,6 @@ class AccountViewModelTest {
     }
 
     @Test
-    fun testBillingUserCancelled() = runTest {
-        // Arrange
-        val result = PurchaseResult.Completed.Cancelled
-        purchaseResult.value = result
-        every { result.toPaymentDialogData() } returns null
-
-        // Act, Assert
-        viewModel.uiState.test { assertNull(awaitItem().paymentDialogData) }
-    }
-
-    @Test
-    fun testBillingPurchaseSuccess() = runTest {
-        // Arrange
-        val result = PurchaseResult.Completed.Success
-        val expectedData: PaymentDialogData = mockk()
-        purchaseResult.value = result
-        every { result.toPaymentDialogData() } returns expectedData
-
-        // Act, Assert
-        viewModel.uiState.test { assertEquals(expectedData, awaitItem().paymentDialogData) }
-    }
-
-    @Test
     fun testStartBillingPayment() {
         // Arrange
         val mockProductId = ProductId("MOCK")
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ChangelogViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ChangelogViewModelTest.kt
index e223a12539..3350178ca3 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ChangelogViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ChangelogViewModelTest.kt
@@ -8,15 +8,17 @@ import io.mockk.impl.annotations.MockK
 import io.mockk.just
 import io.mockk.mockkStatic
 import io.mockk.unmockkAll
-import io.mockk.verify
-import kotlin.test.assertEquals
+import kotlin.test.assertNotNull
 import kotlinx.coroutines.test.runTest
+import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
 import net.mullvad.mullvadvpn.repository.ChangelogRepository
 import org.junit.After
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
 
 class ChangelogViewModelTest {
+    @get:Rule val testCoroutineRule = TestCoroutineRule()
 
     @MockK private lateinit var mockedChangelogRepository: ChangelogRepository
 
@@ -28,7 +30,6 @@ class ChangelogViewModelTest {
         mockkStatic(EVENT_NOTIFIER_EXTENSION_CLASS)
         every { mockedChangelogRepository.setVersionCodeOfMostRecentChangelogShowed(any()) } just
             Runs
-        viewModel = ChangelogViewModel(mockedChangelogRepository, 1, false)
     }
 
     @After
@@ -37,54 +38,41 @@ class ChangelogViewModelTest {
     }
 
     @Test
-    fun testInitialState() = runTest {
-        // Arrange, Act, Assert
-        viewModel.uiState.test { assertEquals(ChangelogDialogUiState.Hide, awaitItem()) }
+    fun testUpToDateVersionCodeShouldNotEmitChangelog() = runTest {
+        // Arrange
+        every { mockedChangelogRepository.getVersionCodeOfMostRecentChangelogShowed() } returns
+            buildVersionCode
+        viewModel = ChangelogViewModel(mockedChangelogRepository, buildVersionCode, false)
+
+        // If we have the most up to date version code, we should not show the changelog dialog
+        viewModel.uiSideEffect.test { expectNoEvents() }
     }
 
     @Test
-    fun testShowAndDismissChangelogDialog() = runTest {
-        viewModel.uiState.test {
-            // Arrange
-            val fakeList = listOf("test")
-            every { mockedChangelogRepository.getVersionCodeOfMostRecentChangelogShowed() } returns
-                -1
-            every { mockedChangelogRepository.getLastVersionChanges() } returns fakeList
-
-            // Assert initial ui state
-            assertEquals(ChangelogDialogUiState.Hide, awaitItem())
+    fun testNotUpToDateVersionCodeShouldEmitChangelog() = runTest {
+        // Arrange
+        every { mockedChangelogRepository.getVersionCodeOfMostRecentChangelogShowed() } returns -1
+        every { mockedChangelogRepository.getLastVersionChanges() } returns listOf("bla", "bla")
 
-            // Refresh and verify that the dialog should be shown
-            viewModel.refreshChangelogDialogUiState()
-            assertEquals(ChangelogDialogUiState.Show(fakeList), awaitItem())
-
-            // Dismiss dialog and verify that the dialog should be hidden
-            viewModel.dismissChangelogDialog()
-            assertEquals(ChangelogDialogUiState.Hide, awaitItem())
-            verify { mockedChangelogRepository.setVersionCodeOfMostRecentChangelogShowed(1) }
-        }
+        viewModel = ChangelogViewModel(mockedChangelogRepository, buildVersionCode, false)
+        // Given a new version with a change log we should return it
+        viewModel.uiSideEffect.test { assertNotNull(awaitItem()) }
     }
 
     @Test
-    fun testShowCaseChangelogWithEmptyListDialog() = runTest {
-        viewModel.uiState.test {
-            // Arrange
-            val fakeEmptyList = emptyList<String>()
-            every { mockedChangelogRepository.getVersionCodeOfMostRecentChangelogShowed() } returns
-                -1
-            every { mockedChangelogRepository.getLastVersionChanges() } returns fakeEmptyList
-
-            // Assert initial ui state
-            assertEquals(ChangelogDialogUiState.Hide, awaitItem())
+    fun testEmptyChangelogShouldNotEmitChangelog() = runTest {
+        // Arrange
+        every { mockedChangelogRepository.getVersionCodeOfMostRecentChangelogShowed() } returns -1
+        every { mockedChangelogRepository.getLastVersionChanges() } returns emptyList()
 
-            // Refresh and verify that the Ui state remain same due list being empty
-            viewModel.refreshChangelogDialogUiState()
-            expectNoEvents()
-        }
+        viewModel = ChangelogViewModel(mockedChangelogRepository, buildVersionCode, false)
+        // Given a new version with a change log we should not return it
+        viewModel.uiSideEffect.test { expectNoEvents() }
     }
 
     companion object {
         private const val EVENT_NOTIFIER_EXTENSION_CLASS =
             "net.mullvad.talpid.util.EventNotifierExtensionsKt"
+        private const val buildVersionCode = 10
     }
 }
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModelTest.kt
index 345a57df80..35898df4ab 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ConnectViewModelTest.kt
@@ -39,11 +39,11 @@ import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionState
 import net.mullvad.mullvadvpn.ui.serviceconnection.authTokenCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.connectionProxy
+import net.mullvad.mullvadvpn.usecase.OutOfTimeUseCase
 import net.mullvad.mullvadvpn.usecase.PaymentUseCase
 import net.mullvad.mullvadvpn.usecase.RelayListUseCase
 import net.mullvad.mullvadvpn.util.appVersionCallbackFlow
 import net.mullvad.talpid.tunnel.ErrorState
-import net.mullvad.talpid.tunnel.ErrorStateCause
 import net.mullvad.talpid.util.EventNotifier
 import org.junit.After
 import org.junit.Before
@@ -103,6 +103,10 @@ class ConnectViewModelTest {
     // Flows
     private val selectedRelayFlow = MutableStateFlow<RelayItem?>(null)
 
+    // Out Of Time Use Case
+    private val outOfTimeUseCase: OutOfTimeUseCase = mockk()
+    private val outOfTimeViewFlow = MutableStateFlow(false)
+
     @Before
     fun setup() {
         mockkStatic(CACHE_EXTENSION_CLASS)
@@ -136,6 +140,7 @@ class ConnectViewModelTest {
         // Flows
         every { mockRelayListUseCase.selectedRelayItem() } returns selectedRelayFlow
 
+        every { outOfTimeUseCase.isOutOfTime() } returns outOfTimeViewFlow
         viewModel =
             ConnectViewModel(
                 serviceConnectionManager = mockServiceConnectionManager,
@@ -144,6 +149,7 @@ class ConnectViewModelTest {
                 inAppNotificationController = mockInAppNotificationController,
                 relayListUseCase = mockRelayListUseCase,
                 newDeviceNotificationUseCase = mockk(),
+                outOfTimeUseCase = outOfTimeUseCase,
                 paymentUseCase = mockPaymentUseCase
             )
     }
@@ -342,8 +348,6 @@ class ConnectViewModelTest {
     fun testOutOfTimeUiSideEffect() =
         runTest(testCoroutineRule.testDispatcher) {
             // Arrange
-            val errorStateCause = ErrorStateCause.AuthFailed("[EXPIRED_ACCOUNT]")
-            val tunnelRealStateTestItem = TunnelState.Error(ErrorState(errorStateCause, true))
             val deferred = async { viewModel.uiSideEffect.first() }
 
             // Act
@@ -352,12 +356,12 @@ class ConnectViewModelTest {
                 serviceConnectionState.value =
                     ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
                 locationSlot.captured.invoke(mockLocation)
-                eventNotifierTunnelRealState.notify(tunnelRealStateTestItem)
+                outOfTimeViewFlow.value = true
                 awaitItem()
             }
 
             // Assert
-            assertIs<ConnectViewModel.UiSideEffect.OpenOutOfTimeView>(deferred.await())
+            assertIs<ConnectViewModel.UiSideEffect.OutOfTime>(deferred.await())
         }
 
     companion object {
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModelTest.kt
index 7eb35404d0..c402a3103e 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModelTest.kt
@@ -20,6 +20,7 @@ import net.mullvad.mullvadvpn.compose.state.LoginState.Success
 import net.mullvad.mullvadvpn.compose.state.LoginUiState
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
 import net.mullvad.mullvadvpn.model.AccountCreationResult
+import net.mullvad.mullvadvpn.model.AccountExpiry
 import net.mullvad.mullvadvpn.model.AccountHistory
 import net.mullvad.mullvadvpn.model.AccountToken
 import net.mullvad.mullvadvpn.model.DeviceListEvent
@@ -28,6 +29,7 @@ import net.mullvad.mullvadvpn.repository.AccountRepository
 import net.mullvad.mullvadvpn.repository.DeviceRepository
 import net.mullvad.mullvadvpn.usecase.ConnectivityUseCase
 import net.mullvad.mullvadvpn.usecase.NewDeviceNotificationUseCase
+import org.joda.time.DateTime
 import org.junit.Assert.assertEquals
 import org.junit.Before
 import org.junit.Rule
@@ -113,6 +115,8 @@ class LoginViewModelTest {
             val uiStates = loginViewModel.uiState.testIn(backgroundScope)
             val sideEffects = loginViewModel.uiSideEffect.testIn(backgroundScope)
             coEvery { mockedAccountRepository.login(any()) } returns LoginResult.Ok
+            coEvery { mockedAccountRepository.accountExpiryState } returns
+                MutableStateFlow(AccountExpiry.Available(DateTime.now().plusDays(3)))
 
             // Act, Assert
             uiStates.skipDefaultItem()
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/OutOfTimeViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/OutOfTimeViewModelTest.kt
index dad51eab59..0232f12e89 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/OutOfTimeViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/OutOfTimeViewModelTest.kt
@@ -1,6 +1,5 @@
 package net.mullvad.mullvadvpn.viewmodel
 
-import android.app.Activity
 import androidx.lifecycle.viewModelScope
 import app.cash.turbine.test
 import io.mockk.coEvery
@@ -12,18 +11,15 @@ import io.mockk.unmockkAll
 import io.mockk.verify
 import kotlin.test.assertEquals
 import kotlin.test.assertIs
-import kotlin.test.assertNull
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.test.runTest
-import net.mullvad.mullvadvpn.compose.dialog.payment.PaymentDialogData
 import net.mullvad.mullvadvpn.compose.state.OutOfTimeUiState
 import net.mullvad.mullvadvpn.compose.state.PaymentState
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
 import net.mullvad.mullvadvpn.lib.common.test.assertLists
 import net.mullvad.mullvadvpn.lib.payment.model.PaymentAvailability
 import net.mullvad.mullvadvpn.lib.payment.model.PaymentProduct
-import net.mullvad.mullvadvpn.lib.payment.model.ProductId
 import net.mullvad.mullvadvpn.lib.payment.model.PurchaseResult
 import net.mullvad.mullvadvpn.model.AccountExpiry
 import net.mullvad.mullvadvpn.model.DeviceState
@@ -37,8 +33,8 @@ import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionState
 import net.mullvad.mullvadvpn.ui.serviceconnection.authTokenCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.connectionProxy
+import net.mullvad.mullvadvpn.usecase.OutOfTimeUseCase
 import net.mullvad.mullvadvpn.usecase.PaymentUseCase
-import net.mullvad.mullvadvpn.util.toPaymentDialogData
 import net.mullvad.talpid.util.EventNotifier
 import org.joda.time.DateTime
 import org.joda.time.ReadableInstant
@@ -50,12 +46,13 @@ import org.junit.Test
 class OutOfTimeViewModelTest {
     @get:Rule val testCoroutineRule = TestCoroutineRule()
 
-    private val serviceConnectionState =
+    private val serviceConnectionStateFlow =
         MutableStateFlow<ServiceConnectionState>(ServiceConnectionState.Disconnected)
-    private val accountExpiryState = MutableStateFlow<AccountExpiry>(AccountExpiry.Missing)
-    private val deviceState = MutableStateFlow<DeviceState>(DeviceState.Initial)
-    private val paymentAvailability = MutableStateFlow<PaymentAvailability?>(null)
-    private val purchaseResult = MutableStateFlow<PurchaseResult?>(null)
+    private val accountExpiryStateFlow = MutableStateFlow<AccountExpiry>(AccountExpiry.Missing)
+    private val deviceStateFlow = MutableStateFlow<DeviceState>(DeviceState.Initial)
+    private val paymentAvailabilityFlow = MutableStateFlow<PaymentAvailability?>(null)
+    private val purchaseResultFlow = MutableStateFlow<PurchaseResult?>(null)
+    private val outOfTimeFlow = MutableStateFlow(true)
 
     // Service connections
     private val mockServiceConnectionContainer: ServiceConnectionContainer = mockk()
@@ -68,6 +65,7 @@ class OutOfTimeViewModelTest {
     private val mockDeviceRepository: DeviceRepository = mockk()
     private val mockServiceConnectionManager: ServiceConnectionManager = mockk()
     private val mockPaymentUseCase: PaymentUseCase = mockk(relaxed = true)
+    private val mockOutOfTimeUseCase: OutOfTimeUseCase = mockk(relaxed = true)
 
     private lateinit var viewModel: OutOfTimeViewModel
 
@@ -76,19 +74,21 @@ class OutOfTimeViewModelTest {
         mockkStatic(SERVICE_CONNECTION_MANAGER_EXTENSIONS)
         mockkStatic(PURCHASE_RESULT_EXTENSIONS_CLASS)
 
-        every { mockServiceConnectionManager.connectionState } returns serviceConnectionState
+        every { mockServiceConnectionManager.connectionState } returns serviceConnectionStateFlow
 
         every { mockServiceConnectionContainer.connectionProxy } returns mockConnectionProxy
 
         every { mockConnectionProxy.onStateChange } returns eventNotifierTunnelRealState
 
-        every { mockAccountRepository.accountExpiryState } returns accountExpiryState
+        every { mockAccountRepository.accountExpiryState } returns accountExpiryStateFlow
 
-        every { mockDeviceRepository.deviceState } returns deviceState
+        every { mockDeviceRepository.deviceState } returns deviceStateFlow
 
-        coEvery { mockPaymentUseCase.purchaseResult } returns purchaseResult
+        coEvery { mockPaymentUseCase.purchaseResult } returns purchaseResultFlow
 
-        coEvery { mockPaymentUseCase.paymentAvailability } returns paymentAvailability
+        coEvery { mockPaymentUseCase.paymentAvailability } returns paymentAvailabilityFlow
+
+        coEvery { mockOutOfTimeUseCase.isOutOfTime() } returns outOfTimeFlow
 
         viewModel =
             OutOfTimeViewModel(
@@ -96,6 +96,7 @@ class OutOfTimeViewModelTest {
                 serviceConnectionManager = mockServiceConnectionManager,
                 deviceRepository = mockDeviceRepository,
                 paymentUseCase = mockPaymentUseCase,
+                outOfTimeUseCase = mockOutOfTimeUseCase,
                 pollAccountExpiry = false
             )
     }
@@ -134,7 +135,7 @@ class OutOfTimeViewModelTest {
             viewModel.uiState.test {
                 assertEquals(OutOfTimeUiState(deviceName = ""), awaitItem())
                 eventNotifierTunnelRealState.notify(tunnelRealStateTestItem)
-                serviceConnectionState.value =
+                serviceConnectionStateFlow.value =
                     ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
                 val result = awaitItem()
                 assertEquals(tunnelRealStateTestItem, result.tunnelState)
@@ -150,7 +151,7 @@ class OutOfTimeViewModelTest {
 
             // Act, Assert
             viewModel.uiSideEffect.test {
-                accountExpiryState.value = AccountExpiry.Available(mockExpiryDate)
+                outOfTimeFlow.value = false
                 val action = awaitItem()
                 assertIs<OutOfTimeViewModel.UiSideEffect.OpenConnectScreen>(action)
             }
@@ -174,8 +175,8 @@ class OutOfTimeViewModelTest {
     fun testBillingProductsUnavailableState() = runTest {
         // Arrange
         val productsUnavailable = PaymentAvailability.ProductsUnavailable
-        paymentAvailability.value = productsUnavailable
-        serviceConnectionState.value =
+        paymentAvailabilityFlow.value = productsUnavailable
+        serviceConnectionStateFlow.value =
             ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
 
         // Act, Assert
@@ -189,8 +190,8 @@ class OutOfTimeViewModelTest {
     fun testBillingProductsGenericErrorState() = runTest {
         // Arrange
         val paymentAvailabilityError = PaymentAvailability.Error.Other(mockk())
-        paymentAvailability.value = paymentAvailabilityError
-        serviceConnectionState.value =
+        paymentAvailabilityFlow.value = paymentAvailabilityError
+        serviceConnectionStateFlow.value =
             ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
 
         // Act, Assert
@@ -204,8 +205,8 @@ class OutOfTimeViewModelTest {
     fun testBillingProductsBillingErrorState() = runTest {
         // Arrange
         val paymentAvailabilityError = PaymentAvailability.Error.BillingUnavailable
-        paymentAvailability.value = paymentAvailabilityError
-        serviceConnectionState.value =
+        paymentAvailabilityFlow.value = paymentAvailabilityError
+        serviceConnectionStateFlow.value =
             ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
 
         // Act, Assert
@@ -221,8 +222,8 @@ class OutOfTimeViewModelTest {
         val mockProduct: PaymentProduct = mockk()
         val expectedProductList = listOf(mockProduct)
         val productsAvailable = PaymentAvailability.ProductsAvailable(listOf(mockProduct))
-        paymentAvailability.value = productsAvailable
-        serviceConnectionState.value =
+        paymentAvailabilityFlow.value = productsAvailable
+        serviceConnectionStateFlow.value =
             ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
 
         // Act, Assert
@@ -234,44 +235,6 @@ class OutOfTimeViewModelTest {
     }
 
     @Test
-    fun testBillingUserCancelled() = runTest {
-        // Arrange
-        val result = PurchaseResult.Completed.Cancelled
-        purchaseResult.value = result
-        every { result.toPaymentDialogData() } returns null
-
-        // Act, Assert
-        viewModel.uiState.test { assertNull(awaitItem().paymentDialogData) }
-    }
-
-    @Test
-    fun testBillingPurchaseSuccess() = runTest {
-        // Arrange
-        val result = PurchaseResult.Completed.Success
-        val expectedData: PaymentDialogData = mockk()
-        purchaseResult.value = result
-        serviceConnectionState.value =
-            ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
-        every { result.toPaymentDialogData() } returns expectedData
-
-        // Act, Assert
-        viewModel.uiState.test { assertEquals(expectedData, awaitItem().paymentDialogData) }
-    }
-
-    @Test
-    fun testStartBillingPayment() {
-        // Arrange
-        val mockProductId = ProductId("MOCK")
-        val mockActivityProvider = mockk<() -> Activity>()
-
-        // Act
-        viewModel.startBillingPayment(mockProductId, mockActivityProvider)
-
-        // Assert
-        coVerify { mockPaymentUseCase.purchaseProduct(mockProductId, mockActivityProvider) }
-    }
-
-    @Test
     fun testOnClosePurchaseResultDialogSuccessful() {
         // Arrange
 
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/PaymentViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/PaymentViewModelTest.kt
new file mode 100644
index 0000000000..665e23c3d4
--- /dev/null
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/PaymentViewModelTest.kt
@@ -0,0 +1,70 @@
+package net.mullvad.mullvadvpn.viewmodel
+
+import app.cash.turbine.test
+import io.mockk.coEvery
+import io.mockk.mockk
+import io.mockk.unmockkAll
+import kotlin.test.assertEquals
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.test.runTest
+import net.mullvad.mullvadvpn.compose.dialog.payment.PaymentDialogData
+import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
+import net.mullvad.mullvadvpn.lib.payment.model.PurchaseResult
+import net.mullvad.mullvadvpn.usecase.PaymentUseCase
+import net.mullvad.mullvadvpn.util.toPaymentDialogData
+import org.junit.After
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+
+class PaymentViewModelTest {
+    @get:Rule val testCoroutineRule = TestCoroutineRule()
+
+    private val mockPaymentUseCase: PaymentUseCase = mockk(relaxed = true)
+
+    private val purchaseResult = MutableStateFlow<PurchaseResult?>(null)
+
+    private lateinit var viewModel: PaymentViewModel
+
+    @Before
+    fun setUp() {
+        coEvery { mockPaymentUseCase.purchaseResult } returns purchaseResult
+
+        viewModel = PaymentViewModel(paymentUseCase = mockPaymentUseCase)
+    }
+
+    @After
+    fun tearDown() {
+        unmockkAll()
+    }
+
+    @Test
+    fun testBillingUserCancelled() = runTest {
+        // Arrange
+        val result = PurchaseResult.Completed.Cancelled
+        purchaseResult.value = result
+
+        // Act, Assert
+        viewModel.uiState.test {
+            assertEquals(PaymentDialogData(), awaitItem().paymentDialogData)
+            purchaseResult.value = result
+        }
+
+        viewModel.uiSideEffect.test {
+            assertEquals(PaymentUiSideEffect.PaymentCancelled, awaitItem())
+        }
+    }
+
+    @Test
+    fun testBillingPurchaseSuccess() = runTest {
+        // Arrange
+        val result = PurchaseResult.Completed.Success
+
+        // Act, Assert
+        viewModel.uiState.test {
+            awaitItem()
+            purchaseResult.value = result
+            assertEquals(result.toPaymentDialogData(), awaitItem().paymentDialogData)
+        }
+    }
+}
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ReportProblemModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ReportProblemViewModelTest.kt
index e0b1c5274c..5726c6249c 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ReportProblemModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/ReportProblemViewModelTest.kt
@@ -20,10 +20,11 @@ import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
 
-class ReportProblemModelTest {
+class ReportProblemViewModelTest {
     @get:Rule val testCoroutineRule = TestCoroutineRule()
 
     @MockK private lateinit var mockMullvadProblemReport: MullvadProblemReport
+
     @MockK(relaxed = true) private lateinit var mockProblemReportRepository: ProblemReportRepository
 
     private val problemReportFlow = MutableStateFlow(UserReport("", ""))
@@ -87,16 +88,30 @@ class ReportProblemModelTest {
         coEvery { mockMullvadProblemReport.sendReport(any()) } returns
             SendProblemReportResult.Success
         val email = ""
+        val description = "My description"
+
+        coEvery { mockProblemReportRepository.setDescription(any()) } answers
+            {
+                problemReportFlow.value = problemReportFlow.value.copy(description = arg(0))
+            }
 
         // Act, Assert
         viewModel.uiState.test {
-            assertEquals(ReportProblemUiState(false, null), awaitItem())
-            viewModel.sendReport(email, "My description")
-            assertEquals(ReportProblemUiState(true, null), awaitItem())
-            viewModel.sendReport(email, "My description")
-            assertEquals(ReportProblemUiState(false, SendingReportUiState.Sending), awaitItem())
+            assertEquals(ReportProblemUiState(), awaitItem())
+            viewModel.updateDescription(description)
+            assertEquals(ReportProblemUiState(description = description), awaitItem())
+
+            viewModel.sendReport(email, description, true)
+            assertEquals(
+                ReportProblemUiState(SendingReportUiState.Sending, email, description),
+                awaitItem()
+            )
             assertEquals(
-                ReportProblemUiState(false, SendingReportUiState.Success(null)),
+                ReportProblemUiState(
+                    SendingReportUiState.Success(null),
+                    "",
+                    "",
+                ),
                 awaitItem()
             )
         }
@@ -109,16 +124,44 @@ class ReportProblemModelTest {
         coEvery { mockMullvadProblemReport.sendReport(any()) } returns
             SendProblemReportResult.Success
         val email = "my@email.com"
+        val description = "My description"
+
+        // This might look a bit weird, and is not optimal. An alternative would be to use the real
+        // ProblemReportRepository, but that would complicate the other tests. This is a compromise.
+        coEvery { mockProblemReportRepository.setEmail(any()) } answers
+            {
+                problemReportFlow.value = problemReportFlow.value.copy(email = arg(0))
+            }
+        coEvery { mockProblemReportRepository.setDescription(any()) } answers
+            {
+                problemReportFlow.value = problemReportFlow.value.copy(description = arg(0))
+            }
 
         // Act, Assert
         viewModel.uiState.test {
-            assertEquals(awaitItem(), ReportProblemUiState(false, null))
-            viewModel.sendReport(email, "My description")
+            assertEquals(awaitItem(), ReportProblemUiState(null, "", ""))
+            viewModel.updateEmail(email)
+            awaitItem()
+            viewModel.updateDescription(description)
+            awaitItem()
+
+            viewModel.sendReport(email, description)
 
-            assertEquals(awaitItem(), ReportProblemUiState(false, SendingReportUiState.Sending))
             assertEquals(
-                awaitItem(),
-                ReportProblemUiState(false, SendingReportUiState.Success(email))
+                ReportProblemUiState(
+                    SendingReportUiState.Sending,
+                    email,
+                    description,
+                ),
+                awaitItem()
+            )
+            assertEquals(
+                ReportProblemUiState(
+                    SendingReportUiState.Success(email),
+                    "",
+                    "",
+                ),
+                awaitItem()
             )
         }
     }
@@ -146,19 +189,4 @@ class ReportProblemModelTest {
         // Assert
         verify { mockProblemReportRepository.setDescription(description) }
     }
-
-    @Test
-    fun testUpdateProblemReport() = runTest {
-        // Arrange
-        val userReport = UserReport("my@email.com", "My description")
-
-        // Act, Assert
-        viewModel.uiState.test {
-            awaitItem()
-            problemReportFlow.value = userReport
-            val result = awaitItem()
-            assertEquals(userReport.email, result.email)
-            assertEquals(userReport.description, result.description)
-        }
-    }
 }
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModelTest.kt
index 74d7d80c19..5ad1af1182 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModelTest.kt
@@ -120,10 +120,10 @@ class SelectLocationViewModelTest {
         every { mockRelayListUseCase.updateSelectedRelayLocation(mockLocation) } returns Unit
 
         // Act, Assert
-        viewModel.uiCloseAction.test {
+        viewModel.uiSideEffect.test {
             viewModel.selectRelay(mockRelayItem)
             // Await an empty item
-            assertEquals(Unit, awaitItem())
+            assertEquals(SelectLocationSideEffect.CloseScreen, awaitItem())
             verify {
                 connectionProxyMock.connect()
                 mockRelayListUseCase.updateSelectedRelayLocation(mockLocation)
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/VpnSettingsViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/VpnSettingsViewModelTest.kt
index f8736eb823..0ac13777cd 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/VpnSettingsViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/VpnSettingsViewModelTest.kt
@@ -9,15 +9,11 @@ import io.mockk.unmockkAll
 import io.mockk.verify
 import kotlin.test.assertEquals
 import kotlin.test.assertIs
-import kotlin.test.assertTrue
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.runTest
-import net.mullvad.mullvadvpn.compose.state.VpnSettingsDialog
-import net.mullvad.mullvadvpn.compose.state.VpnSettingsUiState
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
-import net.mullvad.mullvadvpn.lib.common.test.assertLists
 import net.mullvad.mullvadvpn.model.Constraint
 import net.mullvad.mullvadvpn.model.Port
 import net.mullvad.mullvadvpn.model.PortRange
@@ -31,7 +27,6 @@ import net.mullvad.mullvadvpn.model.WireguardTunnelOptions
 import net.mullvad.mullvadvpn.repository.SettingsRepository
 import net.mullvad.mullvadvpn.usecase.PortRangeUseCase
 import net.mullvad.mullvadvpn.usecase.RelayListUseCase
-import org.apache.commons.validator.routines.InetAddressValidator
 import org.junit.After
 import org.junit.Before
 import org.junit.Rule
@@ -41,7 +36,6 @@ class VpnSettingsViewModelTest {
     @get:Rule val testCoroutineRule = TestCoroutineRule()
 
     private val mockSettingsRepository: SettingsRepository = mockk()
-    private val mockInetAddressValidator: InetAddressValidator = mockk()
     private val mockResources: Resources = mockk()
     private val mockPortRangeUseCase: PortRangeUseCase = mockk()
     private val mockRelayListUseCase: RelayListUseCase = mockk()
@@ -59,7 +53,6 @@ class VpnSettingsViewModelTest {
         viewModel =
             VpnSettingsViewModel(
                 repository = mockSettingsRepository,
-                inetAddressValidator = mockInetAddressValidator,
                 resources = mockResources,
                 portRangeUseCase = mockPortRangeUseCase,
                 relayListUseCase = mockRelayListUseCase,
@@ -133,6 +126,7 @@ class VpnSettingsViewModelTest {
         viewModel.uiState.test {
             assertIs<Constraint.Any<Port>>(awaitItem().selectedWireguardPort)
             mockSettingsUpdate.value = mockSettings
+            assertEquals(expectedPort, awaitItem().customWireguardPort)
             assertEquals(expectedPort, awaitItem().selectedWireguardPort)
         }
     }
@@ -152,23 +146,4 @@ class VpnSettingsViewModelTest {
             mockRelayListUseCase.updateSelectedWireguardConstraints(wireguardConstraints)
         }
     }
-
-    @Test
-    fun test_update_port_range_state() = runTest {
-        // Arrange
-        val expectedPortRange = listOf<PortRange>(mockk(), mockk())
-        val mockSettings: Settings = mockk(relaxed = true)
-
-        every { mockSettings.relaySettings } returns mockk<RelaySettings.Normal>(relaxed = true)
-        portRangeFlow.value = expectedPortRange
-
-        // Act, Assert
-        viewModel.uiState.test {
-            assertIs<VpnSettingsUiState>(awaitItem())
-            viewModel.onWireguardPortInfoClicked()
-            val state = awaitItem()
-            assertTrue { state.dialog is VpnSettingsDialog.WireguardPortInfo }
-            assertLists(expectedPortRange, state.availablePortRanges)
-        }
-    }
 }
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/WelcomeViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/WelcomeViewModelTest.kt
index e958df9337..433a9f5709 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/WelcomeViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/WelcomeViewModelTest.kt
@@ -1,28 +1,23 @@
 package net.mullvad.mullvadvpn.viewmodel
 
-import android.app.Activity
 import androidx.lifecycle.viewModelScope
 import app.cash.turbine.test
 import io.mockk.coEvery
-import io.mockk.coVerify
 import io.mockk.every
 import io.mockk.mockk
 import io.mockk.mockkStatic
 import io.mockk.unmockkAll
 import kotlin.test.assertEquals
 import kotlin.test.assertIs
-import kotlin.test.assertNull
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.test.runTest
-import net.mullvad.mullvadvpn.compose.dialog.payment.PaymentDialogData
 import net.mullvad.mullvadvpn.compose.state.PaymentState
 import net.mullvad.mullvadvpn.compose.state.WelcomeUiState
 import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
 import net.mullvad.mullvadvpn.lib.common.test.assertLists
 import net.mullvad.mullvadvpn.lib.payment.model.PaymentAvailability
 import net.mullvad.mullvadvpn.lib.payment.model.PaymentProduct
-import net.mullvad.mullvadvpn.lib.payment.model.ProductId
 import net.mullvad.mullvadvpn.lib.payment.model.PurchaseResult
 import net.mullvad.mullvadvpn.model.AccountAndDevice
 import net.mullvad.mullvadvpn.model.AccountExpiry
@@ -37,8 +32,8 @@ import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionContainer
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionState
 import net.mullvad.mullvadvpn.ui.serviceconnection.authTokenCache
+import net.mullvad.mullvadvpn.usecase.OutOfTimeUseCase
 import net.mullvad.mullvadvpn.usecase.PaymentUseCase
-import net.mullvad.mullvadvpn.util.toPaymentDialogData
 import net.mullvad.talpid.util.EventNotifier
 import org.joda.time.DateTime
 import org.joda.time.ReadableInstant
@@ -50,12 +45,13 @@ import org.junit.Test
 class WelcomeViewModelTest {
     @get:Rule val testCoroutineRule = TestCoroutineRule()
 
-    private val serviceConnectionState =
+    private val serviceConnectionStateFlow =
         MutableStateFlow<ServiceConnectionState>(ServiceConnectionState.Disconnected)
-    private val deviceState = MutableStateFlow<DeviceState>(DeviceState.Initial)
-    private val accountExpiryState = MutableStateFlow<AccountExpiry>(AccountExpiry.Missing)
-    private val purchaseResult = MutableStateFlow<PurchaseResult?>(null)
-    private val paymentAvailability = MutableStateFlow<PaymentAvailability?>(null)
+    private val deviceStateFlow = MutableStateFlow<DeviceState>(DeviceState.Initial)
+    private val accountExpiryStateFlow = MutableStateFlow<AccountExpiry>(AccountExpiry.Missing)
+    private val purchaseResultFlow = MutableStateFlow<PurchaseResult?>(null)
+    private val paymentAvailabilityFlow = MutableStateFlow<PaymentAvailability?>(null)
+    private val outOfTimeFlow = MutableStateFlow(true)
 
     // Service connections
     private val mockServiceConnectionContainer: ServiceConnectionContainer = mockk()
@@ -68,6 +64,7 @@ class WelcomeViewModelTest {
     private val mockDeviceRepository: DeviceRepository = mockk()
     private val mockServiceConnectionManager: ServiceConnectionManager = mockk()
     private val mockPaymentUseCase: PaymentUseCase = mockk(relaxed = true)
+    private val mockOutOfTimeUseCase: OutOfTimeUseCase = mockk(relaxed = true)
 
     private lateinit var viewModel: WelcomeViewModel
 
@@ -76,19 +73,21 @@ class WelcomeViewModelTest {
         mockkStatic(SERVICE_CONNECTION_MANAGER_EXTENSIONS)
         mockkStatic(PURCHASE_RESULT_EXTENSIONS_CLASS)
 
-        every { mockDeviceRepository.deviceState } returns deviceState
+        every { mockDeviceRepository.deviceState } returns deviceStateFlow
 
-        every { mockServiceConnectionManager.connectionState } returns serviceConnectionState
+        every { mockServiceConnectionManager.connectionState } returns serviceConnectionStateFlow
 
         every { mockServiceConnectionContainer.connectionProxy } returns mockConnectionProxy
 
         every { mockConnectionProxy.onUiStateChange } returns eventNotifierTunnelUiState
 
-        every { mockAccountRepository.accountExpiryState } returns accountExpiryState
+        every { mockAccountRepository.accountExpiryState } returns accountExpiryStateFlow
 
-        coEvery { mockPaymentUseCase.purchaseResult } returns purchaseResult
+        coEvery { mockPaymentUseCase.purchaseResult } returns purchaseResultFlow
 
-        coEvery { mockPaymentUseCase.paymentAvailability } returns paymentAvailability
+        coEvery { mockPaymentUseCase.paymentAvailability } returns paymentAvailabilityFlow
+
+        coEvery { mockOutOfTimeUseCase.isOutOfTime() } returns outOfTimeFlow
 
         viewModel =
             WelcomeViewModel(
@@ -96,6 +95,7 @@ class WelcomeViewModelTest {
                 deviceRepository = mockDeviceRepository,
                 serviceConnectionManager = mockServiceConnectionManager,
                 paymentUseCase = mockPaymentUseCase,
+                outOfTimeUseCase = mockOutOfTimeUseCase,
                 pollAccountExpiry = false
             )
     }
@@ -134,7 +134,7 @@ class WelcomeViewModelTest {
             viewModel.uiState.test {
                 assertEquals(WelcomeUiState(), awaitItem())
                 eventNotifierTunnelUiState.notify(tunnelUiStateTestItem)
-                serviceConnectionState.value =
+                serviceConnectionStateFlow.value =
                     ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
                 val result = awaitItem()
                 assertEquals(tunnelUiStateTestItem, result.tunnelState)
@@ -142,27 +142,26 @@ class WelcomeViewModelTest {
         }
 
     @Test
-    fun testUpdateAccountNumber() =
-        runTest(testCoroutineRule.testDispatcher) {
-            // Arrange
-            val expectedAccountNumber = "4444555566667777"
-            val device: Device = mockk()
-            every { device.displayName() } returns ""
+    fun testUpdateAccountNumber() = runTest {
+        // Arrange
+        val expectedAccountNumber = "4444555566667777"
+        val device: Device = mockk()
+        every { device.displayName() } returns ""
 
-            // Act, Assert
-            viewModel.uiState.test {
-                assertEquals(WelcomeUiState(), awaitItem())
-                serviceConnectionState.value =
-                    ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
-                deviceState.value =
-                    DeviceState.LoggedIn(
-                        accountAndDevice =
-                            AccountAndDevice(account_token = expectedAccountNumber, device = device)
-                    )
-                val result = awaitItem()
-                assertEquals(expectedAccountNumber, result.accountNumber)
-            }
+        // Act, Assert
+        viewModel.uiState.test {
+            assertEquals(WelcomeUiState(), awaitItem())
+            paymentAvailabilityFlow.value = null
+            deviceStateFlow.value =
+                DeviceState.LoggedIn(
+                    accountAndDevice =
+                        AccountAndDevice(account_token = expectedAccountNumber, device = device)
+                )
+            serviceConnectionStateFlow.value =
+                ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
+            assertEquals(expectedAccountNumber, awaitItem().accountNumber)
         }
+    }
 
     @Test
     fun testOpenConnectScreen() =
@@ -173,7 +172,7 @@ class WelcomeViewModelTest {
 
             // Act, Assert
             viewModel.uiSideEffect.test {
-                accountExpiryState.value = AccountExpiry.Available(mockExpiryDate)
+                outOfTimeFlow.value = false
                 val action = awaitItem()
                 assertIs<WelcomeViewModel.UiSideEffect.OpenConnectScreen>(action)
             }
@@ -188,8 +187,8 @@ class WelcomeViewModelTest {
         viewModel.uiState.test {
             // Default item
             awaitItem()
-            paymentAvailability.tryEmit(productsUnavailable)
-            serviceConnectionState.value =
+            paymentAvailabilityFlow.tryEmit(productsUnavailable)
+            serviceConnectionStateFlow.value =
                 ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
             val result = awaitItem().billingPaymentState
             assertIs<PaymentState.NoPayment>(result)
@@ -200,8 +199,8 @@ class WelcomeViewModelTest {
     fun testBillingProductsGenericErrorState() = runTest {
         // Arrange
         val paymentOtherError = PaymentAvailability.Error.Other(mockk())
-        paymentAvailability.tryEmit(paymentOtherError)
-        serviceConnectionState.value =
+        paymentAvailabilityFlow.tryEmit(paymentOtherError)
+        serviceConnectionStateFlow.value =
             ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
 
         // Act, Assert
@@ -215,8 +214,8 @@ class WelcomeViewModelTest {
     fun testBillingProductsBillingErrorState() = runTest {
         // Arrange
         val paymentBillingError = PaymentAvailability.Error.BillingUnavailable
-        paymentAvailability.value = paymentBillingError
-        serviceConnectionState.value =
+        paymentAvailabilityFlow.value = paymentBillingError
+        serviceConnectionStateFlow.value =
             ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
 
         // Act, Assert
@@ -232,8 +231,8 @@ class WelcomeViewModelTest {
         val mockProduct: PaymentProduct = mockk()
         val expectedProductList = listOf(mockProduct)
         val productsAvailable = PaymentAvailability.ProductsAvailable(listOf(mockProduct))
-        paymentAvailability.value = productsAvailable
-        serviceConnectionState.value =
+        paymentAvailabilityFlow.value = productsAvailable
+        serviceConnectionStateFlow.value =
             ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
 
         // Act, Assert
@@ -244,46 +243,6 @@ class WelcomeViewModelTest {
         }
     }
 
-    @Test
-    fun testBillingUserCancelled() = runTest {
-        // Arrange
-        val result = PurchaseResult.Completed.Cancelled
-        purchaseResult.value = result
-        serviceConnectionState.value =
-            ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
-        every { result.toPaymentDialogData() } returns null
-
-        // Act, Assert
-        viewModel.uiState.test { assertNull(awaitItem().paymentDialogData) }
-    }
-
-    @Test
-    fun testBillingPurchaseSuccess() = runTest {
-        // Arrange
-        val result = PurchaseResult.Completed.Success
-        val expectedData: PaymentDialogData = mockk()
-        purchaseResult.value = result
-        serviceConnectionState.value =
-            ServiceConnectionState.ConnectedReady(mockServiceConnectionContainer)
-        every { result.toPaymentDialogData() } returns expectedData
-
-        // Act, Assert
-        viewModel.uiState.test { assertEquals(expectedData, awaitItem().paymentDialogData) }
-    }
-
-    @Test
-    fun testStartBillingPayment() {
-        // Arrange
-        val mockProductId = ProductId("MOCK")
-        val mockActivityProvider = mockk<() -> Activity>()
-
-        // Act
-        viewModel.startBillingPayment(mockProductId, mockActivityProvider)
-
-        // Assert
-        coVerify { mockPaymentUseCase.purchaseProduct(mockProductId, mockActivityProvider) }
-    }
-
     companion object {
         private const val SERVICE_CONNECTION_MANAGER_EXTENSIONS =
             "net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManagerExtensionsKt"