diff --git a/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt
index f20bc2e..ba6fc8b 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt
@@ -50,7 +50,7 @@
         homeScreenCustomizationOptionEntries: List<Pair<CustomizationOption, View>>,
         customizationOptionFloatingSheetViewMap: Map<CustomizationOption, View>?,
         viewModel: CustomizationOptionsViewModel,
-        lifecycleOwner: LifecycleOwner
+        lifecycleOwner: LifecycleOwner,
     ) {
         defaultCustomizationOptionsBinder.bind(
             view,
@@ -104,7 +104,7 @@
                 }
 
                 launch {
-                    viewModel.keyguardQuickAffordanceSummery.collect { summary ->
+                    viewModel.keyguardQuickAffordancePickerViewModel2.summary.collect { summary ->
                         optionShortcutDescription?.let {
                             TextViewBinder.bind(
                                 view = it,
diff --git a/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
new file mode 100644
index 0000000..6bddf4a
--- /dev/null
+++ b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
@@ -0,0 +1,66 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wallpaper.customization.ui.binder
+
+import android.widget.Button
+import android.widget.FrameLayout
+import android.widget.Toolbar
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.LifecycleOwner
+import androidx.lifecycle.lifecycleScope
+import androidx.lifecycle.repeatOnLifecycle
+import com.android.wallpaper.customization.ui.viewmodel.ThemePickerCustomizationOptionsViewModel
+import com.android.wallpaper.picker.customization.ui.binder.DefaultToolbarBinder
+import com.android.wallpaper.picker.customization.ui.binder.ToolbarBinder
+import com.android.wallpaper.picker.customization.ui.viewmodel.CustomizationOptionsViewModel
+import javax.inject.Inject
+import javax.inject.Singleton
+import kotlinx.coroutines.launch
+
+@Singleton
+class ThemePickerToolbarBinder
+@Inject
+constructor(private val defaultToolbarBinder: DefaultToolbarBinder) : ToolbarBinder {
+
+    override fun bind(
+        navButton: FrameLayout,
+        toolbar: Toolbar,
+        applyButton: Button,
+        viewModel: CustomizationOptionsViewModel,
+        lifecycleOwner: LifecycleOwner,
+    ) {
+        defaultToolbarBinder.bind(navButton, toolbar, applyButton, viewModel, lifecycleOwner)
+
+        if (viewModel !is ThemePickerCustomizationOptionsViewModel) {
+            throw IllegalArgumentException(
+                "viewModel $viewModel is not a ThemePickerCustomizationOptionsViewModel."
+            )
+        }
+
+        lifecycleOwner.lifecycleScope.launch {
+            lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
+                launch {
+                    viewModel.keyguardQuickAffordancePickerViewModel2.onApply.collect { onApply ->
+                        applyButton.setOnClickListener {
+                            onApply?.invoke()?.let { viewModel.deselectOption() }
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt b/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt
index 67beb25..86cdd8a 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt
@@ -25,7 +25,9 @@
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor
 import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordanceSlotViewModel
 import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordanceSummaryViewModel
-import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
+import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END
+import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START
+import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.KEYGUARD_QUICK_AFFORDANCE_ID_NONE
 import com.android.themepicker.R
 import com.android.wallpaper.picker.common.button.ui.viewmodel.ButtonStyle
 import com.android.wallpaper.picker.common.button.ui.viewmodel.ButtonViewModel
@@ -81,60 +83,78 @@
                 started = SharingStarted.WhileSubscribed(),
                 initialValue = "",
             )
+    private val _selectedQuickAffordances = MutableStateFlow<Map<String, String>>(emptyMap())
+    val selectedQuickAffordances: Flow<Map<String, String>> =
+        _selectedQuickAffordances.asStateFlow()
+
+    fun resetPreview() {
+        _selectedQuickAffordances.tryEmit(emptyMap())
+        _selectedSlotId.tryEmit(SLOT_ID_BOTTOM_START)
+    }
 
     /** View-models for each slot, keyed by slot ID. */
-    private val slots: Flow<Map<String, KeyguardQuickAffordanceSlotViewModel>> =
+    private val slots: StateFlow<Map<String, KeyguardQuickAffordanceSlotViewModel>> =
         combine(
-            quickAffordanceInteractor.slots,
-            quickAffordanceInteractor.affordances,
-            quickAffordanceInteractor.selections,
-            selectedSlotId,
-        ) { slots, affordances, selections, selectedSlotId ->
-            slots.associate { slot ->
-                val selectedAffordanceIds =
-                    selections
-                        .filter { selection -> selection.slotId == slot.id }
-                        .map { selection -> selection.affordanceId }
-                        .toSet()
-                val selectedAffordances =
-                    affordances.filter { affordance ->
-                        selectedAffordanceIds.contains(affordance.id)
-                    }
-                val isSelected = selectedSlotId == slot.id
-                slot.id to
-                    KeyguardQuickAffordanceSlotViewModel(
-                        name = getSlotName(slot.id),
-                        isSelected = isSelected,
-                        selectedQuickAffordances =
-                            selectedAffordances.map { affordanceModel ->
-                                OptionItemViewModel<Icon>(
-                                    key =
-                                        MutableStateFlow("${slot.id}::${affordanceModel.id}")
-                                            as StateFlow<String>,
-                                    payload =
-                                        Icon.Loaded(
-                                            drawable =
-                                                getAffordanceIcon(affordanceModel.iconResourceId),
-                                            contentDescription =
-                                                Text.Loaded(getSlotContentDescription(slot.id)),
-                                        ),
-                                    text = Text.Loaded(affordanceModel.name),
-                                    isSelected = MutableStateFlow(true) as StateFlow<Boolean>,
-                                    onClicked = flowOf(null),
-                                    onLongClicked = null,
-                                    isEnabled = true,
-                                )
-                            },
-                        maxSelectedQuickAffordances = slot.maxSelectedQuickAffordances,
-                        onClicked =
-                            if (isSelected) {
-                                null
-                            } else {
-                                { _selectedSlotId.tryEmit(slot.id) }
-                            },
-                    )
+                quickAffordanceInteractor.slots,
+                quickAffordanceInteractor.affordances,
+                quickAffordanceInteractor.selections,
+                selectedQuickAffordances,
+                selectedSlotId,
+            ) { slots, affordances, selections, selectedQuickAffordances, selectedSlotId ->
+                slots.associate { slot ->
+                    val selectedAffordanceIds =
+                        selectedQuickAffordances[slot.id]?.let { setOf(it) }
+                            ?: selections
+                                .filter { selection -> selection.slotId == slot.id }
+                                .map { selection -> selection.affordanceId }
+                                .toSet()
+                    val selectedAffordances =
+                        affordances.filter { affordance ->
+                            selectedAffordanceIds.contains(affordance.id)
+                        }
+
+                    val isSelected = selectedSlotId == slot.id
+                    slot.id to
+                        KeyguardQuickAffordanceSlotViewModel(
+                            name = getSlotName(slot.id),
+                            isSelected = isSelected,
+                            selectedQuickAffordances =
+                                selectedAffordances.map { affordanceModel ->
+                                    OptionItemViewModel<Icon>(
+                                        key =
+                                            MutableStateFlow("${slot.id}::${affordanceModel.id}")
+                                                as StateFlow<String>,
+                                        payload =
+                                            Icon.Loaded(
+                                                drawable =
+                                                    getAffordanceIcon(
+                                                        affordanceModel.iconResourceId
+                                                    ),
+                                                contentDescription =
+                                                    Text.Loaded(getSlotContentDescription(slot.id)),
+                                            ),
+                                        text = Text.Loaded(affordanceModel.name),
+                                        isSelected = MutableStateFlow(true) as StateFlow<Boolean>,
+                                        onClicked = flowOf(null),
+                                        onLongClicked = null,
+                                        isEnabled = true,
+                                    )
+                                },
+                            maxSelectedQuickAffordances = slot.maxSelectedQuickAffordances,
+                            onClicked =
+                                if (isSelected) {
+                                    null
+                                } else {
+                                    { _selectedSlotId.tryEmit(slot.id) }
+                                },
+                        )
+                }
             }
-        }
+            .stateIn(
+                scope = viewModelScope,
+                started = SharingStarted.WhileSubscribed(),
+                initialValue = emptyMap(),
+            )
 
     val tabs: Flow<List<FloatingToolbarTabViewModel>> =
         slots.map { slotById ->
@@ -171,7 +191,17 @@
     /** The list of all available quick affordances for the selected slot. */
     val quickAffordances: Flow<List<OptionItemViewModel<Icon>>> =
         quickAffordanceInteractor.affordances.map { affordances ->
-            val isNoneSelected = selectedAffordanceIds.map { it.isEmpty() }.stateIn(viewModelScope)
+            val isNoneSelected =
+                combine(
+                        selectedSlotId,
+                        selectedQuickAffordances,
+                        selectedAffordanceIds,
+                    ) { selectedSlotId, selectedQuickAffordances, selectedAffordanceIds ->
+                        selectedQuickAffordances[selectedSlotId]?.let {
+                            it == KEYGUARD_QUICK_AFFORDANCE_ID_NONE
+                        } ?: selectedAffordanceIds.isEmpty()
+                    }
+                    .stateIn(viewModelScope)
             listOf(
                 none(
                     slotId = selectedSlotId,
@@ -183,15 +213,11 @@
                         ) { isSelected, selectedSlotId ->
                             if (!isSelected) {
                                 {
-                                    viewModelScope.launch {
-                                        quickAffordanceInteractor.unselectAllFromSlot(
-                                            selectedSlotId
-                                        )
-                                        logger.logShortcutApplied(
-                                            shortcut = "none",
-                                            shortcutSlotId = selectedSlotId,
-                                        )
-                                    }
+                                    val newMap =
+                                        _selectedQuickAffordances.value.toMutableMap().apply {
+                                            put(selectedSlotId, KEYGUARD_QUICK_AFFORDANCE_ID_NONE)
+                                        }
+                                    _selectedQuickAffordances.tryEmit(newMap)
                                 }
                             } else {
                                 null
@@ -202,8 +228,15 @@
                 affordances.map { affordance ->
                     val affordanceIcon = getAffordanceIcon(affordance.iconResourceId)
                     val isSelectedFlow: StateFlow<Boolean> =
-                        selectedAffordanceIds
-                            .map { it.contains(affordance.id) }
+                        combine(
+                                selectedSlotId,
+                                selectedQuickAffordances,
+                                selectedAffordanceIds,
+                            ) { selectedSlotId, selectedQuickAffordances, selectedAffordanceIds ->
+                                selectedQuickAffordances[selectedSlotId]?.let {
+                                    it == affordance.id
+                                } ?: selectedAffordanceIds.contains(affordance.id)
+                            }
                             .stateIn(viewModelScope)
                     OptionItemViewModel<Icon>(
                         key =
@@ -221,16 +254,11 @@
                                 ) { isSelected, selectedSlotId ->
                                     if (!isSelected) {
                                         {
-                                            viewModelScope.launch {
-                                                quickAffordanceInteractor.select(
-                                                    slotId = selectedSlotId,
-                                                    affordanceId = affordance.id,
-                                                )
-                                                logger.logShortcutApplied(
-                                                    shortcut = affordance.id,
-                                                    shortcutSlotId = selectedSlotId,
-                                                )
-                                            }
+                                            val newMap =
+                                                _selectedQuickAffordances.value
+                                                    .toMutableMap()
+                                                    .apply { put(selectedSlotId, affordance.id) }
+                                            _selectedQuickAffordances.tryEmit(newMap)
                                         }
                                     } else {
                                         null
@@ -258,6 +286,34 @@
                 }
         }
 
+    val onApply: Flow<(() -> Unit)?> =
+        selectedQuickAffordances.map {
+            if (it.isEmpty()) {
+                null
+            } else {
+                {
+                    it.forEach { entry ->
+                        val slotId = entry.key
+                        val affordanceId = entry.value
+                        viewModelScope.launch {
+                            if (slotId == KEYGUARD_QUICK_AFFORDANCE_ID_NONE) {
+                                quickAffordanceInteractor.unselectAllFromSlot(slotId)
+                            } else {
+                                quickAffordanceInteractor.select(
+                                    slotId = slotId,
+                                    affordanceId = affordanceId
+                                )
+                            }
+                            logger.logShortcutApplied(
+                                shortcut = affordanceId,
+                                shortcutSlotId = slotId,
+                            )
+                        }
+                    }
+                }
+            }
+        }
+
     private val _dialog = MutableStateFlow<DialogViewModel?>(null)
     /**
      * The current dialog to show. If `null`, no dialog should be shown.
@@ -365,10 +421,8 @@
     private fun getSlotName(slotId: String): String {
         return applicationContext.getString(
             when (slotId) {
-                KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START ->
-                    R.string.keyguard_slot_name_bottom_start
-                KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END ->
-                    R.string.keyguard_slot_name_bottom_end
+                SLOT_ID_BOTTOM_START -> R.string.keyguard_slot_name_bottom_start
+                SLOT_ID_BOTTOM_END -> R.string.keyguard_slot_name_bottom_end
                 else -> error("No name for slot with ID of \"$slotId\"!")
             }
         )
@@ -377,10 +431,8 @@
     private fun getSlotContentDescription(slotId: String): String {
         return applicationContext.getString(
             when (slotId) {
-                KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START ->
-                    R.string.keyguard_slot_name_bottom_start
-                KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END ->
-                    R.string.keyguard_slot_name_bottom_end
+                SLOT_ID_BOTTOM_START -> R.string.keyguard_slot_name_bottom_start
+                SLOT_ID_BOTTOM_END -> R.string.keyguard_slot_name_bottom_end
                 else -> error("No accessibility label for slot with ID \"$slotId\"!")
             }
         )
@@ -393,15 +445,9 @@
     val summary: Flow<KeyguardQuickAffordanceSummaryViewModel> =
         slots.map { slots ->
             val icon2 =
-                (slots[KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END]
-                        ?.selectedQuickAffordances
-                        ?.firstOrNull())
-                    ?.payload
+                (slots[SLOT_ID_BOTTOM_END]?.selectedQuickAffordances?.firstOrNull())?.payload
             val icon1 =
-                (slots[KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START]
-                        ?.selectedQuickAffordances
-                        ?.firstOrNull())
-                    ?.payload
+                (slots[SLOT_ID_BOTTOM_START]?.selectedQuickAffordances?.firstOrNull())?.payload
 
             KeyguardQuickAffordanceSummaryViewModel(
                 description = toDescriptionText(applicationContext, slots),
@@ -424,15 +470,9 @@
         slots: Map<String, KeyguardQuickAffordanceSlotViewModel>,
     ): Text {
         val bottomStartAffordanceName =
-            slots[KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START]
-                ?.selectedQuickAffordances
-                ?.firstOrNull()
-                ?.text
+            slots[SLOT_ID_BOTTOM_START]?.selectedQuickAffordances?.firstOrNull()?.text
         val bottomEndAffordanceName =
-            slots[KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END]
-                ?.selectedQuickAffordances
-                ?.firstOrNull()
-                ?.text
+            slots[SLOT_ID_BOTTOM_END]?.selectedQuickAffordances?.firstOrNull()?.text
 
         return when {
             bottomStartAffordanceName != null && bottomEndAffordanceName != null -> {
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
index 79671d6..01c1c80 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
@@ -48,7 +48,10 @@
 
     override val selectedOption = defaultCustomizationOptionsViewModel.selectedOption
 
-    override fun deselectOption(): Boolean = defaultCustomizationOptionsViewModel.deselectOption()
+    override fun deselectOption(): Boolean {
+        keyguardQuickAffordancePickerViewModel2.resetPreview()
+        return defaultCustomizationOptionsViewModel.deselectOption()
+    }
 
     val onCustomizeClockClicked: Flow<(() -> Unit)?> =
         selectedOption.map {
@@ -77,8 +80,6 @@
             }
         }
 
-    val keyguardQuickAffordanceSummery = keyguardQuickAffordancePickerViewModel2.summary
-
     val onCustomizeColorsClicked: Flow<(() -> Unit)?> =
         selectedOption.map {
             if (it == null) {
diff --git a/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt b/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
new file mode 100644
index 0000000..a80febd
--- /dev/null
+++ b/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
@@ -0,0 +1,124 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wallpaper.picker.common.preview.ui.binder
+
+import android.os.Bundle
+import android.os.Message
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.LifecycleOwner
+import androidx.lifecycle.lifecycleScope
+import androidx.lifecycle.repeatOnLifecycle
+import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END
+import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START
+import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.KEY_INITIALLY_SELECTED_SLOT_ID
+import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.KEY_QUICK_AFFORDANCE_ID
+import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.KEY_SLOT_ID
+import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.MESSAGE_ID_DEFAULT_PREVIEW
+import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.MESSAGE_ID_PREVIEW_QUICK_AFFORDANCE_SELECTED
+import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.MESSAGE_ID_SLOT_SELECTED
+import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.MESSAGE_ID_START_CUSTOMIZING_QUICK_AFFORDANCES
+import com.android.wallpaper.customization.ui.util.ThemePickerCustomizationOptionUtil.ThemePickerLockCustomizationOption
+import com.android.wallpaper.customization.ui.viewmodel.ThemePickerCustomizationOptionsViewModel
+import com.android.wallpaper.picker.common.preview.ui.binder.WorkspaceCallbackBinder.Companion.sendMessage
+import com.android.wallpaper.picker.customization.ui.viewmodel.CustomizationOptionsViewModel
+import javax.inject.Inject
+import javax.inject.Singleton
+import kotlinx.coroutines.launch
+
+@Singleton
+class ThemePickerWorkspaceCallbackBinder
+@Inject
+constructor(private val defaultWorkspaceCallbackBinder: DefaultWorkspaceCallbackBinder) :
+    WorkspaceCallbackBinder {
+
+    override fun bind(
+        workspaceCallback: Message,
+        viewModel: CustomizationOptionsViewModel,
+        lifecycleOwner: LifecycleOwner,
+    ) {
+        defaultWorkspaceCallbackBinder.bind(
+            workspaceCallback = workspaceCallback,
+            viewModel = viewModel,
+            lifecycleOwner = lifecycleOwner,
+        )
+
+        if (viewModel !is ThemePickerCustomizationOptionsViewModel) {
+            throw IllegalArgumentException(
+                "viewModel $viewModel is not a ThemePickerCustomizationOptionsViewModel."
+            )
+        }
+
+        lifecycleOwner.lifecycleScope.launch {
+            lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
+                launch {
+                    viewModel.selectedOption.collect {
+                        when (it) {
+                            ThemePickerLockCustomizationOption.SHORTCUTS ->
+                                workspaceCallback.sendMessage(
+                                    MESSAGE_ID_START_CUSTOMIZING_QUICK_AFFORDANCES,
+                                    Bundle().apply {
+                                        putString(
+                                            KEY_INITIALLY_SELECTED_SLOT_ID,
+                                            SLOT_ID_BOTTOM_START,
+                                        )
+                                    }
+                                )
+                            else ->
+                                workspaceCallback.sendMessage(
+                                    MESSAGE_ID_DEFAULT_PREVIEW,
+                                    Bundle.EMPTY,
+                                )
+                        }
+                    }
+                }
+
+                launch {
+                    viewModel.keyguardQuickAffordancePickerViewModel2.selectedSlotId.collect {
+                        workspaceCallback.sendMessage(
+                            MESSAGE_ID_SLOT_SELECTED,
+                            Bundle().apply { putString(KEY_SLOT_ID, it) },
+                        )
+                    }
+                }
+
+                launch {
+                    viewModel.keyguardQuickAffordancePickerViewModel2.selectedQuickAffordances
+                        .collect {
+                            it[SLOT_ID_BOTTOM_START]?.let {
+                                workspaceCallback.sendMessage(
+                                    MESSAGE_ID_PREVIEW_QUICK_AFFORDANCE_SELECTED,
+                                    Bundle().apply {
+                                        putString(KEY_SLOT_ID, SLOT_ID_BOTTOM_START)
+                                        putString(KEY_QUICK_AFFORDANCE_ID, it)
+                                    },
+                                )
+                            }
+                            it[SLOT_ID_BOTTOM_END]?.let {
+                                workspaceCallback.sendMessage(
+                                    MESSAGE_ID_PREVIEW_QUICK_AFFORDANCE_SELECTED,
+                                    Bundle().apply {
+                                        putString(KEY_SLOT_ID, SLOT_ID_BOTTOM_END)
+                                        putString(KEY_QUICK_AFFORDANCE_ID, it)
+                                    },
+                                )
+                            }
+                        }
+                }
+            }
+        }
+    }
+}
diff --git a/src_override/com/android/wallpaper/modules/ThemePickerAppModule.kt b/src_override/com/android/wallpaper/modules/ThemePickerAppModule.kt
index 8cda50e..4426b49 100644
--- a/src_override/com/android/wallpaper/modules/ThemePickerAppModule.kt
+++ b/src_override/com/android/wallpaper/modules/ThemePickerAppModule.kt
@@ -36,6 +36,7 @@
 import com.android.systemui.shared.settings.data.repository.SystemSettingsRepository
 import com.android.systemui.shared.settings.data.repository.SystemSettingsRepositoryImpl
 import com.android.wallpaper.customization.ui.binder.ThemePickerCustomizationOptionsBinder
+import com.android.wallpaper.customization.ui.binder.ThemePickerToolbarBinder
 import com.android.wallpaper.effects.DefaultEffectsController
 import com.android.wallpaper.effects.EffectsController
 import com.android.wallpaper.module.DefaultPartnerProvider
@@ -50,7 +51,10 @@
 import com.android.wallpaper.picker.category.domain.interactor.implementations.DefaultCategoriesLoadingStatusInteractor
 import com.android.wallpaper.picker.category.ui.view.providers.IndividualPickerFactory
 import com.android.wallpaper.picker.category.ui.view.providers.implementation.DefaultIndividualPickerFactory
+import com.android.wallpaper.picker.common.preview.ui.binder.ThemePickerWorkspaceCallbackBinder
+import com.android.wallpaper.picker.common.preview.ui.binder.WorkspaceCallbackBinder
 import com.android.wallpaper.picker.customization.ui.binder.CustomizationOptionsBinder
+import com.android.wallpaper.picker.customization.ui.binder.ToolbarBinder
 import com.android.wallpaper.picker.di.modules.BackgroundDispatcher
 import com.android.wallpaper.picker.di.modules.MainDispatcher
 import com.android.wallpaper.picker.preview.ui.util.DefaultImageEffectDialogUtil
@@ -129,6 +133,8 @@
     @Singleton
     abstract fun bindThemesUserEventLogger(impl: ThemesUserEventLoggerImpl): ThemesUserEventLogger
 
+    @Binds @Singleton abstract fun bindToolbarBinder(impl: ThemePickerToolbarBinder): ToolbarBinder
+
     @Binds
     @Singleton
     abstract fun bindUserEventLogger(impl: ThemesUserEventLoggerImpl): UserEventLogger
@@ -145,6 +151,12 @@
         impl: DefaultCustomizationPreferences
     ): WallpaperPreferences
 
+    @Binds
+    @Singleton
+    abstract fun bindWorkspaceCallbackBinder(
+        impl: ThemePickerWorkspaceCallbackBinder
+    ): WorkspaceCallbackBinder
+
     companion object {
 
         @Provides
diff --git a/tests/common/src/com/android/customization/module/logging/TestThemesUserEventLogger.kt b/tests/common/src/com/android/customization/module/logging/TestThemesUserEventLogger.kt
index 8e9dacd..4651067 100644
--- a/tests/common/src/com/android/customization/module/logging/TestThemesUserEventLogger.kt
+++ b/tests/common/src/com/android/customization/module/logging/TestThemesUserEventLogger.kt
@@ -31,11 +31,15 @@
     @ColorSource
     var themeColorSource: Int = StyleEnums.COLOR_SOURCE_UNSPECIFIED
         private set
+
     var themeColorStyle: Int = -1
         private set
+
     var themeSeedColor: Int = -1
         private set
 
+    var shortcutLogs: List<Pair<String, String>> = emptyList()
+
     override fun logThemeColorApplied(@ColorSource source: Int, style: Int, seedColor: Int) {
         this.themeColorSource = source
         this.themeColorStyle = style
@@ -56,7 +60,9 @@
 
     override fun logLockScreenNotificationApplied(showLockScreenNotifications: Boolean) {}
 
-    override fun logShortcutApplied(shortcut: String, shortcutSlotId: String) {}
+    override fun logShortcutApplied(shortcut: String, shortcutSlotId: String) {
+        shortcutLogs = shortcutLogs.toMutableList().apply { add(shortcut to shortcutSlotId) }
+    }
 
     override fun logDarkThemeApplied(useDarkTheme: Boolean) {}
 
diff --git a/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt b/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt
index 3c46cb7..ac49994 100644
--- a/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt
+++ b/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt
@@ -35,6 +35,7 @@
 import com.android.systemui.shared.customization.data.content.CustomizationProviderClientImpl
 import com.android.systemui.shared.settings.data.repository.SecureSettingsRepository
 import com.android.systemui.shared.settings.data.repository.SecureSettingsRepositoryImpl
+import com.android.wallpaper.customization.ui.binder.ThemePickerToolbarBinder
 import com.android.wallpaper.effects.EffectsController
 import com.android.wallpaper.effects.FakeEffectsController
 import com.android.wallpaper.module.Injector
@@ -44,8 +45,11 @@
 import com.android.wallpaper.module.logging.UserEventLogger
 import com.android.wallpaper.modules.ThemePickerAppModule
 import com.android.wallpaper.network.Requester
+import com.android.wallpaper.picker.common.preview.ui.binder.ThemePickerWorkspaceCallbackBinder
+import com.android.wallpaper.picker.common.preview.ui.binder.WorkspaceCallbackBinder
 import com.android.wallpaper.picker.customization.ui.binder.CustomizationOptionsBinder
 import com.android.wallpaper.picker.customization.ui.binder.DefaultCustomizationOptionsBinder
+import com.android.wallpaper.picker.customization.ui.binder.ToolbarBinder
 import com.android.wallpaper.picker.di.modules.BackgroundDispatcher
 import com.android.wallpaper.picker.di.modules.MainDispatcher
 import com.android.wallpaper.picker.preview.ui.util.DefaultImageEffectDialogUtil
@@ -114,6 +118,8 @@
     @Singleton
     abstract fun bindThemesUserEventLogger(impl: TestThemesUserEventLogger): ThemesUserEventLogger
 
+    @Binds @Singleton abstract fun bindToolbarBinder(impl: ThemePickerToolbarBinder): ToolbarBinder
+
     @Binds @Singleton abstract fun bindUserEventLogger(impl: TestUserEventLogger): UserEventLogger
 
     @Binds
@@ -128,6 +134,12 @@
         impl: TestDefaultCustomizationPreferences
     ): WallpaperPreferences
 
+    @Binds
+    @Singleton
+    abstract fun bindWorkspaceCallbackBinder(
+        impl: ThemePickerWorkspaceCallbackBinder
+    ): WorkspaceCallbackBinder
+
     companion object {
 
         @Provides
diff --git a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2Test.kt b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2Test.kt
new file mode 100644
index 0000000..19d74f6
--- /dev/null
+++ b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2Test.kt
@@ -0,0 +1,418 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wallpaper.customization.ui.viewmodel
+
+import android.content.Context
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.filters.SmallTest
+import com.android.customization.module.logging.TestThemesUserEventLogger
+import com.android.customization.picker.quickaffordance.data.repository.KeyguardQuickAffordancePickerRepository
+import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor
+import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordanceSnapshotRestorer
+import com.android.systemui.shared.customization.data.content.CustomizationProviderClient
+import com.android.systemui.shared.customization.data.content.FakeCustomizationProviderClient
+import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
+import com.android.themepicker.R
+import com.android.wallpaper.picker.common.icon.ui.viewmodel.Icon
+import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
+import com.android.wallpaper.picker.customization.ui.viewmodel.FloatingToolbarTabViewModel
+import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.emptyFlow
+import kotlinx.coroutines.test.TestDispatcher
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.resetMain
+import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.test.setMain
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(RobolectricTestRunner::class)
+class KeyguardQuickAffordancePickerViewModel2Test {
+
+    private val logger = TestThemesUserEventLogger()
+
+    private lateinit var underTest: KeyguardQuickAffordancePickerViewModel2
+
+    private lateinit var context: Context
+    private lateinit var testDispatcher: TestDispatcher
+    private lateinit var testScope: TestScope
+    private lateinit var client: FakeCustomizationProviderClient
+
+    @Before
+    fun setUp() {
+        context = ApplicationProvider.getApplicationContext()
+        testDispatcher = UnconfinedTestDispatcher()
+        Dispatchers.setMain(testDispatcher)
+        testScope = TestScope(testDispatcher)
+        client = FakeCustomizationProviderClient()
+        val quickAffordanceInteractor =
+            KeyguardQuickAffordancePickerInteractor(
+                repository =
+                    KeyguardQuickAffordancePickerRepository(
+                        client = client,
+                        mainScope = testScope.backgroundScope,
+                    ),
+                client = client,
+                snapshotRestorer = KeyguardQuickAffordanceSnapshotRestorer(client),
+            )
+        underTest =
+            KeyguardQuickAffordancePickerViewModel2(
+                applicationContext = context,
+                quickAffordanceInteractor = quickAffordanceInteractor,
+                logger = logger,
+                viewModelScope = testScope.backgroundScope,
+            )
+    }
+
+    @After
+    fun tearDown() {
+        Dispatchers.resetMain()
+    }
+
+    @Test
+    fun selectedSlotIdUpdates_whenClickingOnTabsAndCallingResetPreview() =
+        testScope.runTest {
+            val selectedSlotId = collectLastValue(underTest.selectedSlotId)
+
+            val tabs = collectLastValue(underTest.tabs)
+
+            // Default selected slot ID is bottom_start
+            assertThat(selectedSlotId())
+                .isEqualTo(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START)
+
+            // Click on tab1
+            val tab1 = tabs()?.get(1) ?: throw NullPointerException("secondTab should not be null.")
+            tab1.onClick?.invoke()
+            assertThat(selectedSlotId()).isEqualTo(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END)
+
+            underTest.resetPreview()
+            assertThat(selectedSlotId())
+                .isEqualTo(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START)
+        }
+
+    @Test
+    fun selectedQuickAffordancesMapUpdates_whenClickingOnQuickAffordanceOptionsAndCallingResetPreview() =
+        testScope.runTest {
+            val selectedQuickAffordances = collectLastValue(underTest.selectedQuickAffordances)
+
+            val tabs = collectLastValue(underTest.tabs)
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
+
+            // Default selectedQuickAffordances is an empty map
+            assertThat(selectedQuickAffordances()).isEqualTo(emptyMap<String, String>())
+
+            // Click on quick affordance 1 when selected slot ID is bottom_start
+            val onClickAffordance1 =
+                collectLastValue(quickAffordances()?.get(1)?.onClicked ?: emptyFlow())
+            onClickAffordance1()?.invoke()
+            assertThat(selectedQuickAffordances())
+                .isEqualTo(
+                    mapOf(
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
+                            FakeCustomizationProviderClient.AFFORDANCE_1
+                    )
+                )
+
+            // Click on tab 1 to change the selected slot ID to bottom_end and click on quick
+            // affordance 2
+            tabs()?.get(1)?.onClick?.invoke()
+            val onClickAffordance2 =
+                collectLastValue(quickAffordances()?.get(2)?.onClicked ?: emptyFlow())
+            onClickAffordance2()?.invoke()
+            assertThat(selectedQuickAffordances())
+                .isEqualTo(
+                    mapOf(
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
+                            FakeCustomizationProviderClient.AFFORDANCE_1,
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END to
+                            FakeCustomizationProviderClient.AFFORDANCE_2
+                    )
+                )
+
+            underTest.resetPreview()
+            assertThat(selectedQuickAffordances()).isEqualTo(emptyMap<String, String>())
+        }
+
+    @Test
+    fun tabsUpdates_whenClickingOnTabsAndQuickAffordanceOptions() =
+        testScope.runTest {
+            val tabs = collectLastValue(underTest.tabs)
+
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
+
+            // Default state of the 2 tabs
+            assertTabUiState(
+                tab = tabs()?.get(0),
+                icon = Icon.Resource(R.drawable.link_off, null),
+                text = "Left shortcut",
+                isSelected = true,
+            )
+            assertTabUiState(
+                tab = tabs()?.get(1),
+                icon = Icon.Resource(R.drawable.link_off, null),
+                text = "Right shortcut",
+                isSelected = false,
+            )
+
+            // Click on tab 1
+            tabs()?.get(1)?.onClick?.invoke()
+            assertTabUiState(
+                tab = tabs()?.get(0),
+                icon = Icon.Resource(R.drawable.link_off, null),
+                text = "Left shortcut",
+                isSelected = false,
+            )
+            val tab1 = tabs()?.get(1)
+            assertTabUiState(
+                tab = tab1,
+                icon = Icon.Resource(R.drawable.link_off, null),
+                text = "Right shortcut",
+                isSelected = true,
+            )
+
+            // Click on quick affordance 1 when tab 1 is selected. Icon should change
+            val clickOnQuickAffordance1 =
+                collectLastValue(quickAffordances()?.get(1)?.onClicked ?: emptyFlow())
+            clickOnQuickAffordance1()?.invoke()
+            assertTabUiState(
+                tab = tabs()?.get(1),
+                icon =
+                    Icon.Loaded(
+                        FakeCustomizationProviderClient.ICON_1,
+                        Text.Loaded("Right shortcut")
+                    ),
+                text = "Right shortcut",
+                isSelected = true,
+            )
+        }
+
+    @Test
+    fun quickAffordancesUpdates_whenClickingOnTabsAndQuickAffordanceOptions() =
+        testScope.runTest {
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
+
+            val tabs = collectLastValue(underTest.tabs)
+
+            // The default quickAffordances snapshot
+            assertThat(quickAffordances()?.size).isEqualTo(4)
+            assertQuickAffordance(
+                testScope = this,
+                quickAffordance = quickAffordances()?.get(0),
+                key = "${KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START}::none",
+                icon = Icon.Resource(R.drawable.link_off, null),
+                text = Text.Resource(R.string.keyguard_affordance_none),
+                isSelected = true,
+            )
+            assertQuickAffordance(
+                testScope = this,
+                quickAffordance = quickAffordances()?.get(1),
+                key =
+                    "${KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START}::${FakeCustomizationProviderClient.AFFORDANCE_1}",
+                icon = Icon.Loaded(FakeCustomizationProviderClient.ICON_1, null),
+                text = Text.Loaded(FakeCustomizationProviderClient.AFFORDANCE_1),
+                isSelected = false,
+            )
+            assertQuickAffordance(
+                testScope = this,
+                quickAffordance = quickAffordances()?.get(2),
+                key =
+                    "${KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START}::${FakeCustomizationProviderClient.AFFORDANCE_2}",
+                icon = Icon.Loaded(FakeCustomizationProviderClient.ICON_2, null),
+                text = Text.Loaded(FakeCustomizationProviderClient.AFFORDANCE_2),
+                isSelected = false,
+            )
+            assertQuickAffordance(
+                testScope = this,
+                quickAffordance = quickAffordances()?.get(3),
+                key =
+                    "${KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START}::${FakeCustomizationProviderClient.AFFORDANCE_3}",
+                icon = Icon.Loaded(FakeCustomizationProviderClient.ICON_3, null),
+                text = Text.Loaded(FakeCustomizationProviderClient.AFFORDANCE_3),
+                isSelected = false,
+            )
+
+            // Click on quick affordance 2. Quick affordance 0 will be unselected and quick
+            // affordance 2 will be selected.
+            val onClickQuickAffordance2 =
+                collectLastValue(quickAffordances()?.get(2)?.onClicked ?: emptyFlow())
+            onClickQuickAffordance2()?.invoke()
+            assertQuickAffordance(
+                testScope = this,
+                quickAffordance = quickAffordances()?.get(0),
+                key = "${KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START}::none",
+                icon = Icon.Resource(R.drawable.link_off, null),
+                text = Text.Resource(R.string.keyguard_affordance_none),
+                isSelected = false,
+            )
+            assertQuickAffordance(
+                testScope = this,
+                quickAffordance = quickAffordances()?.get(2),
+                key =
+                    "${KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START}::${FakeCustomizationProviderClient.AFFORDANCE_2}",
+                icon = Icon.Loaded(FakeCustomizationProviderClient.ICON_2, null),
+                text = Text.Loaded(FakeCustomizationProviderClient.AFFORDANCE_2),
+                isSelected = true,
+            )
+
+            tabs()?.get(1)?.onClick?.invoke()
+            assertQuickAffordance(
+                testScope = this,
+                quickAffordance = quickAffordances()?.get(0),
+                key = "${KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END}::none",
+                icon = Icon.Resource(R.drawable.link_off, null),
+                text = Text.Resource(R.string.keyguard_affordance_none),
+                isSelected = true,
+            )
+            assertQuickAffordance(
+                testScope = this,
+                quickAffordance = quickAffordances()?.get(1),
+                key =
+                    "${KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END}::${FakeCustomizationProviderClient.AFFORDANCE_1}",
+                icon = Icon.Loaded(FakeCustomizationProviderClient.ICON_1, null),
+                text = Text.Loaded(FakeCustomizationProviderClient.AFFORDANCE_1),
+                isSelected = false,
+            )
+            assertQuickAffordance(
+                testScope = this,
+                quickAffordance = quickAffordances()?.get(2),
+                key =
+                    "${KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END}::${FakeCustomizationProviderClient.AFFORDANCE_2}",
+                icon = Icon.Loaded(FakeCustomizationProviderClient.ICON_2, null),
+                text = Text.Loaded(FakeCustomizationProviderClient.AFFORDANCE_2),
+                isSelected = false,
+            )
+            assertQuickAffordance(
+                testScope = this,
+                quickAffordance = quickAffordances()?.get(3),
+                key =
+                    "${KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END}::${FakeCustomizationProviderClient.AFFORDANCE_3}",
+                icon = Icon.Loaded(FakeCustomizationProviderClient.ICON_3, null),
+                text = Text.Loaded(FakeCustomizationProviderClient.AFFORDANCE_3),
+                isSelected = false,
+            )
+
+            // When tab 1 is selected, click on quick affordance 3. Quick affordance 0 will be
+            // unselected and quick affordance 3 will be selected.
+            val onClickQuickAffordance3 =
+                collectLastValue(quickAffordances()?.get(3)?.onClicked ?: emptyFlow())
+            onClickQuickAffordance3()?.invoke()
+            assertQuickAffordance(
+                testScope = this,
+                quickAffordance = quickAffordances()?.get(0),
+                key = "${KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END}::none",
+                icon = Icon.Resource(R.drawable.link_off, null),
+                text = Text.Resource(R.string.keyguard_affordance_none),
+                isSelected = false,
+            )
+            assertQuickAffordance(
+                testScope = this,
+                quickAffordance = quickAffordances()?.get(3),
+                key =
+                    "${KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END}::${FakeCustomizationProviderClient.AFFORDANCE_3}",
+                icon = Icon.Loaded(FakeCustomizationProviderClient.ICON_3, null),
+                text = Text.Loaded(FakeCustomizationProviderClient.AFFORDANCE_3),
+                isSelected = true,
+            )
+        }
+
+    @Test
+    fun loggerShouldLogAndClientShouldUpdate_whenOnApply() =
+        testScope.runTest {
+            val onApply = collectLastValue(underTest.onApply)
+
+            val tabs = collectLastValue(underTest.tabs)
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
+
+            // Select the preview quick affordances
+            val onClickAffordance1 =
+                collectLastValue(quickAffordances()?.get(1)?.onClicked ?: emptyFlow())
+            onClickAffordance1()?.invoke()
+            tabs()?.get(1)?.onClick?.invoke()
+            val onClickAffordance2 =
+                collectLastValue(quickAffordances()?.get(2)?.onClicked ?: emptyFlow())
+            onClickAffordance2()?.invoke()
+
+            onApply()?.invoke()
+            assertThat(client.querySelections())
+                .isEqualTo(
+                    listOf(
+                        CustomizationProviderClient.Selection(
+                            slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START,
+                            affordanceId = FakeCustomizationProviderClient.AFFORDANCE_1,
+                            affordanceName = FakeCustomizationProviderClient.AFFORDANCE_1,
+                        ),
+                        CustomizationProviderClient.Selection(
+                            slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END,
+                            affordanceId = FakeCustomizationProviderClient.AFFORDANCE_2,
+                            affordanceName = FakeCustomizationProviderClient.AFFORDANCE_2,
+                        ),
+                    )
+                )
+            assertThat(logger.shortcutLogs)
+                .isEqualTo(
+                    listOf(
+                        FakeCustomizationProviderClient.AFFORDANCE_1 to
+                            KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START,
+                        FakeCustomizationProviderClient.AFFORDANCE_2 to
+                            KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END,
+                    )
+                )
+        }
+
+    private fun assertTabUiState(
+        tab: FloatingToolbarTabViewModel?,
+        icon: Icon?,
+        text: String,
+        isSelected: Boolean,
+    ) {
+        if (tab == null) {
+            throw NullPointerException("tab is null.")
+        }
+        assertThat(tab.icon).isEqualTo(icon)
+        assertThat(tab.text).isEqualTo(text)
+        assertThat(tab.isSelected).isEqualTo(isSelected)
+    }
+
+    private fun assertQuickAffordance(
+        testScope: TestScope,
+        quickAffordance: OptionItemViewModel<Icon>?,
+        key: String,
+        icon: Icon,
+        text: Text,
+        isSelected: Boolean,
+    ) {
+        if (quickAffordance == null) {
+            throw NullPointerException("quickAffordance is null.")
+        }
+        assertThat(testScope.collectLastValue(quickAffordance.key)()).isEqualTo(key)
+        assertThat(quickAffordance.payload).isEqualTo(icon)
+        assertThat(quickAffordance.text).isEqualTo(text)
+        assertThat(quickAffordance.isTextUserVisible).isEqualTo(true)
+        assertThat(testScope.collectLastValue(quickAffordance.isSelected)()).isEqualTo(isSelected)
+        assertThat(quickAffordance.isEnabled).isEqualTo(true)
+    }
+}