Fix multi-account switching during authentication and auto-sync

app/build.gradle.kts

@@ -201,6 +201,7 @@ dependencies {
     testImplementation(libs.junit)
     testImplementation(libs.androidx.test.core)
     testImplementation(libs.robolectric)
+    testImplementation(libs.work.testing)
     testImplementation(platform(libs.androidx.compose.bom))
     testImplementation(libs.androidx.compose.ui.test.junit4)
     androidTestImplementation(platform(libs.androidx.compose.bom))

app/src/main/java/com/anod/appwatcher/accounts/AuthAccountInitializer.kt

@@ -5,8 +5,14 @@ import android.os.Build
 import com.anod.appwatcher.preferences.Preferences
 import finsky.api.DfeApi
 import info.anodsplace.applog.AppLog
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.NonCancellable
+import kotlinx.coroutines.currentCoroutineContext
+import kotlinx.coroutines.ensureActive
+import kotlinx.coroutines.joinAll
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
+import kotlinx.coroutines.withContext
 
 class AuthTokenUnavailableException : IllegalStateException("Unable to retrieve authentication token")
 class AccountSessionBusyException : IllegalStateException("A Play Store synchronization is in progress")
@@ -36,6 +42,8 @@ class AuthAccountInitializer(
     private val dfeApi: DfeApi,
     private val playSessionCoordinator: PlaySessionCoordinator
 ) {
+    private val initializations = mutableSetOf<Job>()
+
     private val deviceRegistration = DeviceRegistration(
         preferences = preferences,
         dfeApi = dfeApi,
@@ -45,16 +53,41 @@ class AuthAccountInitializer(
     suspend fun initialize(
         account: Account,
         userInitiated: Boolean
-    ): AuthAccount =
-        if (userInitiated) {
-            playSessionCoordinator.withUserInitiatedSession {
-                initializeInSession(account, userInitiated = true)
+    ): AuthAccount {
+        val context = currentCoroutineContext()
+        context.ensureActive()
+        val job = checkNotNull(context[Job]) {
+            "Account initialization requires an active coroutine job"
+        }
+        val previousJobs = synchronized(this) {
+            val previous = initializations.toList()
+            if (userInitiated) {
+                previous.forEach { it.cancel() }
             }
-        } else {
-            playSessionCoordinator.withSession {
-                initializeInSession(account, userInitiated = false)
+            initializations.add(job)
+            previous
+        }
+        return try {
+            // Canceled initializations may still be finishing device registration.
+            withContext(NonCancellable) {
+                previousJobs.joinAll()
+            }
+            context.ensureActive()
+            if (userInitiated) {
+                playSessionCoordinator.withUserInitiatedSession {
+                    initializeInSession(account, userInitiated = true)
+                }
+            } else {
+                playSessionCoordinator.withSession {
+                    initializeInSession(account, userInitiated = false)
+                }
+            }
+        } finally {
+            synchronized(this) {
+                initializations.remove(job)
             }
         }
+    }
 
     private suspend fun initializeInSession(
         account: Account,

app/src/main/java/com/anod/appwatcher/search/SearchViewModel.kt

@@ -299,8 +299,9 @@ class SearchViewModel(initialState: SearchViewState) : BaseFlowViewModel<SearchV
         userInitiated: Boolean,
         resumingInteractiveAuth: Boolean
     ) {
-        if (accountInitializationJob?.isActive == true) {
-            return
+        if (userInitiated) {
+            pendingAccountInitialization = null
+            accountInitializationJob?.cancel()
         }
         accountInitializationJob = viewModelScope.launch {
             try {

app/src/main/java/com/anod/appwatcher/watchlist/MainViewModel.kt

@@ -7,6 +7,7 @@ import android.widget.Toast
 import androidx.annotation.StringRes
 import androidx.compose.runtime.Immutable
 import androidx.lifecycle.viewModelScope
+import androidx.work.Operation
 import com.anod.appwatcher.R
 import com.anod.appwatcher.accounts.AccountSelectionResult
 import com.anod.appwatcher.accounts.AccountSessionBusyException
@@ -36,6 +37,7 @@ import kotlinx.coroutines.ensureActive
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.filter
+import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.launch
 import org.koin.core.component.KoinComponent
@@ -183,6 +185,9 @@ class MainViewModel : BaseFlowViewModel<MainViewState, MainViewEvent, MainViewAc
     }
 
     private fun onResume() {
+        if (viewState.account != prefs.account) {
+            viewState = viewState.copy(account = prefs.account)
+        }
         initAccount()
         AppLog.d("mark updates as viewed.")
         prefs.isLastUpdatesViewed = true
@@ -229,8 +234,9 @@ class MainViewModel : BaseFlowViewModel<MainViewState, MainViewEvent, MainViewAc
         userInitiated: Boolean,
         resumingInteractiveAuth: Boolean
     ) {
-        if (accountInitializationJob?.isActive == true) {
-            return
+        if (userInitiated) {
+            pendingAccountInitialization = null
+            accountInitializationJob?.cancel()
         }
         val collectReports = prefs.collectCrashReports
         val initializer = authAccountInitializer
@@ -294,7 +300,7 @@ class MainViewModel : BaseFlowViewModel<MainViewState, MainViewEvent, MainViewAc
     private suspend fun scheduleRefresh() {
         SyncScheduler(context)
             .schedule(prefs.isRequiresCharging, prefs.isWifiOnly, prefs.updatesFrequency.toLong(), false)
-            .collect { }
+            .first { it !is Operation.State.IN_PROGRESS }
     }
 
     private fun showAccountErrorToast(errorMessage: String) {

app/src/test/java/com/anod/appwatcher/accounts/AuthAccountInitializerTest.kt

@@ -9,9 +9,11 @@ import java.io.IOException
 import java.util.ArrayDeque
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.CoroutineStart
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.async
 import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.withTimeout
 import org.junit.After
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertFalse
@@ -245,6 +247,49 @@ class AuthAccountInitializerTest {
         assertEquals(1, tokenProvider.requestedAccounts.size)
     }
 
+    @Test
+    fun explicitSelectionCancelsRunningInitializationBehindQueuedAutomaticRequest() = runBlocking {
+        preferences.account = AuthAccount("account@example.com", AuthTokenBlocking.ACCOUNT_TYPE, "", "", "")
+        val checkInStarted = CompletableDeferred<Unit>()
+        val finishCheckIn = CompletableDeferred<Unit>()
+        val dfeApi = FakeDfeApi().apply {
+            beforeCheckIn = {
+                checkInStarted.complete(Unit)
+                finishCheckIn.await()
+            }
+        }
+        val tokenProvider = RecordingTokenProvider("token-a", "token-b")
+        val initializer = AuthAccountInitializer(
+            preferences,
+            AuthTokenBlocking.create(tokenProvider),
+            dfeApi,
+            PlaySessionCoordinator()
+        )
+        val accountA = Account("account@example.com", AuthTokenBlocking.ACCOUNT_TYPE)
+        val accountB = Account("b@example.com", AuthTokenBlocking.ACCOUNT_TYPE)
+        val running = async { initializer.initialize(accountA, userInitiated = true) }
+
+        try {
+            withTimeout(5_000) { checkInStarted.await() }
+            val automatic = async(start = CoroutineStart.UNDISPATCHED) {
+                initializer.initialize(accountA, userInitiated = false)
+            }
+            val selected = async(start = CoroutineStart.UNDISPATCHED) {
+                initializer.initialize(accountB, userInitiated = true)
+            }
+
+            assertTrue("The running initialization must be canceled, not just its queued successor", running.isCancelled)
+            assertTrue(automatic.isCancelled)
+            finishCheckIn.complete(Unit)
+
+            assertEquals(accountB.name, withTimeout(5_000) { selected.await().name })
+            assertEquals(accountB.name, preferences.account?.name)
+            assertEquals(listOf(accountA.name, accountB.name), tokenProvider.requestedAccounts)
+        } finally {
+            finishCheckIn.complete(Unit)
+        }
+    }
+
     private fun completeAccount() = AuthAccount(
         name = "account@example.com",
         type = AuthTokenBlocking.ACCOUNT_TYPE,