Restore QS data even with just one intent

In the case where we are restoring from a device that didn't have a
`qs_auto_tiles` setting, we can still restore the user tiles and assume
that the setting was empty. We listen for user setup complete that
indicates that settings provider has finished restoring.

Restoring just from qs_auto_tiles makes no sense, as the user cares
about the current tiles. In that case, if we only got that intent, we
log an error and ignore it.

Test: atest QSSettingsRestoredBroadcastRepositoryTest
Test: manual, restore from a device with deleted qs_auto_tiles
Fixes: 325253849
Flag: ACONFIG com.android.systemui.qs_new_pipeline NEXTFOOD
Change-Id: I2a9709301a0474d7d6529667638c5c63dfe3c8ef

4 files changed, 206 insertions, 20 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/data/repository/QSSettingsRestoredBroadcastRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/data/repository/QSSettingsRestoredBroadcastRepositoryTest.kt
index ff8a9bd019fb..39851b6f21a7 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/data/repository/QSSettingsRestoredBroadcastRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/data/repository/QSSettingsRestoredBroadcastRepositoryTest.kt
@@ -8,6 +8,7 @@ import com.android.systemui.SysuiTestCase
 import com.android.systemui.broadcast.FakeBroadcastDispatcher
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.qs.pipeline.shared.logging.QSPipelineLogger
+import com.android.systemui.statusbar.policy.FakeDeviceProvisionedController
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.StandardTestDispatcher
@@ -28,6 +29,7 @@ class QSSettingsRestoredBroadcastRepositoryTest : SysuiTestCase() {
     private val testScope = TestScope(dispatcher)
 
     @Mock private lateinit var pipelineLogger: QSPipelineLogger
+    private val deviceProvisionedController = FakeDeviceProvisionedController()
 
     private lateinit var underTest: QSSettingsRestoredBroadcastRepository
 
@@ -38,6 +40,7 @@ class QSSettingsRestoredBroadcastRepositoryTest : SysuiTestCase() {
         underTest =
             QSSettingsRestoredBroadcastRepository(
                 fakeBroadcastDispatcher,
+                deviceProvisionedController,
                 pipelineLogger,
                 testScope.backgroundScope,
                 dispatcher,
@@ -176,6 +179,100 @@ class QSSettingsRestoredBroadcastRepositoryTest : SysuiTestCase() {
             }
         }
 
+    @Test
+    fun restoreAfterUserSetup_singleTilesRestoredBroadcast() =
+        testScope.runTest {
+            runCurrent()
+            val restoreData by collectLastValue(underTest.restoreData)
+            val user = 0
+
+            val tilesIntent =
+                createRestoreIntent(
+                    RestoreType.TILES,
+                    CURRENT_TILES,
+                    RESTORED_TILES,
+                )
+
+            sendIntentForUser(tilesIntent, user)
+
+            deviceProvisionedController.setUserSetup(user)
+
+            with(restoreData!!) {
+                assertThat(restoredTiles).isEqualTo(RESTORED_TILES.toTilesList())
+                assertThat(restoredAutoAddedTiles).isEmpty()
+                assertThat(userId).isEqualTo(user)
+            }
+        }
+
+    @Test
+    fun restoreAfterUserSetup_singleAutoAddRestoredBroadcast_noRestore() =
+        testScope.runTest {
+            runCurrent()
+            val restoreData by collectLastValue(underTest.restoreData)
+            val user = 0
+
+            val autoAddIntent =
+                createRestoreIntent(
+                    RestoreType.AUTOADD,
+                    CURRENT_AUTO_ADDED_TILES,
+                    RESTORED_AUTO_ADDED_TILES,
+                )
+
+            sendIntentForUser(autoAddIntent, user)
+
+            deviceProvisionedController.setUserSetup(user)
+
+            assertThat(restoreData).isNull()
+        }
+
+    @Test
+    fun restoreAfterUserSetup_otherUserFinishedSetup_noRestore() =
+        testScope.runTest {
+            runCurrent()
+            val restoreData by collectLastValue(underTest.restoreData)
+            val user = 0
+
+            val tilesIntent =
+                createRestoreIntent(
+                    RestoreType.TILES,
+                    CURRENT_TILES,
+                    RESTORED_TILES,
+                )
+
+            sendIntentForUser(tilesIntent, user)
+
+            deviceProvisionedController.setUserSetup(user + 1)
+
+            assertThat(restoreData).isNull()
+        }
+
+    @Test
+    fun restoreAfterUserSetup_otherUserFinishedSetup_thenCorrectUser_restored() =
+        testScope.runTest {
+            runCurrent()
+            val restoreData by collectLastValue(underTest.restoreData)
+            val user = 0
+
+            val tilesIntent =
+                createRestoreIntent(
+                    RestoreType.TILES,
+                    CURRENT_TILES,
+                    RESTORED_TILES,
+                )
+
+            sendIntentForUser(tilesIntent, user)
+
+            deviceProvisionedController.setUserSetup(user + 1)
+            runCurrent()
+            deviceProvisionedController.setUserSetup(user)
+
+            with(restoreData!!) {
+                assertThat(restoredTiles).isEqualTo(RESTORED_TILES.toTilesList())
+                assertThat(restoredAutoAddedTiles).isEmpty()
+                assertThat(userId).isEqualTo(user)
+            }
+        }
+
     private fun sendIntentForUser(intent: Intent, userId: Int) {
         fakeBroadcastDispatcher.sendIntentToMatchingReceiversOnly(
             context,
diff --git a/packages/SystemUI/src/com/android/systemui/qs/pipeline/data/repository/QSSettingsRestoredRepository.kt b/packages/SystemUI/src/com/android/systemui/qs/pipeline/data/repository/QSSettingsRestoredRepository.kt
index e718eea09665..a2bb9e9a2f63 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/pipeline/data/repository/QSSettingsRestoredRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/pipeline/data/repository/QSSettingsRestoredRepository.kt
@@ -6,6 +6,8 @@ import android.os.UserHandle
 import android.provider.Settings
 import android.util.Log
 import com.android.systemui.broadcast.BroadcastDispatcher
+import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
+import com.android.systemui.common.shared.model.PackageChangeModel.Empty.user
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dagger.qualifiers.Background
@@ -13,18 +15,28 @@ import com.android.systemui.qs.pipeline.data.model.RestoreData
 import com.android.systemui.qs.pipeline.data.repository.QSSettingsRestoredRepository.Companion.BUFFER_CAPACITY
 import com.android.systemui.qs.pipeline.data.repository.TilesSettingConverter.toTilesList
 import com.android.systemui.qs.pipeline.shared.logging.QSPipelineLogger
+import com.android.systemui.statusbar.policy.DeviceProvisionedController
+import com.android.systemui.util.kotlin.emitOnStart
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.asFlow
 import kotlinx.coroutines.flow.buffer
 import kotlinx.coroutines.flow.catch
 import kotlinx.coroutines.flow.filter
-import kotlinx.coroutines.flow.flow
+import kotlinx.coroutines.flow.flattenConcat
 import kotlinx.coroutines.flow.flowOn
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.mapNotNull
+import kotlinx.coroutines.flow.merge
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.shareIn
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
 
 /** Provides restored data (from Backup and Restore) for Quick Settings pipeline */
 interface QSSettingsRestoredRepository {
@@ -44,34 +56,87 @@ class QSSettingsRestoredBroadcastRepository
 @Inject
 constructor(
     broadcastDispatcher: BroadcastDispatcher,
+    private val deviceProvisionedController: DeviceProvisionedController,
     logger: QSPipelineLogger,
     @Application private val scope: CoroutineScope,
     @Background private val backgroundDispatcher: CoroutineDispatcher,
 ) : QSSettingsRestoredRepository {
 
+    private val onUserSetupChangedForSomeUser =
+        conflatedCallbackFlow {
+                val callback =
+                    object : DeviceProvisionedController.DeviceProvisionedListener {
+                        override fun onUserSetupChanged() {
+                            trySend(Unit)
+                        }
+                    }
+                deviceProvisionedController.addCallback(callback)
+                awaitClose { deviceProvisionedController.removeCallback(callback) }
+            }
+            .emitOnStart()
+
+    @OptIn(ExperimentalCoroutinesApi::class)
     override val restoreData =
-        flow {
+        run {
+                val mutex = Mutex()
                 val firstIntent = mutableMapOf<Int, Intent>()
-                broadcastDispatcher
-                    .broadcastFlow(INTENT_FILTER, UserHandle.ALL) { intent, receiver ->
-                        intent to receiver.sendingUserId
-                    }
-                    .filter { it.first.isCorrectSetting() }
-                    .collect { (intent, user) ->
-                        if (user !in firstIntent) {
-                            firstIntent[user] = intent
-                        } else {
-                            val firstRestored = firstIntent.remove(user)!!
-                            emit(processIntents(user, firstRestored, intent))
+
+                val restoresFromTwoBroadcasts: Flow<RestoreData> =
+                    broadcastDispatcher
+                        .broadcastFlow(INTENT_FILTER, UserHandle.ALL) { intent, receiver ->
+                            intent to receiver.sendingUserId
                         }
-                    }
+                        .filter { it.first.isCorrectSetting() }
+                        .mapNotNull { (intent, user) ->
+                            mutex.withLock {
+                                if (user !in firstIntent) {
+                                    firstIntent[user] = intent
+                                    null
+                                } else {
+                                    val firstRestored = firstIntent.remove(user)!!
+                                    processIntents(user, firstRestored, intent)
+                                }
+                            }
+                        }
+                        .catch { Log.e(TAG, "Error parsing broadcast", it) }
+
+                val restoresFromUserSetup: Flow<RestoreData> =
+                    onUserSetupChangedForSomeUser
+                        .map {
+                            mutex.withLock {
+                                firstIntent
+                                    .filter { (userId, _) ->
+                                        deviceProvisionedController.isUserSetup(userId)
+                                    }
+                                    .onEach { firstIntent.remove(it.key) }
+                                    .map { processSingleIntent(it.key, it.value) }
+                                    .asFlow()
+                            }
+                        }
+                        .flattenConcat()
+                        .catch { Log.e(TAG, "Error parsing tiles intent after user setup", it) }
+                        .onEach { logger.logSettingsRestoredOnUserSetupComplete(it.userId) }
+                merge(restoresFromTwoBroadcasts, restoresFromUserSetup)
             }
-            .catch { Log.e(TAG, "Error parsing broadcast", it) }
             .flowOn(backgroundDispatcher)
             .buffer(BUFFER_CAPACITY)
             .shareIn(scope, SharingStarted.Eagerly)
             .onEach(logger::logSettingsRestored)
 
+    private fun processSingleIntent(user: Int, intent: Intent): RestoreData {
+        intent.validateIntent()
+        if (intent.getStringExtra(Intent.EXTRA_SETTING_NAME) != TILES_SETTING) {
+            throw IllegalStateException(
+                "Single intent restored for user $user is not tiles: $intent"
+            )
+        }
+        return RestoreData(
+            (intent.getStringExtra(Intent.EXTRA_SETTING_NEW_VALUE) ?: "").toTilesList(),
+            emptySet(),
+            user,
+        )
+    }
+
     private fun processIntents(user: Int, intent1: Intent, intent2: Intent): RestoreData {
         intent1.validateIntent()
         intent2.validateIntent()
diff --git a/packages/SystemUI/src/com/android/systemui/qs/pipeline/shared/logging/QSPipelineLogger.kt b/packages/SystemUI/src/com/android/systemui/qs/pipeline/shared/logging/QSPipelineLogger.kt
index 7d2c6c8f70fb..e237ca9f462b 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/pipeline/shared/logging/QSPipelineLogger.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/pipeline/shared/logging/QSPipelineLogger.kt
@@ -221,6 +221,15 @@ constructor(
         )
     }
 
+    fun logSettingsRestoredOnUserSetupComplete(userId: Int) {
+        restoreLogBuffer.log(
+            RESTORE_TAG,
+            LogLevel.DEBUG,
+            { int1 = userId },
+            { "Restored from single intent after user setup complete for user $int1" }
+        )
+    }
+
     fun logSettingsRestored(restoreData: RestoreData) {
         restoreLogBuffer.log(
             RESTORE_TAG,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/FakeDeviceProvisionedController.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/FakeDeviceProvisionedController.kt
index 0c2b115a8af5..aa2c2a2e1ec3 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/FakeDeviceProvisionedController.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/FakeDeviceProvisionedController.kt
@@ -2,30 +2,45 @@ package com.android.systemui.statusbar.policy
 
 class FakeDeviceProvisionedController : DeviceProvisionedController {
     @JvmField var deviceProvisioned = true
+    @JvmField var currentUser = 0
+
+    private val callbacks = mutableSetOf<DeviceProvisionedController.DeviceProvisionedListener>()
+    private val usersSetup = mutableSetOf<Int>()
 
     override fun addCallback(listener: DeviceProvisionedController.DeviceProvisionedListener) {
-        TODO("Not yet implemented")
+        callbacks.add(listener)
     }
 
     override fun removeCallback(listener: DeviceProvisionedController.DeviceProvisionedListener) {
-        TODO("Not yet implemented")
+        callbacks.remove(listener)
     }
 
     override fun isDeviceProvisioned() = deviceProvisioned
 
+    @Deprecated("Deprecated in Java")
     override fun getCurrentUser(): Int {
-        TODO("Not yet implemented")
+        return currentUser
     }
 
     override fun isUserSetup(user: Int): Boolean {
-        TODO("Not yet implemented")
+        return user in usersSetup
     }
 
     override fun isCurrentUserSetup(): Boolean {
-        TODO("Not yet implemented")
+        return currentUser in usersSetup
     }
 
     override fun isFrpActive(): Boolean {
         TODO("Not yet implemented")
     }
+
+    fun setCurrentUser(userId: Int) {
+        currentUser = userId
+        callbacks.toSet().forEach { it.onUserSwitched() }
+    }
+
+    fun setUserSetup(userId: Int) {
+        usersSetup.add(userId)
+        callbacks.toSet().forEach { it.onUserSetupChanged() }
+    }
 }