Fix crash from resetting shortcuts to none

Wallpaper picker crashes from resetting both shortcuts to none, and
behaves incorrectly when resetting one shortcut to none, because slots
without a selection were not accounted for in the snapshot restorer.

Flag: NONE
Bug: 305156127
Test: Manually verified different shortcut reset scenarios
Change-Id: Ib218dccf3f34509b64662eafdf7ac301884ec066

4 files changed, 17 insertions, 5 deletions

diff --git a/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractor.kt b/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractor.kt
index 60801943..3eca6241 100644
--- a/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractor.kt
+++ b/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractor.kt
@@ -64,7 +64,7 @@ class KeyguardQuickAffordancePickerInteractor(
     }
 
     /** Unselects all affordances from the slot with the given ID. */
-    suspend fun unselectAll(slotId: String) {
+    suspend fun unselectAllFromSlot(slotId: String) {
         client.deleteAllSelections(
             slotId = slotId,
         )
@@ -72,6 +72,11 @@ class KeyguardQuickAffordancePickerInteractor(
         snapshotRestorer.get().storeSnapshot()
     }
 
+    /** Unselects all affordances from all slots. */
+    suspend fun unselectAll() {
+        client.querySlots().forEach { client.deleteAllSelections(it.id) }
+    }
+
     /** Returns a [Drawable] for the given resource ID, from the system UI package. */
     suspend fun getAffordanceIcon(
         @DrawableRes iconResourceId: Int,
diff --git a/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordanceSnapshotRestorer.kt b/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordanceSnapshotRestorer.kt
index 3c7928ce..fee0cb51 100644
--- a/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordanceSnapshotRestorer.kt
+++ b/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordanceSnapshotRestorer.kt
@@ -42,9 +42,14 @@ class KeyguardQuickAffordanceSnapshotRestorer(
     }
 
     override suspend fun restoreToSnapshot(snapshot: RestorableSnapshot) {
+        // reset all current selections
+        interactor.unselectAll()
+
+        val allSelections = checkNotNull(snapshot.args[KEY_SELECTIONS])
+        if (allSelections.isEmpty()) return
+
         val selections: List<Pair<String, String>> =
-            checkNotNull(snapshot.args[KEY_SELECTIONS]).split(SELECTION_SEPARATOR).map { selection
-                ->
+            allSelections.split(SELECTION_SEPARATOR).map { selection ->
                 val (slotId, affordanceId) = selection.split(SLOT_AFFORDANCE_SEPARATOR)
                 slotId to affordanceId
             }
diff --git a/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt b/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt
index 42aced9b..260c0d3b 100644
--- a/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt
+++ b/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt
@@ -217,7 +217,9 @@ private constructor(
                             if (!isSelected) {
                                 {
                                     viewModelScope.launch {
-                                        quickAffordanceInteractor.unselectAll(selectedSlotId)
+                                        quickAffordanceInteractor.unselectAllFromSlot(
+                                            selectedSlotId
+                                        )
                                         logger.logShortcutApplied(
                                             shortcut = "none",
                                             shortcutSlotId = selectedSlotId,
diff --git a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
index bf53f61f..4b4790ad 100644
--- a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
@@ -132,7 +132,7 @@ class KeyguardQuickAffordancePickerInteractorTest {
                 affordanceId = FakeCustomizationProviderClient.AFFORDANCE_3,
             )
 
-            underTest.unselectAll(
+            underTest.unselectAllFromSlot(
                 slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END,
             )