Quick affordance interactor uses repository.

Based on a new feature flag, the quick affordance interactor may use the
new quick affordance repository.

Bug: 254853190
Test: added/updated unit tests. Manually verified that, when I hard-code
the left/right selections, I get them to show up on the lock screen.
Clicking on them also works as before.

Change-Id: Ie810d37d1dee97a0dcd23dc3c1ea2141be10750a

5 files changed, 531 insertions, 14 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/flags/Flags.kt b/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
index 844a311a988b..c5d598236c20 100644
--- a/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
+++ b/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
@@ -125,6 +125,14 @@ object Flags {
     // TODO(b/255607168): Tracking Bug
     @JvmField val DOZING_MIGRATION_1 = UnreleasedFlag(213)
 
+    /**
+     * Whether to enable the code powering customizable lock screen quick affordances.
+     *
+     * Note that this flag does not enable individual implementations of quick affordances like the
+     * new camera quick affordance. Look for individual flags for those.
+     */
+    @JvmField val CUSTOMIZABLE_LOCK_SCREEN_QUICK_AFFORDANCES = UnreleasedFlag(214, teamfood = false)
+
     // 300 - power menu
     // TODO(b/254512600): Tracking Bug
     @JvmField val POWER_MENU_LITE = ReleasedFlag(300)
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt
index 13d97aaf28da..92caa89bb0e8 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt
@@ -18,19 +18,30 @@
 package com.android.systemui.keyguard.domain.interactor
 
 import android.content.Intent
+import android.util.Log
 import com.android.internal.widget.LockPatternUtils
 import com.android.systemui.animation.Expandable
 import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.flags.FeatureFlags
+import com.android.systemui.flags.Flags
 import com.android.systemui.keyguard.data.quickaffordance.KeyguardQuickAffordanceConfig
+import com.android.systemui.keyguard.data.repository.KeyguardQuickAffordanceRepository
 import com.android.systemui.keyguard.domain.model.KeyguardQuickAffordanceModel
 import com.android.systemui.keyguard.domain.quickaffordance.KeyguardQuickAffordanceRegistry
+import com.android.systemui.keyguard.shared.model.KeyguardQuickAffordancePickerRepresentation
+import com.android.systemui.keyguard.shared.model.KeyguardSlotPickerRepresentation
 import com.android.systemui.keyguard.shared.quickaffordance.KeyguardQuickAffordancePosition
 import com.android.systemui.plugins.ActivityStarter
 import com.android.systemui.settings.UserTracker
+import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
 import com.android.systemui.statusbar.policy.KeyguardStateController
+import dagger.Lazy
 import javax.inject.Inject
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onStart
 
 @SysUISingleton
@@ -43,7 +54,12 @@ constructor(
     private val keyguardStateController: KeyguardStateController,
     private val userTracker: UserTracker,
     private val activityStarter: ActivityStarter,
+    private val featureFlags: FeatureFlags,
+    private val repository: Lazy<KeyguardQuickAffordanceRepository>,
 ) {
+    private val isUsingRepository: Boolean
+        get() = featureFlags.isEnabled(Flags.CUSTOMIZABLE_LOCK_SCREEN_QUICK_AFFORDANCES)
+
     /** Returns an observable for the quick affordance at the given position. */
     fun quickAffordance(
         position: KeyguardQuickAffordancePosition
@@ -72,7 +88,19 @@ constructor(
         configKey: String,
         expandable: Expandable?,
     ) {
-        @Suppress("UNCHECKED_CAST") val config = registry.get(configKey)
+        @Suppress("UNCHECKED_CAST")
+        val config =
+            if (isUsingRepository) {
+                val (slotId, decodedConfigKey) = configKey.decode()
+                repository.get().selections.value[slotId]?.find { it.key == decodedConfigKey }
+            } else {
+                registry.get(configKey)
+            }
+        if (config == null) {
+            Log.e(TAG, "Affordance config with key of \"$configKey\" not found!")
+            return
+        }
+
         when (val result = config.onTriggered(expandable)) {
             is KeyguardQuickAffordanceConfig.OnTriggeredResult.StartActivity ->
                 launchQuickAffordance(
@@ -84,28 +112,138 @@ constructor(
         }
     }
 
+    /**
+     * Selects an affordance with the given ID on the slot with the given ID.
+     *
+     * @return `true` if the affordance was selected successfully; `false` otherwise.
+     */
+    suspend fun select(slotId: String, affordanceId: String): Boolean {
+        check(isUsingRepository)
+
+        val slots = repository.get().getSlotPickerRepresentations()
+        val slot = slots.find { it.id == slotId } ?: return false
+        val selections =
+            repository.get().getSelections().getOrDefault(slotId, emptyList()).toMutableList()
+        val alreadySelected = selections.remove(affordanceId)
+        if (!alreadySelected) {
+            while (selections.size > 0 && selections.size >= slot.maxSelectedAffordances) {
+                selections.removeAt(0)
+            }
+        }
+
+        selections.add(affordanceId)
+
+        repository
+            .get()
+            .setSelections(
+                slotId = slotId,
+                affordanceIds = selections,
+            )
+
+        return true
+    }
+
+    /**
+     * Unselects one or all affordances from the slot with the given ID.
+     *
+     * @param slotId The ID of the slot.
+     * @param affordanceId The ID of the affordance to remove; if `null`, removes all affordances
+     * from the slot.
+     * @return `true` if the affordance was successfully removed; `false` otherwise (for example, if
+     * the affordance was not on the slot to begin with).
+     */
+    suspend fun unselect(slotId: String, affordanceId: String?): Boolean {
+        check(isUsingRepository)
+
+        val slots = repository.get().getSlotPickerRepresentations()
+        if (slots.find { it.id == slotId } == null) {
+            return false
+        }
+
+        if (affordanceId.isNullOrEmpty()) {
+            return if (
+                repository.get().getSelections().getOrDefault(slotId, emptyList()).isEmpty()
+            ) {
+                false
+            } else {
+                repository.get().setSelections(slotId = slotId, affordanceIds = emptyList())
+                true
+            }
+        }
+
+        val selections =
+            repository.get().getSelections().getOrDefault(slotId, emptyList()).toMutableList()
+        return if (selections.remove(affordanceId)) {
+            repository
+                .get()
+                .setSelections(
+                    slotId = slotId,
+                    affordanceIds = selections,
+                )
+            true
+        } else {
+            false
+        }
+    }
+
+    /** Returns affordance IDs indexed by slot ID, for all known slots. */
+    suspend fun getSelections(): Map<String, List<String>> {
+        check(isUsingRepository)
+
+        val selections = repository.get().getSelections()
+        return repository.get().getSlotPickerRepresentations().associate { slotRepresentation ->
+            slotRepresentation.id to (selections[slotRepresentation.id] ?: emptyList())
+        }
+    }
+
     private fun quickAffordanceInternal(
         position: KeyguardQuickAffordancePosition
     ): Flow<KeyguardQuickAffordanceModel> {
-        val configs = registry.getAll(position)
+        return if (isUsingRepository) {
+            repository
+                .get()
+                .selections
+                .map { it[position.toSlotId()] ?: emptyList() }
+                .flatMapLatest { configs -> combinedConfigs(position, configs) }
+        } else {
+            combinedConfigs(position, registry.getAll(position))
+        }
+    }
+
+    private fun combinedConfigs(
+        position: KeyguardQuickAffordancePosition,
+        configs: List<KeyguardQuickAffordanceConfig>,
+    ): Flow<KeyguardQuickAffordanceModel> {
+        if (configs.isEmpty()) {
+            return flowOf(KeyguardQuickAffordanceModel.Hidden)
+        }
+
         return combine(
             configs.map { config ->
-                // We emit an initial "Hidden" value to make sure that there's always an initial
-                // value and avoid subtle bugs where the downstream isn't receiving any values
-                // because one config implementation is not emitting an initial value. For example,
-                // see b/244296596.
+                // We emit an initial "Hidden" value to make sure that there's always an
+                // initial value and avoid subtle bugs where the downstream isn't receiving
+                // any values because one config implementation is not emitting an initial
+                // value. For example, see b/244296596.
                 config.lockScreenState.onStart {
                     emit(KeyguardQuickAffordanceConfig.LockScreenState.Hidden)
                 }
             }
         ) { states ->
             val index =
-                states.indexOfFirst { it is KeyguardQuickAffordanceConfig.LockScreenState.Visible }
+                states.indexOfFirst { state ->
+                    state is KeyguardQuickAffordanceConfig.LockScreenState.Visible
+                }
             if (index != -1) {
                 val visibleState =
                     states[index] as KeyguardQuickAffordanceConfig.LockScreenState.Visible
+                val configKey = configs[index].key
                 KeyguardQuickAffordanceModel.Visible(
-                    configKey = configs[index].key,
+                    configKey =
+                        if (isUsingRepository) {
+                            configKey.encode(position.toSlotId())
+                        } else {
+                            configKey
+                        },
                     icon = visibleState.icon,
                     activationState = visibleState.activationState,
                 )
@@ -145,4 +283,39 @@ constructor(
             )
         }
     }
+
+    private fun KeyguardQuickAffordancePosition.toSlotId(): String {
+        return when (this) {
+            KeyguardQuickAffordancePosition.BOTTOM_START ->
+                KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START
+            KeyguardQuickAffordancePosition.BOTTOM_END ->
+                KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END
+        }
+    }
+
+    private fun String.encode(slotId: String): String {
+        return "$slotId$DELIMITER$this"
+    }
+
+    private fun String.decode(): Pair<String, String> {
+        val splitUp = this.split(DELIMITER)
+        return Pair(splitUp[0], splitUp[1])
+    }
+
+    fun getAffordancePickerRepresentations(): List<KeyguardQuickAffordancePickerRepresentation> {
+        check(isUsingRepository)
+
+        return repository.get().getAffordancePickerRepresentations()
+    }
+
+    fun getSlotPickerRepresentations(): List<KeyguardSlotPickerRepresentation> {
+        check(isUsingRepository)
+
+        return repository.get().getSlotPickerRepresentations()
+    }
+
+    companion object {
+        private const val TAG = "KeyguardQuickAffordanceInteractor"
+        private const val DELIMITER = "::"
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorParameterizedTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorParameterizedTest.kt
index d6469406920b..8b6603d79244 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorParameterizedTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorParameterizedTest.kt
@@ -25,10 +25,14 @@ import com.android.systemui.animation.ActivityLaunchAnimator
 import com.android.systemui.animation.Expandable
 import com.android.systemui.common.shared.model.ContentDescription
 import com.android.systemui.common.shared.model.Icon
+import com.android.systemui.flags.FakeFeatureFlags
+import com.android.systemui.flags.Flags
 import com.android.systemui.keyguard.data.quickaffordance.BuiltInKeyguardQuickAffordanceKeys
 import com.android.systemui.keyguard.data.quickaffordance.FakeKeyguardQuickAffordanceConfig
 import com.android.systemui.keyguard.data.quickaffordance.KeyguardQuickAffordanceConfig
+import com.android.systemui.keyguard.data.quickaffordance.KeyguardQuickAffordanceSelectionManager
 import com.android.systemui.keyguard.data.repository.FakeKeyguardRepository
+import com.android.systemui.keyguard.data.repository.KeyguardQuickAffordanceRepository
 import com.android.systemui.keyguard.domain.quickaffordance.FakeKeyguardQuickAffordanceRegistry
 import com.android.systemui.keyguard.shared.quickaffordance.KeyguardQuickAffordancePosition
 import com.android.systemui.plugins.ActivityStarter
@@ -37,6 +41,8 @@ import com.android.systemui.statusbar.policy.KeyguardStateController
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.test.runBlockingTest
 import org.junit.Before
 import org.junit.Test
@@ -189,6 +195,8 @@ class KeyguardQuickAffordanceInteractorParameterizedTest : SysuiTestCase() {
                     /* startActivity= */ true,
                 ),
             )
+
+        private val IMMEDIATE = Dispatchers.Main.immediate
     }
 
     @Mock private lateinit var lockPatternUtils: LockPatternUtils
@@ -214,6 +222,19 @@ class KeyguardQuickAffordanceInteractorParameterizedTest : SysuiTestCase() {
 
         homeControls =
             FakeKeyguardQuickAffordanceConfig(BuiltInKeyguardQuickAffordanceKeys.HOME_CONTROLS)
+        val quickAccessWallet =
+            FakeKeyguardQuickAffordanceConfig(
+                BuiltInKeyguardQuickAffordanceKeys.QUICK_ACCESS_WALLET
+            )
+        val qrCodeScanner =
+            FakeKeyguardQuickAffordanceConfig(BuiltInKeyguardQuickAffordanceKeys.QR_CODE_SCANNER)
+        val quickAffordanceRepository =
+            KeyguardQuickAffordanceRepository(
+                scope = CoroutineScope(IMMEDIATE),
+                backgroundDispatcher = IMMEDIATE,
+                selectionManager = KeyguardQuickAffordanceSelectionManager(),
+                configs = setOf(homeControls, quickAccessWallet, qrCodeScanner),
+            )
         underTest =
             KeyguardQuickAffordanceInteractor(
                 keyguardInteractor = KeyguardInteractor(repository = FakeKeyguardRepository()),
@@ -226,12 +247,8 @@ class KeyguardQuickAffordanceInteractorParameterizedTest : SysuiTestCase() {
                                 ),
                             KeyguardQuickAffordancePosition.BOTTOM_END to
                                 listOf(
-                                    FakeKeyguardQuickAffordanceConfig(
-                                        BuiltInKeyguardQuickAffordanceKeys.QUICK_ACCESS_WALLET
-                                    ),
-                                    FakeKeyguardQuickAffordanceConfig(
-                                        BuiltInKeyguardQuickAffordanceKeys.QR_CODE_SCANNER
-                                    ),
+                                    quickAccessWallet,
+                                    qrCodeScanner,
                                 ),
                         ),
                     ),
@@ -239,6 +256,11 @@ class KeyguardQuickAffordanceInteractorParameterizedTest : SysuiTestCase() {
                 keyguardStateController = keyguardStateController,
                 userTracker = userTracker,
                 activityStarter = activityStarter,
+                featureFlags =
+                    FakeFeatureFlags().apply {
+                        set(Flags.CUSTOMIZABLE_LOCK_SCREEN_QUICK_AFFORDANCES, false)
+                    },
+                repository = { quickAffordanceRepository },
             )
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorTest.kt
index a5307504e138..33645354d9f3 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorTest.kt
@@ -22,22 +22,31 @@ import com.android.internal.widget.LockPatternUtils
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.common.shared.model.ContentDescription
 import com.android.systemui.common.shared.model.Icon
+import com.android.systemui.flags.FakeFeatureFlags
+import com.android.systemui.flags.Flags
 import com.android.systemui.keyguard.data.quickaffordance.BuiltInKeyguardQuickAffordanceKeys
 import com.android.systemui.keyguard.data.quickaffordance.FakeKeyguardQuickAffordanceConfig
 import com.android.systemui.keyguard.data.quickaffordance.KeyguardQuickAffordanceConfig
+import com.android.systemui.keyguard.data.quickaffordance.KeyguardQuickAffordanceSelectionManager
 import com.android.systemui.keyguard.data.repository.FakeKeyguardRepository
+import com.android.systemui.keyguard.data.repository.KeyguardQuickAffordanceRepository
 import com.android.systemui.keyguard.domain.model.KeyguardQuickAffordanceModel
 import com.android.systemui.keyguard.domain.quickaffordance.FakeKeyguardQuickAffordanceRegistry
 import com.android.systemui.keyguard.shared.quickaffordance.ActivationState
 import com.android.systemui.keyguard.shared.quickaffordance.KeyguardQuickAffordancePosition
 import com.android.systemui.plugins.ActivityStarter
 import com.android.systemui.settings.UserTracker
+import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
 import com.android.systemui.statusbar.policy.KeyguardStateController
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.runBlockingTest
 import kotlinx.coroutines.yield
 import org.junit.Before
@@ -47,6 +56,7 @@ import org.junit.runners.JUnit4
 import org.mockito.Mock
 import org.mockito.MockitoAnnotations
 
+@OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(JUnit4::class)
 class KeyguardQuickAffordanceInteractorTest : SysuiTestCase() {
@@ -62,6 +72,7 @@ class KeyguardQuickAffordanceInteractorTest : SysuiTestCase() {
     private lateinit var homeControls: FakeKeyguardQuickAffordanceConfig
     private lateinit var quickAccessWallet: FakeKeyguardQuickAffordanceConfig
     private lateinit var qrCodeScanner: FakeKeyguardQuickAffordanceConfig
+    private lateinit var featureFlags: FakeFeatureFlags
 
     @Before
     fun setUp() {
@@ -79,6 +90,18 @@ class KeyguardQuickAffordanceInteractorTest : SysuiTestCase() {
         qrCodeScanner =
             FakeKeyguardQuickAffordanceConfig(BuiltInKeyguardQuickAffordanceKeys.QR_CODE_SCANNER)
 
+        val quickAffordanceRepository =
+            KeyguardQuickAffordanceRepository(
+                scope = CoroutineScope(IMMEDIATE),
+                backgroundDispatcher = IMMEDIATE,
+                selectionManager = KeyguardQuickAffordanceSelectionManager(),
+                configs = setOf(homeControls, quickAccessWallet, qrCodeScanner),
+            )
+        featureFlags =
+            FakeFeatureFlags().apply {
+                set(Flags.CUSTOMIZABLE_LOCK_SCREEN_QUICK_AFFORDANCES, false)
+            }
+
         underTest =
             KeyguardQuickAffordanceInteractor(
                 keyguardInteractor = KeyguardInteractor(repository = repository),
@@ -100,6 +123,8 @@ class KeyguardQuickAffordanceInteractorTest : SysuiTestCase() {
                 keyguardStateController = keyguardStateController,
                 userTracker = userTracker,
                 activityStarter = activityStarter,
+                featureFlags = featureFlags,
+                repository = { quickAffordanceRepository },
             )
     }
 
@@ -203,6 +228,270 @@ class KeyguardQuickAffordanceInteractorTest : SysuiTestCase() {
             job.cancel()
         }
 
+    @Test
+    fun select() =
+        runBlocking(IMMEDIATE) {
+            featureFlags.set(Flags.CUSTOMIZABLE_LOCK_SCREEN_QUICK_AFFORDANCES, true)
+            homeControls.setState(
+                KeyguardQuickAffordanceConfig.LockScreenState.Visible(icon = ICON)
+            )
+            quickAccessWallet.setState(
+                KeyguardQuickAffordanceConfig.LockScreenState.Visible(icon = ICON)
+            )
+            qrCodeScanner.setState(
+                KeyguardQuickAffordanceConfig.LockScreenState.Visible(icon = ICON)
+            )
+
+            assertThat(underTest.getSelections())
+                .isEqualTo(
+                    mapOf<String, List<String>>(
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to emptyList(),
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END to emptyList(),
+                    )
+                )
+
+            var startConfig: KeyguardQuickAffordanceModel? = null
+            val job1 =
+                underTest
+                    .quickAffordance(KeyguardQuickAffordancePosition.BOTTOM_START)
+                    .onEach { startConfig = it }
+                    .launchIn(this)
+            var endConfig: KeyguardQuickAffordanceModel? = null
+            val job2 =
+                underTest
+                    .quickAffordance(KeyguardQuickAffordancePosition.BOTTOM_END)
+                    .onEach { endConfig = it }
+                    .launchIn(this)
+
+            underTest.select(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START, homeControls.key)
+            yield()
+            yield()
+            assertThat(startConfig)
+                .isEqualTo(
+                    KeyguardQuickAffordanceModel.Visible(
+                        configKey =
+                            KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START +
+                                "::${homeControls.key}",
+                        icon = ICON,
+                        activationState = ActivationState.NotSupported,
+                    )
+                )
+            assertThat(endConfig)
+                .isEqualTo(
+                    KeyguardQuickAffordanceModel.Hidden,
+                )
+            assertThat(underTest.getSelections())
+                .isEqualTo(
+                    mapOf(
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
+                            listOf(homeControls.key),
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END to emptyList(),
+                    )
+                )
+
+            underTest.select(
+                KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START,
+                quickAccessWallet.key
+            )
+            yield()
+            yield()
+            assertThat(startConfig)
+                .isEqualTo(
+                    KeyguardQuickAffordanceModel.Visible(
+                        configKey =
+                            KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START +
+                                "::${quickAccessWallet.key}",
+                        icon = ICON,
+                        activationState = ActivationState.NotSupported,
+                    )
+                )
+            assertThat(endConfig)
+                .isEqualTo(
+                    KeyguardQuickAffordanceModel.Hidden,
+                )
+            assertThat(underTest.getSelections())
+                .isEqualTo(
+                    mapOf(
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
+                            listOf(quickAccessWallet.key),
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END to emptyList(),
+                    )
+                )
+
+            underTest.select(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END, qrCodeScanner.key)
+            yield()
+            yield()
+            assertThat(startConfig)
+                .isEqualTo(
+                    KeyguardQuickAffordanceModel.Visible(
+                        configKey =
+                            KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START +
+                                "::${quickAccessWallet.key}",
+                        icon = ICON,
+                        activationState = ActivationState.NotSupported,
+                    )
+                )
+            assertThat(endConfig)
+                .isEqualTo(
+                    KeyguardQuickAffordanceModel.Visible(
+                        configKey =
+                            KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END +
+                                "::${qrCodeScanner.key}",
+                        icon = ICON,
+                        activationState = ActivationState.NotSupported,
+                    )
+                )
+            assertThat(underTest.getSelections())
+                .isEqualTo(
+                    mapOf(
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
+                            listOf(quickAccessWallet.key),
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END to
+                            listOf(qrCodeScanner.key),
+                    )
+                )
+
+            job1.cancel()
+            job2.cancel()
+        }
+
+    @Test
+    fun `unselect - one`() =
+        runBlocking(IMMEDIATE) {
+            featureFlags.set(Flags.CUSTOMIZABLE_LOCK_SCREEN_QUICK_AFFORDANCES, true)
+            homeControls.setState(
+                KeyguardQuickAffordanceConfig.LockScreenState.Visible(icon = ICON)
+            )
+            quickAccessWallet.setState(
+                KeyguardQuickAffordanceConfig.LockScreenState.Visible(icon = ICON)
+            )
+            qrCodeScanner.setState(
+                KeyguardQuickAffordanceConfig.LockScreenState.Visible(icon = ICON)
+            )
+
+            var startConfig: KeyguardQuickAffordanceModel? = null
+            val job1 =
+                underTest
+                    .quickAffordance(KeyguardQuickAffordancePosition.BOTTOM_START)
+                    .onEach { startConfig = it }
+                    .launchIn(this)
+            var endConfig: KeyguardQuickAffordanceModel? = null
+            val job2 =
+                underTest
+                    .quickAffordance(KeyguardQuickAffordancePosition.BOTTOM_END)
+                    .onEach { endConfig = it }
+                    .launchIn(this)
+            underTest.select(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START, homeControls.key)
+            yield()
+            yield()
+            underTest.select(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END, quickAccessWallet.key)
+            yield()
+            yield()
+
+            underTest.unselect(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START, homeControls.key)
+            yield()
+            yield()
+
+            assertThat(startConfig)
+                .isEqualTo(
+                    KeyguardQuickAffordanceModel.Hidden,
+                )
+            assertThat(endConfig)
+                .isEqualTo(
+                    KeyguardQuickAffordanceModel.Visible(
+                        configKey =
+                            KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END +
+                                "::${quickAccessWallet.key}",
+                        icon = ICON,
+                        activationState = ActivationState.NotSupported,
+                    )
+                )
+            assertThat(underTest.getSelections())
+                .isEqualTo(
+                    mapOf(
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to emptyList(),
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END to
+                            listOf(quickAccessWallet.key),
+                    )
+                )
+
+            underTest.unselect(
+                KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END,
+                quickAccessWallet.key
+            )
+            yield()
+            yield()
+
+            assertThat(startConfig)
+                .isEqualTo(
+                    KeyguardQuickAffordanceModel.Hidden,
+                )
+            assertThat(endConfig)
+                .isEqualTo(
+                    KeyguardQuickAffordanceModel.Hidden,
+                )
+            assertThat(underTest.getSelections())
+                .isEqualTo(
+                    mapOf<String, List<String>>(
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to emptyList(),
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END to emptyList(),
+                    )
+                )
+
+            job1.cancel()
+            job2.cancel()
+        }
+
+    @Test
+    fun `unselect - all`() =
+        runBlocking(IMMEDIATE) {
+            featureFlags.set(Flags.CUSTOMIZABLE_LOCK_SCREEN_QUICK_AFFORDANCES, true)
+            homeControls.setState(
+                KeyguardQuickAffordanceConfig.LockScreenState.Visible(icon = ICON)
+            )
+            quickAccessWallet.setState(
+                KeyguardQuickAffordanceConfig.LockScreenState.Visible(icon = ICON)
+            )
+            qrCodeScanner.setState(
+                KeyguardQuickAffordanceConfig.LockScreenState.Visible(icon = ICON)
+            )
+
+            underTest.select(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START, homeControls.key)
+            yield()
+            yield()
+            underTest.select(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END, quickAccessWallet.key)
+            yield()
+            yield()
+
+            underTest.unselect(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START, null)
+            yield()
+            yield()
+
+            assertThat(underTest.getSelections())
+                .isEqualTo(
+                    mapOf(
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to emptyList(),
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END to
+                            listOf(quickAccessWallet.key),
+                    )
+                )
+
+            underTest.unselect(
+                KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END,
+                null,
+            )
+            yield()
+            yield()
+
+            assertThat(underTest.getSelections())
+                .isEqualTo(
+                    mapOf<String, List<String>>(
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to emptyList(),
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END to emptyList(),
+                    )
+                )
+        }
+
     companion object {
         private val ICON: Icon = mock {
             whenever(this.contentDescription)
@@ -213,5 +502,6 @@ class KeyguardQuickAffordanceInteractorTest : SysuiTestCase() {
                 )
         }
         private const val CONTENT_DESCRIPTION_RESOURCE_ID = 1337
+        private val IMMEDIATE = Dispatchers.Main.immediate
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardBottomAreaViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardBottomAreaViewModelTest.kt
index d201d3b2334f..78148c4d3d1b 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardBottomAreaViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardBottomAreaViewModelTest.kt
@@ -23,10 +23,14 @@ import com.android.systemui.SysuiTestCase
 import com.android.systemui.animation.Expandable
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.doze.util.BurnInHelperWrapper
+import com.android.systemui.flags.FakeFeatureFlags
+import com.android.systemui.flags.Flags
 import com.android.systemui.keyguard.data.quickaffordance.BuiltInKeyguardQuickAffordanceKeys
 import com.android.systemui.keyguard.data.quickaffordance.FakeKeyguardQuickAffordanceConfig
 import com.android.systemui.keyguard.data.quickaffordance.KeyguardQuickAffordanceConfig
+import com.android.systemui.keyguard.data.quickaffordance.KeyguardQuickAffordanceSelectionManager
 import com.android.systemui.keyguard.data.repository.FakeKeyguardRepository
+import com.android.systemui.keyguard.data.repository.KeyguardQuickAffordanceRepository
 import com.android.systemui.keyguard.domain.interactor.KeyguardBottomAreaInteractor
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
 import com.android.systemui.keyguard.domain.interactor.KeyguardQuickAffordanceInteractor
@@ -41,6 +45,8 @@ import com.android.systemui.util.mockito.mock
 import com.google.common.truth.Truth.assertThat
 import kotlin.math.max
 import kotlin.math.min
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.test.runBlockingTest
@@ -109,6 +115,18 @@ class KeyguardBottomAreaViewModelTest : SysuiTestCase() {
         whenever(userTracker.userHandle).thenReturn(mock())
         whenever(lockPatternUtils.getStrongAuthForUser(anyInt()))
             .thenReturn(LockPatternUtils.StrongAuthTracker.STRONG_AUTH_NOT_REQUIRED)
+        val quickAffordanceRepository =
+            KeyguardQuickAffordanceRepository(
+                scope = CoroutineScope(IMMEDIATE),
+                backgroundDispatcher = IMMEDIATE,
+                selectionManager = KeyguardQuickAffordanceSelectionManager(),
+                configs =
+                    setOf(
+                        homeControlsQuickAffordanceConfig,
+                        quickAccessWalletAffordanceConfig,
+                        qrCodeScannerAffordanceConfig,
+                    ),
+            )
         underTest =
             KeyguardBottomAreaViewModel(
                 keyguardInteractor = keyguardInteractor,
@@ -120,6 +138,11 @@ class KeyguardBottomAreaViewModelTest : SysuiTestCase() {
                         keyguardStateController = keyguardStateController,
                         userTracker = userTracker,
                         activityStarter = activityStarter,
+                        featureFlags =
+                            FakeFeatureFlags().apply {
+                                set(Flags.CUSTOMIZABLE_LOCK_SCREEN_QUICK_AFFORDANCES, false)
+                            },
+                        repository = { quickAffordanceRepository },
                     ),
                 bottomAreaInteractor = KeyguardBottomAreaInteractor(repository = repository),
                 burnInHelperWrapper = burnInHelperWrapper,
@@ -569,5 +592,6 @@ class KeyguardBottomAreaViewModelTest : SysuiTestCase() {
     companion object {
         private const val DEFAULT_BURN_IN_OFFSET = 5
         private const val RETURNED_BURN_IN_OFFSET = 3
+        private val IMMEDIATE = Dispatchers.Main.immediate
     }
 }