Improve account data fetching

Add a new account data fetch every time a user enters the connect screen
This is limited to at a maximum one fetch every minute.

Add a check that the user is still logged in to the same account
before updating the account data cache.

Change account fetching behavior in the account screen to
fetch on every enter instead of init.

2 files changed, 105 insertions, 9 deletions

diff --git a/android/lib/shared/src/main/kotlin/net/mullvad/mullvadvpn/lib/shared/AccountRepository.kt b/android/lib/shared/src/main/kotlin/net/mullvad/mullvadvpn/lib/shared/AccountRepository.kt
index 2e922a0895..bfb1918875 100644
--- a/android/lib/shared/src/main/kotlin/net/mullvad/mullvadvpn/lib/shared/AccountRepository.kt
+++ b/android/lib/shared/src/main/kotlin/net/mullvad/mullvadvpn/lib/shared/AccountRepository.kt
@@ -1,7 +1,6 @@
 package net.mullvad.mullvadvpn.lib.shared
 
 import arrow.core.Either
-import arrow.core.raise.nullable
 import java.time.ZonedDateTime
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.MutableSharedFlow
@@ -27,6 +26,7 @@ class AccountRepository(
     private val deviceRepository: DeviceRepository,
     val scope: CoroutineScope,
 ) {
+    private var lastSuccessfulAccountDataFetch: ZonedDateTime? = null
 
     private val _mutableAccountDataCache: MutableSharedFlow<AccountData> = MutableSharedFlow()
 
@@ -43,7 +43,10 @@ class AccountRepository(
                 managementService.deviceState.map { deviceState ->
                     when (deviceState) {
                         is DeviceState.LoggedIn -> {
-                            managementService.getAccountData(deviceState.accountNumber).getOrNull()
+                            managementService
+                                .getAccountData(deviceState.accountNumber)
+                                .getOrNull()
+                                ?.also { lastSuccessfulAccountDataFetch = ZonedDateTime.now() }
                         }
                         DeviceState.LoggedOut,
                         DeviceState.Revoked -> null
@@ -72,15 +75,27 @@ class AccountRepository(
     suspend fun clearAccountHistory(): Either<ClearAccountHistoryError, Unit> =
         managementService.clearAccountHistory().onRight { _mutableAccountHistory.value = null }
 
-    suspend fun getAccountData(): AccountData? = nullable {
-        val deviceState = ensureNotNull(deviceRepository.deviceState.value as? DeviceState.LoggedIn)
+    /*
+     * Fetches the account data from the server, and updates the cache.
+     * Unless force is true, it will only fetch if no fetch was made in the last minute.
+     */
+    suspend fun refreshAccountData(ignoreTimeout: Boolean = true) {
+        // Only refresh if logged in
+        val deviceState = deviceRepository.deviceState.value as? DeviceState.LoggedIn ?: return
 
-        val accountData =
-            managementService.getAccountData(deviceState.accountNumber).getOrNull().bind()
+        if (ignoreTimeout || lastSuccessfulAccountDataFetch.canFetchAccountData()) {
+            val accountData =
+                managementService.getAccountData(deviceState.accountNumber).getOrNull()
+            lastSuccessfulAccountDataFetch = ZonedDateTime.now()
 
-        // Update stateflow cache
-        _mutableAccountDataCache.emit(accountData)
-        accountData
+            // Update stateflow cache, only update if device state is still logged in and using the
+            // same account number
+            deviceRepository.deviceState.value?.let {
+                if (it is DeviceState.LoggedIn && it.accountNumber == accountData?.accountNumber) {
+                    _mutableAccountDataCache.emit(accountData)
+                }
+            }
+        }
     }
 
     suspend fun getWebsiteAuthToken(): WebsiteAuthToken? =
@@ -93,4 +108,7 @@ class AccountRepository(
     fun resetIsNewAccount() {
         _isNewAccount.value = false
     }
+
+    private fun ZonedDateTime?.canFetchAccountData(): Boolean =
+        this == null || this.isBefore(ZonedDateTime.now().minusMinutes(1))
 }
diff --git a/android/lib/shared/src/test/kotlin/net/mullvad/mullvadvpn/lib/shared/AccountRepositoryTest.kt b/android/lib/shared/src/test/kotlin/net/mullvad/mullvadvpn/lib/shared/AccountRepositoryTest.kt
new file mode 100644
index 0000000000..ff2ecb88b2
--- /dev/null
+++ b/android/lib/shared/src/test/kotlin/net/mullvad/mullvadvpn/lib/shared/AccountRepositoryTest.kt
@@ -0,0 +1,78 @@
+package net.mullvad.mullvadvpn.lib.shared
+
+import arrow.core.right
+import io.mockk.coEvery
+import io.mockk.coVerify
+import io.mockk.every
+import io.mockk.mockk
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.MainScope
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.test.runTest
+import net.mullvad.mullvadvpn.lib.common.test.TestCoroutineRule
+import net.mullvad.mullvadvpn.lib.daemon.grpc.ManagementService
+import net.mullvad.mullvadvpn.lib.model.AccountData
+import net.mullvad.mullvadvpn.lib.model.AccountNumber
+import net.mullvad.mullvadvpn.lib.model.DeviceState
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.extension.ExtendWith
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@ExtendWith(TestCoroutineRule::class)
+class AccountRepositoryTest {
+
+    private val mockManagementService: ManagementService = mockk()
+    private val mockDeviceRepository: DeviceRepository = mockk()
+
+    private val mockDeviceStateFlow = MutableStateFlow<DeviceState>(DeviceState.LoggedOut)
+
+    private lateinit var accountRepository: AccountRepository
+
+    @BeforeEach
+    fun setup() {
+        every { mockDeviceRepository.deviceState } returns mockDeviceStateFlow
+        every { mockManagementService.deviceState } returns mockDeviceStateFlow
+
+        accountRepository =
+            AccountRepository(
+                managementService = mockManagementService,
+                deviceRepository = mockDeviceRepository,
+                scope = MainScope(),
+            )
+    }
+
+    @Test
+    fun `given force is true should always call managementService getAccountData`() = runTest {
+        // Arrange
+        val accountData: AccountData = mockk()
+        val accountNumber = AccountNumber("1234567890")
+        every { accountData.accountNumber } returns accountNumber
+        coEvery { mockManagementService.getAccountData(accountNumber) } returns accountData.right()
+
+        // Act
+        mockDeviceStateFlow.emit(DeviceState.LoggedIn(accountNumber, mockk(relaxed = true)))
+        accountRepository.refreshAccountData(ignoreTimeout = true)
+
+        // Assert
+        coVerify { mockManagementService.getAccountData(accountNumber) }
+    }
+
+    @Test
+    fun `given last latestAccountDataFetch null should always call managementService getAccountData`() =
+        runTest {
+            // Arrange
+            val accountData: AccountData = mockk()
+            val accountNumber = AccountNumber("1234567890")
+            every { accountData.accountNumber } returns accountNumber
+            coEvery { mockManagementService.getAccountData(accountNumber) } returns
+                accountData.right()
+
+            // Act
+            mockDeviceStateFlow.emit(DeviceState.LoggedIn(accountNumber, mockk(relaxed = true)))
+            accountRepository.refreshAccountData(ignoreTimeout = false)
+
+            // Assert
+            coVerify { mockManagementService.getAccountData(accountNumber) }
+        }
+}