Refactor android app account cache

16 files changed, 214 insertions, 168 deletions

diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
index 4f29198652..47054e660a 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
@@ -6,6 +6,7 @@ import kotlinx.coroutines.Dispatchers
 import net.mullvad.mullvadvpn.applist.ApplicationsIconManager
 import net.mullvad.mullvadvpn.applist.ApplicationsProvider
 import net.mullvad.mullvadvpn.ipc.EventDispatcher
+import net.mullvad.mullvadvpn.ui.serviceconnection.AccountCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.DeviceRepository
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
 import net.mullvad.mullvadvpn.ui.serviceconnection.SplitTunneling
@@ -37,9 +38,11 @@ val uiModule = module {
     }
 
     single { ServiceConnectionManager(androidContext()) }
+
+    single { AccountCache(get()) }
     single { DeviceRepository(get()) }
     viewModel { LoginViewModel(get(), get()) }
-    viewModel { DeviceRevokedViewModel(get()) }
+    viewModel { DeviceRevokedViewModel(get(), get()) }
     viewModel { DeviceListViewModel(get()) }
 }
 
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/AccountFragment.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/AccountFragment.kt
index 8cd506f13d..3251c6df2a 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/AccountFragment.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/AccountFragment.kt
@@ -13,6 +13,7 @@ import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.R
 import net.mullvad.mullvadvpn.model.TunnelState
+import net.mullvad.mullvadvpn.ui.serviceconnection.AccountCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.DeviceRepository
 import net.mullvad.mullvadvpn.ui.widget.Button
 import net.mullvad.mullvadvpn.ui.widget.CopyableInformationView
@@ -25,6 +26,9 @@ import org.joda.time.DateTime
 import org.koin.android.ext.android.inject
 
 class AccountFragment : ServiceDependentFragment(OnNoService.GoBack) {
+
+    // Injected dependencies
+    private val accountCache: AccountCache by inject()
     private val deviceRepository: DeviceRepository by inject()
 
     override val isSecureScreen = true
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/ConnectFragment.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/ConnectFragment.kt
index 43d23445ce..dccba31b2e 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/ConnectFragment.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/ConnectFragment.kt
@@ -13,15 +13,21 @@ import net.mullvad.mullvadvpn.model.TunnelState
 import net.mullvad.mullvadvpn.ui.notification.AccountExpiryNotification
 import net.mullvad.mullvadvpn.ui.notification.TunnelStateNotification
 import net.mullvad.mullvadvpn.ui.notification.VersionInfoNotification
+import net.mullvad.mullvadvpn.ui.serviceconnection.AccountCache
 import net.mullvad.mullvadvpn.ui.widget.HeaderBar
 import net.mullvad.mullvadvpn.ui.widget.NotificationBanner
 import net.mullvad.mullvadvpn.ui.widget.SwitchLocationButton
 import org.joda.time.DateTime
+import org.koin.android.ext.android.inject
 
 val KEY_IS_TUNNEL_INFO_EXPANDED = "is_tunnel_info_expanded"
 
 class ConnectFragment :
     ServiceDependentFragment(OnNoService.GoToLaunchScreen), NavigationBarPainter {
+
+    // Injected dependencies
+    private val accountCache: AccountCache by inject()
+
     private lateinit var actionButton: ConnectActionButton
     private lateinit var switchLocationButton: SwitchLocationButton
     private lateinit var headerBar: HeaderBar
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/OutOfTimeFragment.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/OutOfTimeFragment.kt
index bb65aa2135..f15f841e60 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/OutOfTimeFragment.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/OutOfTimeFragment.kt
@@ -11,6 +11,7 @@ import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.flow.map
 import net.mullvad.mullvadvpn.R
 import net.mullvad.mullvadvpn.model.TunnelState
+import net.mullvad.mullvadvpn.ui.serviceconnection.AccountCache
 import net.mullvad.mullvadvpn.ui.widget.Button
 import net.mullvad.mullvadvpn.ui.widget.HeaderBar
 import net.mullvad.mullvadvpn.ui.widget.RedeemVoucherButton
@@ -18,8 +19,13 @@ import net.mullvad.mullvadvpn.ui.widget.SitePaymentButton
 import net.mullvad.talpid.tunnel.ActionAfterDisconnect
 import net.mullvad.talpid.tunnel.ErrorStateCause
 import org.joda.time.DateTime
+import org.koin.android.ext.android.inject
 
 class OutOfTimeFragment : ServiceDependentFragment(OnNoService.GoToLaunchScreen) {
+
+    // Injected dependencies
+    private val accountCache: AccountCache by inject()
+
     private lateinit var headerBar: HeaderBar
 
     private lateinit var sitePaymentButton: SitePaymentButton
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/RedeemVoucherDialogFragment.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/RedeemVoucherDialogFragment.kt
index dce06a445c..67536dca55 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/RedeemVoucherDialogFragment.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/RedeemVoucherDialogFragment.kt
@@ -29,6 +29,9 @@ import org.koin.android.ext.android.inject
 const val FULL_VOUCHER_CODE_LENGTH = "XXXX-XXXX-XXXX-XXXX".length
 
 class RedeemVoucherDialogFragment : DialogFragment() {
+
+    // Injected dependencies
+    private val accountCache: AccountCache by inject()
     private val serviceConnectionManager: ServiceConnectionManager by inject()
 
     private val jobTracker = JobTracker()
@@ -37,7 +40,6 @@ class RedeemVoucherDialogFragment : DialogFragment() {
     private lateinit var errorMessage: TextView
     private lateinit var voucherInput: EditText
 
-    private var accountCache: AccountCache? = null
     private var accountExpiry: DateTime? = null
     private var redeemButton: Button? = null
     private var voucherRedeemer: VoucherRedeemer? = null
@@ -54,16 +56,11 @@ class RedeemVoucherDialogFragment : DialogFragment() {
         parentActivity = context as MainActivity
 
         serviceConnectionManager.serviceNotifier.subscribe(this) { connection ->
-            accountCache = connection?.accountCache
             voucherRedeemer = connection?.voucherRedeemer
         }
 
-        accountCache?.apply {
-            jobTracker.newUiJob("updateExpiry") {
-                accountCache?.accountExpiryState?.collect { state ->
-                    accountExpiry = state.date()
-                }
-            }
+        jobTracker.newUiJob("updateExpiry") {
+            accountCache.accountExpiryState.collect { accountExpiry = it.date() }
         }
 
         updateRedeemButton()
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceDependentFragment.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceDependentFragment.kt
index 2911dcfb9f..be2998cfad 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceDependentFragment.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/ServiceDependentFragment.kt
@@ -5,7 +5,6 @@ import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
 import net.mullvad.mullvadvpn.R
-import net.mullvad.mullvadvpn.ui.serviceconnection.AccountCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.AppVersionInfoCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.AuthTokenCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.ConnectionProxy
@@ -13,7 +12,6 @@ import net.mullvad.mullvadvpn.ui.serviceconnection.CustomDns
 import net.mullvad.mullvadvpn.ui.serviceconnection.LocationInfoCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.RelayListListener
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionContainer
-import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionDeviceDataSource
 import net.mullvad.mullvadvpn.ui.serviceconnection.SettingsListener
 import net.mullvad.mullvadvpn.ui.serviceconnection.SplitTunneling
 
@@ -33,9 +31,6 @@ abstract class ServiceDependentFragment(private val onNoService: OnNoService) :
 
     private var state = State.Uninitialized
 
-    lateinit var accountCache: AccountCache
-        private set
-
     lateinit var appVersionInfoCache: AppVersionInfoCache
         private set
 
@@ -48,9 +43,6 @@ abstract class ServiceDependentFragment(private val onNoService: OnNoService) :
     lateinit var customDns: CustomDns
         private set
 
-    lateinit var deviceDataSource: ServiceConnectionDeviceDataSource
-        private set
-
     lateinit var locationInfoCache: LocationInfoCache
         private set
 
@@ -66,11 +58,9 @@ abstract class ServiceDependentFragment(private val onNoService: OnNoService) :
     override fun onNewServiceConnection(serviceConnectionContainer: ServiceConnectionContainer) {
         // This method is always either called first or after an `onNoServiceConnection`, so the
         // initialization of the fields doesn't have to be synchronized
-        accountCache = serviceConnectionContainer.accountCache
         appVersionInfoCache = serviceConnectionContainer.appVersionInfoCache
         authTokenCache = serviceConnectionContainer.authTokenCache
         connectionProxy = serviceConnectionContainer.connectionProxy
-        deviceDataSource = serviceConnectionContainer.deviceDataSource
         customDns = serviceConnectionContainer.customDns
         locationInfoCache = serviceConnectionContainer.locationInfoCache
         relayListListener = serviceConnectionContainer.relayListListener
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/SettingsFragment.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/SettingsFragment.kt
index 5d78098eb8..8d119807b8 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/SettingsFragment.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/SettingsFragment.kt
@@ -24,6 +24,7 @@ import net.mullvad.mullvadvpn.ui.widget.NavigateCell
 import org.koin.android.ext.android.inject
 
 class SettingsFragment : ServiceAwareFragment(), StatusBarPainter, NavigationBarPainter {
+    private val accountCache: AccountCache by inject()
     private val deviceRepository: DeviceRepository by inject()
 
     private lateinit var accountMenu: AccountCell
@@ -34,11 +35,9 @@ class SettingsFragment : ServiceAwareFragment(), StatusBarPainter, NavigationBar
 
     private var active = false
 
-    private var accountCache: AccountCache? = null
     private var versionInfoCache: AppVersionInfoCache? = null
 
     override fun onNewServiceConnection(serviceConnectionContainer: ServiceConnectionContainer) {
-        accountCache = serviceConnectionContainer.accountCache
         versionInfoCache = serviceConnectionContainer.appVersionInfoCache
 
         if (active) {
@@ -47,7 +46,6 @@ class SettingsFragment : ServiceAwareFragment(), StatusBarPainter, NavigationBar
     }
 
     override fun onNoServiceConnection() {
-        accountCache = null
         versionInfoCache = null
     }
 
@@ -132,18 +130,16 @@ class SettingsFragment : ServiceAwareFragment(), StatusBarPainter, NavigationBar
     }
 
     private fun configureListeners() {
-        accountCache?.apply {
-            jobTracker.newUiJob("updateAccountExpiry") {
-                accountExpiryState
-                    .map { state -> state.date() }
-                    .collect { expiryDate ->
-                        accountMenu.accountExpiry = expiryDate
-                    }
-            }
-
-            fetchAccountExpiry()
+        jobTracker.newUiJob("updateAccountExpiry") {
+            accountCache.accountExpiryState
+                .map { state -> state.date() }
+                .collect { expiryDate ->
+                    accountMenu.accountExpiry = expiryDate
+                }
         }
 
+        accountCache.fetchAccountExpiry()
+
         versionInfoCache?.onUpdate = {
             jobTracker.newUiJob("updateVersionInfo") {
                 updateVersionInfo()
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/WelcomeFragment.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/WelcomeFragment.kt
index f83660bd06..a5220bd758 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/WelcomeFragment.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/WelcomeFragment.kt
@@ -13,6 +13,7 @@ import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.collect
 import net.mullvad.mullvadvpn.R
 import net.mullvad.mullvadvpn.model.TunnelState
+import net.mullvad.mullvadvpn.ui.serviceconnection.AccountCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.DeviceRepository
 import net.mullvad.mullvadvpn.ui.widget.HeaderBar
 import net.mullvad.mullvadvpn.ui.widget.RedeemVoucherButton
@@ -23,6 +24,9 @@ import org.koin.android.ext.android.inject
 val POLL_INTERVAL: Long = 15 /* s */ * 1000 /* ms */
 
 class WelcomeFragment : ServiceDependentFragment(OnNoService.GoToLaunchScreen) {
+
+    // Injected dependencies
+    private val accountCache: AccountCache by inject()
     private val deviceRepository: DeviceRepository by inject()
 
     private lateinit var accountLabel: TextView
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/AccountCache.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/AccountCache.kt
index ed42b34a7d..c95bf14aeb 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/AccountCache.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/AccountCache.kt
@@ -1,82 +1,103 @@
 package net.mullvad.mullvadvpn.ui.serviceconnection
 
-import android.os.Messenger
-import kotlinx.coroutines.channels.BufferOverflow
-import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.asSharedFlow
+import kotlinx.coroutines.flow.SharedFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.flow.onStart
+import kotlinx.coroutines.flow.shareIn
+import kotlinx.coroutines.flow.stateIn
 import net.mullvad.mullvadvpn.ipc.Event
-import net.mullvad.mullvadvpn.ipc.EventDispatcher
-import net.mullvad.mullvadvpn.ipc.Request
 import net.mullvad.mullvadvpn.model.AccountCreationResult
 import net.mullvad.mullvadvpn.model.AccountExpiry
 import net.mullvad.mullvadvpn.model.AccountHistory
+import net.mullvad.mullvadvpn.util.flatMapReadyConnectionOrDefault
 
-class AccountCache(private val connection: Messenger, eventDispatcher: EventDispatcher) {
+class AccountCache(
+    private val serviceConnectionManager: ServiceConnectionManager,
+    dispatcher: CoroutineDispatcher = Dispatchers.IO
+) {
+    private val dataSource
+        get() = serviceConnectionManager.connectionState.value.readyContainer()?.accountDataSource
 
-    private val _accountCreationEvents = MutableSharedFlow<AccountCreationResult>(
-        extraBufferCapacity = 1,
-        onBufferOverflow = BufferOverflow.DROP_OLDEST
-    )
-    val accountCreationEvents = _accountCreationEvents.asSharedFlow()
+    private val _cachedCreatedAccount = MutableStateFlow<String?>(null)
+    val cachedCreatedAccount = _cachedCreatedAccount.asStateFlow()
 
-    private val _accountExpiryState = MutableStateFlow<AccountExpiry>(AccountExpiry.Missing)
-    val accountExpiryState = _accountExpiryState.asStateFlow()
-
-    private val _accountHistoryEvents = MutableSharedFlow<AccountHistory>(
-        extraBufferCapacity = 1,
-        onBufferOverflow = BufferOverflow.DROP_OLDEST
-    )
-    val accountHistoryEvents = _accountHistoryEvents.asSharedFlow()
-
-    private val _loginEvents = MutableSharedFlow<Event.LoginEvent>(
-        extraBufferCapacity = 1,
-        onBufferOverflow = BufferOverflow.DROP_OLDEST
-    )
-    val loginEvents = _loginEvents.asSharedFlow()
-
-    init {
-        eventDispatcher.apply {
-            registerHandler(Event.AccountCreationEvent::class) { event ->
-                _accountCreationEvents.tryEmit(event.result)
+    val accountCreationEvents: SharedFlow<AccountCreationResult> =
+        serviceConnectionManager.connectionState
+            .flatMapReadyConnectionOrDefault(flowOf()) { state ->
+                state.container.accountDataSource.accountCreationResult
             }
-
-            registerHandler(Event.AccountExpiryEvent::class) { event ->
-                _accountExpiryState.tryEmit(event.expiry)
+            .onEach {
+                _cachedCreatedAccount.value = (it as AccountCreationResult.Success).accountToken
             }
+            .shareIn(
+                CoroutineScope(dispatcher),
+                SharingStarted.WhileSubscribed()
+            )
 
-            registerHandler(Event.AccountHistoryEvent::class) { event ->
-                _accountHistoryEvents.tryEmit(event.history)
-            }
+    val accountExpiryState: StateFlow<AccountExpiry> = serviceConnectionManager.connectionState
+        .flatMapReadyConnectionOrDefault(flowOf()) { state ->
+            state.container.accountDataSource.accountExpiry
+        }
+        .onStart {
+            fetchAccountExpiry()
+        }
+        .stateIn(
+            CoroutineScope(dispatcher),
+            SharingStarted.WhileSubscribed(),
+            AccountExpiry.Missing
+        )
 
-            registerHandler(Event.LoginEvent::class) { event ->
-                _loginEvents.tryEmit(event)
-            }
+    val accountHistoryEvents: StateFlow<AccountHistory> = serviceConnectionManager.connectionState
+        .flatMapReadyConnectionOrDefault(flowOf()) { state ->
+            state.container.accountDataSource.accountHistory
         }
-    }
+        .onStart {
+            fetchAccountHistory()
+        }
+        .stateIn(
+            CoroutineScope(dispatcher),
+            SharingStarted.WhileSubscribed(),
+            AccountHistory.Missing
+        )
+
+    val loginEvents: SharedFlow<Event.LoginEvent> = serviceConnectionManager.connectionState
+        .flatMapReadyConnectionOrDefault(flowOf()) { state ->
+            state.container.accountDataSource.loginEvents
+        }
+        .shareIn(
+            CoroutineScope(dispatcher),
+            SharingStarted.WhileSubscribed()
+        )
 
-    fun createNewAccount() {
-        connection.send(Request.CreateAccount.message)
+    fun createAccount() {
+        dataSource?.createAccount()
     }
 
-    fun login(account: String) {
-        connection.send(Request.Login(account).message)
+    fun login(accountToken: String) {
+        dataSource?.login(accountToken)
     }
 
     fun logout() {
-        connection.send(Request.Logout.message)
+        dataSource?.logout()
     }
 
     fun fetchAccountExpiry() {
-        connection.send(Request.FetchAccountExpiry.message)
+        dataSource?.fetchAccountExpiry()
     }
 
     fun fetchAccountHistory() {
-        connection.send(Request.FetchAccountHistory.message)
+        dataSource?.fetchAccountHistory()
     }
 
     fun clearAccountHistory() {
-        connection.send(Request.ClearAccountHistory.message)
+        dataSource?.clearAccountHistory()
     }
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionAccountDataSource.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionAccountDataSource.kt
new file mode 100644
index 0000000000..23edd1efbf
--- /dev/null
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionAccountDataSource.kt
@@ -0,0 +1,60 @@
+package net.mullvad.mullvadvpn.ui.serviceconnection
+
+import android.os.Messenger
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.callbackFlow
+import net.mullvad.mullvadvpn.ipc.Event
+import net.mullvad.mullvadvpn.ipc.EventDispatcher
+import net.mullvad.mullvadvpn.ipc.Request
+
+class ServiceConnectionAccountDataSource(
+    private val connection: Messenger,
+    private val dispatcher: EventDispatcher
+) {
+    val accountCreationResult = callbackFlow {
+        val handler: (Event.AccountCreationEvent) -> Unit = { event ->
+            trySend(event.result)
+        }
+        dispatcher.registerHandler(Event.AccountCreationEvent::class, handler)
+        awaitClose {
+            // The current dispatcher doesn't support unregistration of handlers.
+        }
+    }
+
+    val accountExpiry = callbackFlow {
+        val handler: (Event.AccountExpiryEvent) -> Unit = { event ->
+            trySend(event.expiry)
+        }
+        dispatcher.registerHandler(Event.AccountExpiryEvent::class, handler)
+        awaitClose {
+            // The current dispatcher doesn't support unregistration of handlers.
+        }
+    }
+
+    val accountHistory = callbackFlow {
+        val handler: (Event.AccountHistoryEvent) -> Unit = { event ->
+            trySend(event.history)
+        }
+        dispatcher.registerHandler(Event.AccountHistoryEvent::class, handler)
+        awaitClose {
+            // The current dispatcher doesn't support unregistration of handlers.
+        }
+    }
+
+    val loginEvents = callbackFlow {
+        val handler: (Event.LoginEvent) -> Unit = { event ->
+            trySend(event)
+        }
+        dispatcher.registerHandler(Event.LoginEvent::class, handler)
+        awaitClose {
+            // The current dispatcher doesn't support unregistration of handlers.
+        }
+    }
+
+    fun createAccount() = connection.send(Request.CreateAccount.message)
+    fun login(accountToken: String) = connection.send(Request.Login(accountToken).message)
+    fun logout() = connection.send(Request.Logout.message)
+    fun fetchAccountExpiry() = connection.send(Request.FetchAccountExpiry.message)
+    fun fetchAccountHistory() = connection.send(Request.FetchAccountHistory.message)
+    fun clearAccountHistory() = connection.send(Request.ClearAccountHistory.message)
+}
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionContainer.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionContainer.kt
index 737e03d4e3..130dc1a8c9 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionContainer.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionContainer.kt
@@ -33,7 +33,7 @@ class ServiceConnectionContainer(
         named(SERVICE_CONNECTION_SCOPE), this
     )
 
-    val accountCache = AccountCache(connection, dispatcher)
+    val accountDataSource = ServiceConnectionAccountDataSource(connection, dispatcher)
     val authTokenCache = AuthTokenCache(connection, dispatcher)
     val connectionProxy = ConnectionProxy(connection, dispatcher)
     val deviceDataSource = ServiceConnectionDeviceDataSource(connection, dispatcher)
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/util/FlowUtils.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/util/FlowUtils.kt
index df750e64b4..a0a809f8ff 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/util/FlowUtils.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/util/FlowUtils.kt
@@ -13,8 +13,10 @@ import kotlinx.coroutines.channels.SendChannel
 import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.callbackFlow
+import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.take
 import net.mullvad.mullvadvpn.model.ServiceResult
+import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionState
 
 fun <T> SendChannel<T>.safeOffer(element: T): Boolean {
     return runCatching { offer(element) }.getOrDefault(false)
@@ -63,3 +65,16 @@ fun Context.bindServiceFlow(intent: Intent, flags: Int = 0): Flow<ServiceResult>
         }
     }
 }
+
+fun <R> Flow<ServiceConnectionState>.flatMapReadyConnectionOrDefault(
+    default: Flow<R>,
+    transform: (value: ServiceConnectionState.ConnectedReady) -> Flow<R>
+): Flow<R> {
+    return flatMapLatest { state ->
+        if (state is ServiceConnectionState.ConnectedReady) {
+            transform.invoke(state)
+        } else {
+            default
+        }
+    }
+}
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModel.kt
index d1749ee249..8487433362 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModel.kt
@@ -10,14 +10,16 @@ import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.stateIn
 import net.mullvad.mullvadvpn.compose.state.DeviceRevokedUiState
+import net.mullvad.mullvadvpn.ui.serviceconnection.AccountCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionContainer
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
 import net.mullvad.talpid.util.callbackFlowFromSubscription
 
-// TODO: Refactor AccountCache and ConnectionProxy and inject those rather than injecting
+// TODO: Refactor ConnectionProxy to be easily injectable rather than injecting
 //  ServiceConnectionManager here.
 class DeviceRevokedViewModel(
     private val serviceConnectionManager: ServiceConnectionManager,
+    private val accountCache: AccountCache,
     dispatcher: CoroutineDispatcher = Dispatchers.IO
 ) : ViewModel() {
 
@@ -46,7 +48,7 @@ class DeviceRevokedViewModel(
             if (container.connectionProxy.state.isSecured()) {
                 container.connectionProxy.disconnect()
             }
-            container.accountCache.logout()
+            accountCache.logout()
         }
     }
 
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt
index bcfd042580..ff391c668b 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt
@@ -3,55 +3,27 @@ package net.mullvad.mullvadvpn.viewmodel
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.CoroutineDispatcher
-import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.emptyFlow
 import kotlinx.coroutines.flow.first
-import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onStart
-import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.model.AccountCreationResult
-import net.mullvad.mullvadvpn.model.AccountHistory
 import net.mullvad.mullvadvpn.model.LoginResult
 import net.mullvad.mullvadvpn.ui.serviceconnection.AccountCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.DeviceRepository
-import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
-import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionState
 
 class LoginViewModel(
+    private val accountCache: AccountCache,
     private val deviceRepository: DeviceRepository,
-    private val serviceConnectionManager: ServiceConnectionManager,
     private val dispatcher: CoroutineDispatcher = Dispatchers.IO
 ) : ViewModel() {
     private val _uiState = MutableStateFlow<LoginUiState>(LoginUiState.Default)
     val uiState: StateFlow<LoginUiState> = _uiState
 
-    private val accountCache: AccountCache?
-        get() {
-            return serviceConnectionManager.connectionState.value.readyContainer()?.accountCache
-        }
-
-    val accountHistory = serviceConnectionManager.connectionState
-        .flatMapLatest { state ->
-            if (state is ServiceConnectionState.ConnectedReady) {
-                state.container.accountCache.accountHistoryEvents
-                    .onStart {
-                        state.container.accountCache.fetchAccountHistory()
-                    }
-            } else {
-                emptyFlow()
-            }
-        }
-        .stateIn(
-            scope = CoroutineScope(dispatcher),
-            started = SharingStarted.WhileSubscribed(),
-            initialValue = AccountHistory.Missing
-        )
+    val accountHistory = accountCache.accountHistoryEvents
 
     sealed class LoginUiState {
         object Default : LoginUiState()
@@ -69,54 +41,29 @@ class LoginViewModel(
         data class OtherError(val errorMessage: String) : LoginUiState()
     }
 
-    fun clearAccountHistory() {
-        accountCache.tryPerformAction(
-            errorMessageIfAccountCacheNotAvailable = SERVICE_NOT_CONNECTED_ERROR_MESSAGE
-        ) { cache ->
-            cache.clearAccountHistory()
-        }
-    }
+    fun clearAccountHistory() = accountCache.clearAccountHistory()
 
     fun clearState() {
         _uiState.value = LoginUiState.Default
     }
 
     fun createAccount() {
-        accountCache.tryPerformAction(
-            errorMessageIfAccountCacheNotAvailable = SERVICE_NOT_CONNECTED_ERROR_MESSAGE
-        ) { cache ->
-            _uiState.value = LoginUiState.CreatingAccount
-            viewModelScope.launch(dispatcher) {
-                _uiState.value = cache.accountCreationEvents
-                    .onStart { cache.createNewAccount() }
-                    .first()
-                    .mapToUiState()
-            }
+        _uiState.value = LoginUiState.CreatingAccount
+        viewModelScope.launch(dispatcher) {
+            _uiState.value = accountCache.accountCreationEvents
+                .onStart { accountCache.createAccount() }
+                .first()
+                .mapToUiState()
         }
     }
 
     fun login(accountToken: String) {
-        accountCache.tryPerformAction(
-            errorMessageIfAccountCacheNotAvailable = SERVICE_NOT_CONNECTED_ERROR_MESSAGE
-        ) { cache ->
-            _uiState.value = LoginUiState.Loading
-            viewModelScope.launch(dispatcher) {
-                _uiState.value = cache.loginEvents
-                    .onStart { cache.login(accountToken) }
-                    .map { it.result.mapToUiState(accountToken) }
-                    .first()
-            }
-        }
-    }
-
-    private fun AccountCache?.tryPerformAction(
-        errorMessageIfAccountCacheNotAvailable: String,
-        action: (AccountCache) -> Unit
-    ) {
-        if (this != null) {
-            action(this)
-        } else {
-            _uiState.value = LoginUiState.OtherError(errorMessageIfAccountCacheNotAvailable)
+        _uiState.value = LoginUiState.Loading
+        viewModelScope.launch(dispatcher) {
+            _uiState.value = accountCache.loginEvents
+                .onStart { accountCache.login(accountToken) }
+                .map { it.result.mapToUiState(accountToken) }
+                .first()
         }
     }
 
@@ -142,8 +89,4 @@ class LoginViewModel(
             else -> LoginUiState.OtherError(errorMessage = this.toString())
         }
     }
-
-    companion object {
-        private const val SERVICE_NOT_CONNECTED_ERROR_MESSAGE = "Not connected to service!"
-    }
 }
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModelTest.kt
index f507aa5b35..6dd1c586c9 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceRevokedViewModelTest.kt
@@ -17,6 +17,7 @@ import kotlinx.coroutines.test.TestCoroutineDispatcher
 import kotlinx.coroutines.test.runBlockingTest
 import net.mullvad.mullvadvpn.compose.state.DeviceRevokedUiState
 import net.mullvad.mullvadvpn.model.TunnelState
+import net.mullvad.mullvadvpn.ui.serviceconnection.AccountCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.ConnectionProxy
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionContainer
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
@@ -30,6 +31,9 @@ import org.junit.Test
 class DeviceRevokedViewModelTest {
 
     @MockK
+    private lateinit var mockedAccountCache: AccountCache
+
+    @MockK
     private lateinit var mockedServiceConnectionManager: ServiceConnectionManager
 
     private val serviceConnectionState =
@@ -44,6 +48,7 @@ class DeviceRevokedViewModelTest {
         every { mockedServiceConnectionManager.connectionState } returns serviceConnectionState
         viewModel = DeviceRevokedViewModel(
             mockedServiceConnectionManager,
+            mockedAccountCache,
             TestCoroutineDispatcher()
         )
     }
@@ -100,7 +105,7 @@ class DeviceRevokedViewModelTest {
         val mockedContainer = mockk<ServiceConnectionContainer>().also {
             every { it.connectionProxy.state } returns TunnelState.Disconnected
             every { it.connectionProxy.disconnect() } just Runs
-            every { it.accountCache.logout() } just Runs
+            every { mockedAccountCache.logout() } just Runs
         }
         serviceConnectionState.value = ServiceConnectionState.ConnectedReady(mockedContainer)
 
@@ -109,7 +114,7 @@ class DeviceRevokedViewModelTest {
 
         // Assert
         verify {
-            mockedContainer.accountCache.logout()
+            mockedAccountCache.logout()
         }
     }
 
@@ -119,7 +124,7 @@ class DeviceRevokedViewModelTest {
         val mockedContainer = mockk<ServiceConnectionContainer>().also {
             every { it.connectionProxy.state } returns TunnelState.Connected(mockk(), mockk())
             every { it.connectionProxy.disconnect() } just Runs
-            every { it.accountCache.logout() } just Runs
+            every { mockedAccountCache.logout() } just Runs
         }
         serviceConnectionState.value = ServiceConnectionState.ConnectedReady(mockedContainer)
 
@@ -129,7 +134,7 @@ class DeviceRevokedViewModelTest {
         // Assert
         verifyOrder {
             mockedContainer.connectionProxy.disconnect()
-            mockedContainer.accountCache.logout()
+            mockedAccountCache.logout()
         }
     }
 
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 59c52aab7d..6b44989f52 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
@@ -22,7 +22,6 @@ import net.mullvad.mullvadvpn.model.LoginResult
 import net.mullvad.mullvadvpn.ui.serviceconnection.AccountCache
 import net.mullvad.mullvadvpn.ui.serviceconnection.DeviceRepository
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionContainer
-import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionState
 import org.junit.Before
 import org.junit.Test
@@ -36,15 +35,12 @@ class LoginViewModelTest {
     private lateinit var mockedDeviceRepository: DeviceRepository
 
     @MockK
-    private lateinit var mockedServiceConnectionManager: ServiceConnectionManager
-
-    @MockK
     private lateinit var mockedServiceConnectionContainer: ServiceConnectionContainer
 
     private lateinit var loginViewModel: LoginViewModel
 
     private val accountCreationTestEvents = MutableSharedFlow<AccountCreationResult>()
-    private val accountHistoryTestEvents = MutableSharedFlow<AccountHistory>()
+    private val accountHistoryTestEvents = MutableStateFlow<AccountHistory>(AccountHistory.Missing)
     private val loginTestEvents = MutableSharedFlow<Event.LoginEvent>()
 
     private val serviceConnectionState =
@@ -58,15 +54,13 @@ class LoginViewModelTest {
         every { mockedAccountCache.accountCreationEvents } returns accountCreationTestEvents
         every { mockedAccountCache.accountHistoryEvents } returns accountHistoryTestEvents
         every { mockedAccountCache.loginEvents } returns loginTestEvents
-        every { mockedServiceConnectionManager.connectionState } returns serviceConnectionState
-        every { mockedServiceConnectionContainer.accountCache } returns mockedAccountCache
 
         serviceConnectionState.value =
             ServiceConnectionState.ConnectedReady(mockedServiceConnectionContainer)
 
         loginViewModel = LoginViewModel(
+            mockedAccountCache,
             mockedDeviceRepository,
-            mockedServiceConnectionManager,
             TestCoroutineDispatcher()
         )
     }