diff --git a/src/com/android/customization/model/grid/DefaultShapeGridManager.kt b/src/com/android/customization/model/grid/DefaultShapeGridManager.kt
index f0d9efa..8de5ff2 100644
--- a/src/com/android/customization/model/grid/DefaultShapeGridManager.kt
+++ b/src/com/android/customization/model/grid/DefaultShapeGridManager.kt
@@ -39,91 +39,56 @@
         context.getString(R.string.grid_control_metadata_name)
     private val previewUtils: PreviewUtils = PreviewUtils(context, authorityMetadataKey)
 
+    override suspend fun isGridOptionAvailable(): Boolean {
+        return previewUtils.supportsPreview() && (getGridOptions()?.size ?: 0) > 1
+    }
+
     override suspend fun getGridOptions(): List<GridOptionModel>? =
         withContext(bgDispatcher) {
-            if (previewUtils.supportsPreview()) {
-                context.contentResolver
-                    .query(previewUtils.getUri(GRID_OPTIONS), null, null, null, null)
-                    ?.use { cursor ->
-                        buildList {
-                            while (cursor.moveToNext()) {
-                                val rows = cursor.getInt(cursor.getColumnIndex(COL_ROWS))
-                                val cols = cursor.getInt(cursor.getColumnIndex(COL_COLS))
-                                add(
-                                    GridOptionModel(
-                                        key = cursor.getString(cursor.getColumnIndex(COL_GRID_KEY)),
-                                        title =
-                                            context.getString(
-                                                com.android.themepicker.R.string.grid_title_pattern,
-                                                cols,
-                                                rows,
-                                            ),
-                                        isCurrent =
-                                            cursor
-                                                .getString(cursor.getColumnIndex(COL_IS_DEFAULT))
-                                                .toBoolean(),
-                                        rows = rows,
-                                        cols = cols,
-                                    )
+            context.contentResolver
+                .query(previewUtils.getUri(LIST_OPTIONS), null, null, null, null)
+                ?.use { cursor ->
+                    buildList {
+                        while (cursor.moveToNext()) {
+                            val rows = cursor.getInt(cursor.getColumnIndex(COL_ROWS))
+                            val cols = cursor.getInt(cursor.getColumnIndex(COL_COLS))
+                            add(
+                                GridOptionModel(
+                                    key = cursor.getString(cursor.getColumnIndex(COL_NAME)),
+                                    title =
+                                        context.getString(
+                                            com.android.themepicker.R.string.grid_title_pattern,
+                                            cols,
+                                            rows,
+                                        ),
+                                    isCurrent =
+                                        cursor
+                                            .getString(cursor.getColumnIndex(COL_IS_DEFAULT))
+                                            .toBoolean(),
+                                    rows = rows,
+                                    cols = cols,
                                 )
-                            }
+                            )
                         }
                     }
-            } else {
-                null
-            }
+                }
         }
 
-    override suspend fun getShapeOptions(): List<ShapeOptionModel>? =
-        withContext(bgDispatcher) {
-            if (previewUtils.supportsPreview()) {
-                context.contentResolver
-                    .query(previewUtils.getUri(SHAPE_OPTIONS), null, null, null, null)
-                    ?.use { cursor ->
-                        buildList {
-                            while (cursor.moveToNext()) {
-                                add(
-                                    ShapeOptionModel(
-                                        key =
-                                            cursor.getString(cursor.getColumnIndex(COL_SHAPE_KEY)),
-                                        title = cursor.getString(cursor.getColumnIndex(COL_TITLE)),
-                                        path = cursor.getString(cursor.getColumnIndex(COL_PATH)),
-                                        isCurrent =
-                                            cursor
-                                                .getString(cursor.getColumnIndex(COL_IS_DEFAULT))
-                                                .toBoolean(),
-                                    )
-                                )
-                            }
-                        }
-                    }
-            } else {
-                null
-            }
-        }
-
-    override fun applyShapeGridOption(shapeKey: String, gridKey: String): Int {
+    override fun applyGridOption(gridName: String): Int {
         return context.contentResolver.update(
-            previewUtils.getUri(SHAPE_GRID),
-            ContentValues().apply {
-                put(COL_SHAPE_KEY, shapeKey)
-                put(COL_GRID_KEY, gridKey)
-            },
+            previewUtils.getUri(DEFAULT_GRID),
+            ContentValues().apply { put("name", gridName) },
             null,
             null,
         )
     }
 
     companion object {
-        const val SHAPE_OPTIONS: String = "shape_options"
-        const val GRID_OPTIONS: String = "grid_options"
-        const val SHAPE_GRID: String = "default_grid"
-        const val COL_SHAPE_KEY: String = "shape_key"
-        const val COL_GRID_KEY: String = "name"
+        const val LIST_OPTIONS: String = "list_options"
+        const val DEFAULT_GRID: String = "default_grid"
+        const val COL_NAME: String = "name"
         const val COL_ROWS: String = "rows"
         const val COL_COLS: String = "cols"
         const val COL_IS_DEFAULT: String = "is_default"
-        const val COL_TITLE: String = "title"
-        const val COL_PATH: String = "path"
     }
 }
diff --git a/src/com/android/customization/model/grid/ShapeGridManager.kt b/src/com/android/customization/model/grid/ShapeGridManager.kt
index 0a23346..1603ed9 100644
--- a/src/com/android/customization/model/grid/ShapeGridManager.kt
+++ b/src/com/android/customization/model/grid/ShapeGridManager.kt
@@ -18,9 +18,9 @@
 
 interface ShapeGridManager {
 
+    suspend fun isGridOptionAvailable(): Boolean
+
     suspend fun getGridOptions(): List<GridOptionModel>?
 
-    suspend fun getShapeOptions(): List<ShapeOptionModel>?
-
-    fun applyShapeGridOption(shapeKey: String, gridKey: String): Int
+    fun applyGridOption(gridName: String): Int
 }
diff --git a/src/com/android/customization/model/grid/ShapeOptionModel.kt b/src/com/android/customization/model/grid/ShapeOptionModel.kt
deleted file mode 100644
index c3ed192..0000000
--- a/src/com/android/customization/model/grid/ShapeOptionModel.kt
+++ /dev/null
@@ -1,24 +0,0 @@
-/*
- * 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.customization.model.grid
-
-data class ShapeOptionModel(
-    val key: String,
-    val title: String,
-    val path: String,
-    val isCurrent: Boolean,
-)
diff --git a/src/com/android/customization/picker/grid/data/repository/ShapeGridRepository.kt b/src/com/android/customization/picker/grid/data/repository/ShapeGridRepository.kt
index 86c455e..9ae4d18 100644
--- a/src/com/android/customization/picker/grid/data/repository/ShapeGridRepository.kt
+++ b/src/com/android/customization/picker/grid/data/repository/ShapeGridRepository.kt
@@ -19,7 +19,6 @@
 
 import com.android.customization.model.grid.GridOptionModel
 import com.android.customization.model.grid.ShapeGridManager
-import com.android.customization.model.grid.ShapeOptionModel
 import com.android.wallpaper.picker.di.modules.BackgroundDispatcher
 import javax.inject.Inject
 import javax.inject.Singleton
@@ -42,31 +41,27 @@
     @BackgroundDispatcher private val bgDispatcher: CoroutineDispatcher,
 ) {
 
-    private val _shapeOptions = MutableStateFlow<List<ShapeOptionModel>?>(null)
+    suspend fun isGridOptionAvailable(): Boolean =
+        withContext(bgDispatcher) { manager.isGridOptionAvailable() }
+
     private val _gridOptions = MutableStateFlow<List<GridOptionModel>?>(null)
 
     init {
         bgScope.launch {
-            _gridOptions.value = manager.getGridOptions()
-            _shapeOptions.value = manager.getShapeOptions()
+            val options = manager.getGridOptions()
+            _gridOptions.value = options
         }
     }
 
-    val shapeOptions: StateFlow<List<ShapeOptionModel>?> = _shapeOptions.asStateFlow()
-
-    val selectedShapeOption: Flow<ShapeOptionModel?> =
-        shapeOptions.map { shapeOptions -> shapeOptions?.firstOrNull { it.isCurrent } }
-
     val gridOptions: StateFlow<List<GridOptionModel>?> = _gridOptions.asStateFlow()
 
     val selectedGridOption: Flow<GridOptionModel?> =
         gridOptions.map { gridOptions -> gridOptions?.firstOrNull { it.isCurrent } }
 
-    suspend fun applySelectedOption(shapeKey: String, gridKey: String) =
+    suspend fun applySelectedOption(key: String) =
         withContext(bgDispatcher) {
-            manager.applyShapeGridOption(shapeKey, gridKey)
-            // After applying, we should query and update shape and grid options again.
+            manager.applyGridOption(key)
+            // After applying new grid option, we should query and update the grid options again.
             _gridOptions.value = manager.getGridOptions()
-            _shapeOptions.value = manager.getShapeOptions()
         }
 }
diff --git a/src/com/android/customization/picker/grid/domain/interactor/ShapeGridInteractor.kt b/src/com/android/customization/picker/grid/domain/interactor/ShapeGridInteractor.kt
index 8c4522e..c9c2b0c 100644
--- a/src/com/android/customization/picker/grid/domain/interactor/ShapeGridInteractor.kt
+++ b/src/com/android/customization/picker/grid/domain/interactor/ShapeGridInteractor.kt
@@ -23,15 +23,11 @@
 
 @Singleton
 class ShapeGridInteractor @Inject constructor(private val repository: ShapeGridRepository) {
-
-    val shapeOptions = repository.shapeOptions
-
-    val selectedShapeOption = repository.selectedShapeOption
+    suspend fun isGridOptionAvailable(): Boolean = repository.isGridOptionAvailable()
 
     val gridOptions = repository.gridOptions
 
     val selectedGridOption = repository.selectedGridOption
 
-    suspend fun applySelectedOption(shapeKey: String, gridKey: String) =
-        repository.applySelectedOption(shapeKey, gridKey)
+    suspend fun applySelectedOption(key: String) = repository.applySelectedOption(key)
 }
diff --git a/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
index 91705dc..6282842 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
@@ -43,8 +43,16 @@
         applyButton: Button,
         viewModel: CustomizationOptionsViewModel,
         lifecycleOwner: LifecycleOwner,
+        onNavBack: () -> Unit,
     ) {
-        defaultToolbarBinder.bind(navButton, toolbar, applyButton, viewModel, lifecycleOwner)
+        defaultToolbarBinder.bind(
+            navButton,
+            toolbar,
+            applyButton,
+            viewModel,
+            lifecycleOwner,
+            onNavBack,
+        )
 
         if (viewModel !is ThemePickerCustomizationOptionsViewModel) {
             throw IllegalArgumentException(
@@ -56,7 +64,12 @@
             lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
                 launch {
                     viewModel.onApplyButtonClicked.collect { onApplyButtonClicked ->
-                        applyButton.setOnClickListener { onApplyButtonClicked?.invoke() }
+                        applyButton.setOnClickListener {
+                            if (onApplyButtonClicked != null) {
+                                onApplyButtonClicked.invoke()
+                                onNavBack.invoke()
+                            }
+                        }
                     }
                 }
 
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModel.kt
index 7f3c4cb..7be934d 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModel.kt
@@ -20,7 +20,6 @@
 import android.content.res.Resources
 import com.android.customization.model.ResourceConstants
 import com.android.customization.model.grid.GridOptionModel
-import com.android.customization.model.grid.ShapeOptionModel
 import com.android.customization.picker.grid.domain.interactor.ShapeGridInteractor
 import com.android.customization.picker.grid.ui.viewmodel.GridIconViewModel
 import com.android.customization.picker.grid.ui.viewmodel.ShapeIconViewModel
@@ -42,6 +41,7 @@
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.filterNotNull
+import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.flow.stateIn
@@ -59,7 +59,36 @@
         GRID,
     }
 
-    //// Tabs
+    // TODO (b/362237825): Get the shape options from the Launcher
+    val shapePaths =
+        listOf(
+            Pair(
+                "arch",
+                "M100 83.46C100 85.471 100 86.476 99.9 87.321 99.116 93.916 93.916 99.116 87.321 99.9 86.476 100 85.471 100 83.46 100H16.54C14.529 100 13.524 100 12.679 99.9 6.084 99.116.884 93.916.1 87.321 0 86.476 0 85.471 0 83.46L0 50C0 22.386 22.386 0 50 0 77.614 0 100 22.386 100 50V83.46Z",
+            ),
+            Pair(
+                "4-sided-cookie",
+                "M63.605 3C84.733-6.176 106.176 15.268 97 36.395L95.483 39.888C92.681 46.338 92.681 53.662 95.483 60.112L97 63.605C106.176 84.732 84.733 106.176 63.605 97L60.112 95.483C53.662 92.681 46.338 92.681 39.888 95.483L36.395 97C15.267 106.176-6.176 84.732 3 63.605L4.517 60.112C7.319 53.662 7.319 46.338 4.517 39.888L3 36.395C-6.176 15.268 15.267-6.176 36.395 3L39.888 4.517C46.338 7.319 53.662 7.319 60.112 4.517L63.605 3Z",
+            ),
+            Pair(
+                "7-sided-cookie",
+                "M35.209 4.878C36.326 3.895 36.884 3.404 37.397 3.006 44.82-2.742 55.18-2.742 62.603 3.006 63.116 3.404 63.674 3.895 64.791 4.878 65.164 5.207 65.351 5.371 65.539 5.529 68.167 7.734 71.303 9.248 74.663 9.932 74.902 9.981 75.147 10.025 75.637 10.113 77.1 10.375 77.831 10.506 78.461 10.66 87.573 12.893 94.032 21.011 94.176 30.412 94.186 31.062 94.151 31.805 94.08 33.293 94.057 33.791 94.045 34.04 94.039 34.285 93.958 37.72 94.732 41.121 96.293 44.18 96.404 44.399 96.522 44.618 96.759 45.056 97.467 46.366 97.821 47.021 98.093 47.611 102.032 56.143 99.727 66.266 92.484 72.24 91.983 72.653 91.381 73.089 90.177 73.961 89.774 74.254 89.572 74.4 89.377 74.548 86.647 76.626 84.477 79.353 83.063 82.483 82.962 82.707 82.865 82.936 82.671 83.395 82.091 84.766 81.8 85.451 81.51 86.033 77.31 94.44 67.977 98.945 58.801 96.994 58.166 96.859 57.451 96.659 56.019 96.259 55.54 96.125 55.3 96.058 55.063 95.998 51.74 95.154 48.26 95.154 44.937 95.998 44.699 96.058 44.46 96.125 43.981 96.259 42.549 96.659 41.834 96.859 41.199 96.994 32.023 98.945 22.69 94.44 18.49 86.033 18.2 85.451 17.909 84.766 17.329 83.395 17.135 82.936 17.038 82.707 16.937 82.483 15.523 79.353 13.353 76.626 10.623 74.548 10.428 74.4 10.226 74.254 9.823 73.961 8.619 73.089 8.017 72.653 7.516 72.24.273 66.266-2.032 56.143 1.907 47.611 2.179 47.021 2.533 46.366 3.241 45.056 3.478 44.618 3.596 44.399 3.707 44.18 5.268 41.121 6.042 37.72 5.961 34.285 5.955 34.04 5.943 33.791 5.92 33.293 5.849 31.805 5.814 31.062 5.824 30.412 5.968 21.011 12.427 12.893 21.539 10.66 22.169 10.506 22.9 10.375 24.363 10.113 24.853 10.025 25.098 9.981 25.337 9.932 28.697 9.248 31.833 7.734 34.461 5.529 34.649 5.371 34.836 5.207 35.209 4.878Z",
+            ),
+            Pair(
+                "sunny",
+                "M42.846 4.873C46.084-.531 53.916-.531 57.154 4.873L60.796 10.951C62.685 14.103 66.414 15.647 69.978 14.754L76.851 13.032C82.962 11.5 88.5 17.038 86.968 23.149L85.246 30.022C84.353 33.586 85.897 37.315 89.049 39.204L95.127 42.846C100.531 46.084 100.531 53.916 95.127 57.154L89.049 60.796C85.897 62.685 84.353 66.414 85.246 69.978L86.968 76.851C88.5 82.962 82.962 88.5 76.851 86.968L69.978 85.246C66.414 84.353 62.685 85.898 60.796 89.049L57.154 95.127C53.916 100.531 46.084 100.531 42.846 95.127L39.204 89.049C37.315 85.898 33.586 84.353 30.022 85.246L23.149 86.968C17.038 88.5 11.5 82.962 13.032 76.851L14.754 69.978C15.647 66.414 14.103 62.685 10.951 60.796L4.873 57.154C-.531 53.916-.531 46.084 4.873 42.846L10.951 39.204C14.103 37.315 15.647 33.586 14.754 30.022L13.032 23.149C11.5 17.038 17.038 11.5 23.149 13.032L30.022 14.754C33.586 15.647 37.315 14.103 39.204 10.951L42.846 4.873Z",
+            ),
+            Pair(
+                "circle",
+                "M99.18 50C99.18 77.162 77.162 99.18 50 99.18 22.838 99.18.82 77.162.82 50 .82 22.839 22.838.82 50 .82 77.162.82 99.18 22.839 99.18 50Z",
+            ),
+            Pair(
+                "square",
+                "M99.18 53.689C99.18 67.434 99.18 74.306 97.022 79.758 93.897 87.649 87.649 93.897 79.758 97.022 74.306 99.18 67.434 99.18 53.689 99.18H46.311C32.566 99.18 25.694 99.18 20.242 97.022 12.351 93.897 6.103 87.649 2.978 79.758.82 74.306.82 67.434.82 53.689L.82 46.311C.82 32.566.82 25.694 2.978 20.242 6.103 12.351 12.351 6.103 20.242 2.978 25.694.82 32.566.82 46.311.82L53.689.82C67.434.82 74.306.82 79.758 2.978 87.649 6.103 93.897 12.351 97.022 20.242 99.18 25.694 99.18 32.566 99.18 46.311V53.689Z",
+            ),
+        )
+
+    // Tabs
     private val _selectedTab = MutableStateFlow(Tab.SHAPE)
     val selectedTab: StateFlow<Tab> = _selectedTab.asStateFlow()
     val tabs: Flow<List<FloatingToolbarTabViewModel>> =
@@ -88,14 +117,8 @@
             )
         }
 
-    //// Shape
-
     // The currently-set system shape option
-    val selectedShapeKey =
-        interactor.selectedShapeOption
-            .filterNotNull()
-            .map { it.key }
-            .shareIn(scope = viewModelScope, started = SharingStarted.Lazily, replay = 1)
+    val selectedShapeKey = flowOf(shapePaths[0].first)
     private val overridingShapeKey = MutableStateFlow<String?>(null)
     // If the overriding key is null, use the currently-set system shape option
     val previewingShapeKey =
@@ -105,24 +128,21 @@
         }
 
     val shapeOptions: Flow<List<OptionItemViewModel<ShapeIconViewModel>>> =
-        interactor.shapeOptions
-            .filterNotNull()
-            .map { shapeOptions -> shapeOptions.map { toShapeOptionItemViewModel(it) } }
+        MutableStateFlow(shapePaths.map { toShapeOptionItemViewModel(it.first, it.second) })
             .shareIn(scope = viewModelScope, started = SharingStarted.Lazily, replay = 1)
 
-    //// Grid
-
     // The currently-set system grid option
     val selectedGridOption =
         interactor.selectedGridOption
             .filterNotNull()
             .map { toGridOptionItemViewModel(it) }
             .shareIn(scope = viewModelScope, started = SharingStarted.Lazily, replay = 1)
-    private val overridingGridKey = MutableStateFlow<String?>(null)
+    private val overridingGridOptionKey = MutableStateFlow<String?>(null)
     // If the overriding key is null, use the currently-set system grid option
-    val previewingGridKey =
-        combine(overridingGridKey, selectedGridOption) { overridingGridOptionKey, selectedGridOption
-            ->
+    val previewingGridOptionKey =
+        combine(overridingGridOptionKey, selectedGridOption) {
+            overridingGridOptionKey,
+            selectedGridOption ->
             overridingGridOptionKey ?: selectedGridOption.key.value
         }
 
@@ -133,33 +153,29 @@
             .shareIn(scope = viewModelScope, started = SharingStarted.Lazily, replay = 1)
 
     val onApply: Flow<(suspend () -> Unit)?> =
-        combine(previewingGridKey, selectedGridOption, previewingShapeKey, selectedShapeKey) {
+        combine(previewingGridOptionKey, selectedGridOption) {
             previewingGridOptionKey,
-            selectedGridOption,
-            previewingShapeKey,
-            selectedShapeKey ->
-            if (
-                previewingGridOptionKey == selectedGridOption.key.value &&
-                    previewingShapeKey == selectedShapeKey
-            ) {
+            selectedGridOption ->
+            if (previewingGridOptionKey == selectedGridOption.key.value) {
                 null
             } else {
-                { interactor.applySelectedOption(previewingShapeKey, previewingGridOptionKey) }
+                { interactor.applySelectedOption(previewingGridOptionKey) }
             }
         }
 
     fun resetPreview() {
         overridingShapeKey.value = null
-        overridingGridKey.value = null
+        overridingGridOptionKey.value = null
         _selectedTab.value = Tab.SHAPE
     }
 
     private fun toShapeOptionItemViewModel(
-        option: ShapeOptionModel
+        key: String,
+        path: String,
     ): OptionItemViewModel<ShapeIconViewModel> {
         val isSelected =
             previewingShapeKey
-                .map { it == option.key }
+                .map { it == key }
                 .stateIn(
                     scope = viewModelScope,
                     started = SharingStarted.Lazily,
@@ -167,14 +183,14 @@
                 )
 
         return OptionItemViewModel(
-            key = MutableStateFlow(option.key),
-            payload = ShapeIconViewModel(option.key, option.path),
-            text = Text.Loaded(option.title),
+            key = MutableStateFlow(key),
+            payload = ShapeIconViewModel(key, path),
+            text = Text.Loaded(key),
             isSelected = isSelected,
             onClicked =
                 isSelected.map {
                     if (!it) {
-                        { overridingShapeKey.value = option.key }
+                        { overridingShapeKey.value = key }
                     } else {
                         null
                     }
@@ -195,7 +211,7 @@
                     )
             )
         val isSelected =
-            previewingGridKey
+            previewingGridOptionKey
                 .map { it == option.key }
                 .stateIn(
                     scope = viewModelScope,
@@ -212,7 +228,7 @@
             onClicked =
                 isSelected.map {
                     if (!it) {
-                        { overridingGridKey.value = option.key }
+                        { overridingGridOptionKey.value = option.key }
                     } else {
                         null
                     }
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
index d54fd59..2e500e2 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
@@ -60,13 +60,20 @@
 
     override val selectedOption = defaultCustomizationOptionsViewModel.selectedOption
 
-    override fun deselectOption(): Boolean {
-        keyguardQuickAffordancePickerViewModel2.resetPreview()
-        shapeGridPickerViewModel.resetPreview()
-        clockPickerViewModel.resetPreview()
-        colorPickerViewModel2.resetPreview()
-        darkModeViewModel.resetPreview()
-        return defaultCustomizationOptionsViewModel.deselectOption()
+    override fun handleBackPressed(): Boolean {
+        val isBackPressedHandled = defaultCustomizationOptionsViewModel.handleBackPressed()
+
+        if (isBackPressedHandled) {
+            // If isBackPressedHandled is handled by DefaultCustomizationOptionsViewModel, it means
+            // we navigate back to the main screen from a secondary screen. Reset preview.
+            keyguardQuickAffordancePickerViewModel2.resetPreview()
+            shapeGridPickerViewModel.resetPreview()
+            clockPickerViewModel.resetPreview()
+            colorPickerViewModel2.resetPreview()
+            darkModeViewModel.resetPreview()
+        }
+
+        return isBackPressedHandled
     }
 
     val onCustomizeClockClicked: Flow<(() -> Unit)?> =
@@ -147,16 +154,12 @@
                 }
             }
             .map { onApply ->
-                {
-                    if (onApply != null) {
-                        viewModelScope.launch {
-                            onApply()
-                            // We only wait until onApply() is done to execute deselectOption()
-                            deselectOption()
-                        }
-                    } else {
-                        null
+                if (onApply != null) {
+                    fun() {
+                        viewModelScope.launch { onApply() }
                     }
+                } else {
+                    null
                 }
             }
             .stateIn(viewModelScope, SharingStarted.Eagerly, 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
index dea3351..0dc38e8 100644
--- a/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
+++ b/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
@@ -24,8 +24,6 @@
 import androidx.lifecycle.LifecycleOwner
 import androidx.lifecycle.lifecycleScope
 import androidx.lifecycle.repeatOnLifecycle
-import com.android.customization.model.grid.DefaultShapeGridManager.Companion.COL_GRID_KEY
-import com.android.customization.model.grid.DefaultShapeGridManager.Companion.COL_SHAPE_KEY
 import com.android.customization.picker.color.data.util.MaterialColorsGenerator
 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
@@ -139,19 +137,10 @@
                 lifecycleOwner.lifecycleScope.launch {
                     lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
                         launch {
-                            viewModel.shapeGridPickerViewModel.previewingShapeKey.collect {
-                                workspaceCallback.sendMessage(
-                                    MESSAGE_ID_UPDATE_SHAPE,
-                                    bundleOf(COL_SHAPE_KEY to it),
-                                )
-                            }
-                        }
-
-                        launch {
-                            viewModel.shapeGridPickerViewModel.previewingGridKey.collect {
+                            viewModel.shapeGridPickerViewModel.previewingGridOptionKey.collect {
                                 workspaceCallback.sendMessage(
                                     MESSAGE_ID_UPDATE_GRID,
-                                    bundleOf(COL_GRID_KEY to it),
+                                    bundleOf(KEY_GRID_NAME to it),
                                 )
                             }
                         }
@@ -183,8 +172,8 @@
     }
 
     companion object {
-        const val MESSAGE_ID_UPDATE_SHAPE = 2586
         const val MESSAGE_ID_UPDATE_GRID = 7414
+        const val KEY_GRID_NAME = "grid_name"
 
         const val MESSAGE_ID_UPDATE_COLOR = 856
         const val KEY_COLOR_RESOURCE_IDS: String = "color_resource_ids"
diff --git a/tests/common/src/com/android/customization/model/grid/FakeShapeGridManager.kt b/tests/common/src/com/android/customization/model/grid/FakeShapeGridManager.kt
index b1f044a..a5a94c3 100644
--- a/tests/common/src/com/android/customization/model/grid/FakeShapeGridManager.kt
+++ b/tests/common/src/com/android/customization/model/grid/FakeShapeGridManager.kt
@@ -22,17 +22,16 @@
 @Singleton
 class FakeShapeGridManager @Inject constructor() : ShapeGridManager {
 
+    var isGridOptionAvailable: Boolean = true
+
     private var gridOptions: List<GridOptionModel>? = DEFAULT_GRID_OPTION_LIST
 
-    private var shapeOptions: List<ShapeOptionModel>? = DEFAULT_SHAPE_OPTION_LIST
+    override suspend fun isGridOptionAvailable(): Boolean = isGridOptionAvailable
 
     override suspend fun getGridOptions(): List<GridOptionModel>? = gridOptions
 
-    override suspend fun getShapeOptions(): List<ShapeOptionModel>? = shapeOptions
-
-    override fun applyShapeGridOption(shapeKey: String, gridKey: String): Int {
-        shapeOptions = shapeOptions?.map { it.copy(isCurrent = it.key == shapeKey) }
-        gridOptions = gridOptions?.map { it.copy(isCurrent = it.key == gridKey) }
+    override fun applyGridOption(gridName: String): Int {
+        gridOptions = gridOptions?.map { it.copy(isCurrent = it.key == gridName) }
         return 0
     }
 
@@ -54,51 +53,5 @@
                     cols = 4,
                 ),
             )
-
-        val DEFAULT_SHAPE_OPTION_LIST =
-            listOf(
-                ShapeOptionModel(
-                    key = "arch",
-                    title = "arch",
-                    path =
-                        "M100 83.46C100 85.471 100 86.476 99.9 87.321 99.116 93.916 93.916 99.116 87.321 99.9 86.476 100 85.471 100 83.46 100H16.54C14.529 100 13.524 100 12.679 99.9 6.084 99.116.884 93.916.1 87.321 0 86.476 0 85.471 0 83.46L0 50C0 22.386 22.386 0 50 0 77.614 0 100 22.386 100 50V83.46Z",
-                    isCurrent = true,
-                ),
-                ShapeOptionModel(
-                    key = "4-sided-cookie",
-                    title = "4-sided-cookie",
-                    path =
-                        "M63.605 3C84.733-6.176 106.176 15.268 97 36.395L95.483 39.888C92.681 46.338 92.681 53.662 95.483 60.112L97 63.605C106.176 84.732 84.733 106.176 63.605 97L60.112 95.483C53.662 92.681 46.338 92.681 39.888 95.483L36.395 97C15.267 106.176-6.176 84.732 3 63.605L4.517 60.112C7.319 53.662 7.319 46.338 4.517 39.888L3 36.395C-6.176 15.268 15.267-6.176 36.395 3L39.888 4.517C46.338 7.319 53.662 7.319 60.112 4.517L63.605 3Z",
-                    isCurrent = false,
-                ),
-                ShapeOptionModel(
-                    key = "7-sided-cookie",
-                    title = "7-sided-cookie",
-                    path =
-                        "M35.209 4.878C36.326 3.895 36.884 3.404 37.397 3.006 44.82-2.742 55.18-2.742 62.603 3.006 63.116 3.404 63.674 3.895 64.791 4.878 65.164 5.207 65.351 5.371 65.539 5.529 68.167 7.734 71.303 9.248 74.663 9.932 74.902 9.981 75.147 10.025 75.637 10.113 77.1 10.375 77.831 10.506 78.461 10.66 87.573 12.893 94.032 21.011 94.176 30.412 94.186 31.062 94.151 31.805 94.08 33.293 94.057 33.791 94.045 34.04 94.039 34.285 93.958 37.72 94.732 41.121 96.293 44.18 96.404 44.399 96.522 44.618 96.759 45.056 97.467 46.366 97.821 47.021 98.093 47.611 102.032 56.143 99.727 66.266 92.484 72.24 91.983 72.653 91.381 73.089 90.177 73.961 89.774 74.254 89.572 74.4 89.377 74.548 86.647 76.626 84.477 79.353 83.063 82.483 82.962 82.707 82.865 82.936 82.671 83.395 82.091 84.766 81.8 85.451 81.51 86.033 77.31 94.44 67.977 98.945 58.801 96.994 58.166 96.859 57.451 96.659 56.019 96.259 55.54 96.125 55.3 96.058 55.063 95.998 51.74 95.154 48.26 95.154 44.937 95.998 44.699 96.058 44.46 96.125 43.981 96.259 42.549 96.659 41.834 96.859 41.199 96.994 32.023 98.945 22.69 94.44 18.49 86.033 18.2 85.451 17.909 84.766 17.329 83.395 17.135 82.936 17.038 82.707 16.937 82.483 15.523 79.353 13.353 76.626 10.623 74.548 10.428 74.4 10.226 74.254 9.823 73.961 8.619 73.089 8.017 72.653 7.516 72.24.273 66.266-2.032 56.143 1.907 47.611 2.179 47.021 2.533 46.366 3.241 45.056 3.478 44.618 3.596 44.399 3.707 44.18 5.268 41.121 6.042 37.72 5.961 34.285 5.955 34.04 5.943 33.791 5.92 33.293 5.849 31.805 5.814 31.062 5.824 30.412 5.968 21.011 12.427 12.893 21.539 10.66 22.169 10.506 22.9 10.375 24.363 10.113 24.853 10.025 25.098 9.981 25.337 9.932 28.697 9.248 31.833 7.734 34.461 5.529 34.649 5.371 34.836 5.207 35.209 4.878Z",
-                    isCurrent = false,
-                ),
-                ShapeOptionModel(
-                    key = "sunny",
-                    title = "sunny",
-                    path =
-                        "M42.846 4.873C46.084-.531 53.916-.531 57.154 4.873L60.796 10.951C62.685 14.103 66.414 15.647 69.978 14.754L76.851 13.032C82.962 11.5 88.5 17.038 86.968 23.149L85.246 30.022C84.353 33.586 85.897 37.315 89.049 39.204L95.127 42.846C100.531 46.084 100.531 53.916 95.127 57.154L89.049 60.796C85.897 62.685 84.353 66.414 85.246 69.978L86.968 76.851C88.5 82.962 82.962 88.5 76.851 86.968L69.978 85.246C66.414 84.353 62.685 85.898 60.796 89.049L57.154 95.127C53.916 100.531 46.084 100.531 42.846 95.127L39.204 89.049C37.315 85.898 33.586 84.353 30.022 85.246L23.149 86.968C17.038 88.5 11.5 82.962 13.032 76.851L14.754 69.978C15.647 66.414 14.103 62.685 10.951 60.796L4.873 57.154C-.531 53.916-.531 46.084 4.873 42.846L10.951 39.204C14.103 37.315 15.647 33.586 14.754 30.022L13.032 23.149C11.5 17.038 17.038 11.5 23.149 13.032L30.022 14.754C33.586 15.647 37.315 14.103 39.204 10.951L42.846 4.873Z",
-                    isCurrent = false,
-                ),
-                ShapeOptionModel(
-                    key = "circle",
-                    title = "circle",
-                    path =
-                        "M99.18 50C99.18 77.162 77.162 99.18 50 99.18 22.838 99.18.82 77.162.82 50 .82 22.839 22.838.82 50 .82 77.162.82 99.18 22.839 99.18 50Z",
-                    isCurrent = false,
-                ),
-                ShapeOptionModel(
-                    key = "square",
-                    title = "square",
-                    path =
-                        "M99.18 53.689C99.18 67.434 99.18 74.306 97.022 79.758 93.897 87.649 87.649 93.897 79.758 97.022 74.306 99.18 67.434 99.18 53.689 99.18H46.311C32.566 99.18 25.694 99.18 20.242 97.022 12.351 93.897 6.103 87.649 2.978 79.758.82 74.306.82 67.434.82 53.689L.82 46.311C.82 32.566.82 25.694 2.978 20.242 6.103 12.351 12.351 6.103 20.242 2.978 25.694.82 32.566.82 46.311.82L53.689.82C67.434.82 74.306.82 79.758 2.978 87.649 6.103 93.897 12.351 97.022 20.242 99.18 25.694 99.18 32.566 99.18 46.311V53.689Z",
-                    isCurrent = false,
-                ),
-            )
     }
 }
diff --git a/tests/robotests/src/com/android/customization/picker/grid/data/repository/ShapeGridRepositoryTest.kt b/tests/robotests/src/com/android/customization/picker/grid/data/repository/ShapeGridRepositoryTest.kt
index 985d983..61f8743 100644
--- a/tests/robotests/src/com/android/customization/picker/grid/data/repository/ShapeGridRepositoryTest.kt
+++ b/tests/robotests/src/com/android/customization/picker/grid/data/repository/ShapeGridRepositoryTest.kt
@@ -69,87 +69,52 @@
     }
 
     @Test
-    fun shapeOptions_default() =
+    fun isGridOptionAvailable_false() =
         testScope.runTest {
-            val gridOptions = collectLastValue(underTest.shapeOptions)
-
-            assertThat(gridOptions()).isEqualTo(FakeShapeGridManager.DEFAULT_SHAPE_OPTION_LIST)
+            gridOptionsManager.isGridOptionAvailable = false
+            assertThat(underTest.isGridOptionAvailable()).isFalse()
         }
 
     @Test
-    fun shapeOptions_shouldUpdateAfterApplyShapeGridOption() =
+    fun isGridOptionAvailable_true() =
         testScope.runTest {
-            val shapeOptions = collectLastValue(underTest.shapeOptions)
-
-            underTest.applySelectedOption("circle", "practical")
-
-            assertThat(shapeOptions())
-                .isEqualTo(
-                    FakeShapeGridManager.DEFAULT_SHAPE_OPTION_LIST.map {
-                        it.copy(isCurrent = (it.key == "circle"))
-                    }
-                )
-        }
-
-    @Test
-    fun selectedShapeOption_default() =
-        testScope.runTest {
-            val selectedGridOption = collectLastValue(underTest.selectedShapeOption)
-
-            assertThat(selectedGridOption())
-                .isEqualTo(FakeShapeGridManager.DEFAULT_SHAPE_OPTION_LIST[0])
-        }
-
-    @Test
-    fun selectedShapeOption_shouldUpdateAfterApplyShapeGridOption() =
-        testScope.runTest {
-            val selectedShapeOption = collectLastValue(underTest.selectedShapeOption)
-
-            underTest.applySelectedOption("circle", "practical")
-
-            assertThat(selectedShapeOption())
-                .isEqualTo(FakeShapeGridManager.DEFAULT_SHAPE_OPTION_LIST[4].copy(isCurrent = true))
+            gridOptionsManager.isGridOptionAvailable = true
+            assertThat(underTest.isGridOptionAvailable()).isTrue()
         }
 
     @Test
     fun gridOptions_default() =
         testScope.runTest {
             val gridOptions = collectLastValue(underTest.gridOptions)
-
             assertThat(gridOptions()).isEqualTo(FakeShapeGridManager.DEFAULT_GRID_OPTION_LIST)
         }
 
     @Test
-    fun gridOptions_shouldUpdateAfterApplyShapeGridOption() =
-        testScope.runTest {
-            val gridOptions = collectLastValue(underTest.gridOptions)
-
-            underTest.applySelectedOption("circle", "practical")
-
-            assertThat(gridOptions())
-                .isEqualTo(
-                    FakeShapeGridManager.DEFAULT_GRID_OPTION_LIST.map {
-                        it.copy(isCurrent = (it.key == "practical"))
-                    }
-                )
-        }
-
-    @Test
     fun selectedGridOption_default() =
         testScope.runTest {
             val selectedGridOption = collectLastValue(underTest.selectedGridOption)
-
             assertThat(selectedGridOption())
                 .isEqualTo(FakeShapeGridManager.DEFAULT_GRID_OPTION_LIST[0])
         }
 
     @Test
-    fun selectedGridOption_shouldUpdateAfterApplyShapeGridOption() =
+    fun gridOptions_shouldUpdateAfterApplyGridOption() =
+        testScope.runTest {
+            val gridOptions = collectLastValue(underTest.gridOptions)
+            underTest.applySelectedOption("practical")
+            assertThat(gridOptions())
+                .isEqualTo(
+                    FakeShapeGridManager.DEFAULT_GRID_OPTION_LIST.map {
+                        it.copy(isCurrent = it.key == "practical")
+                    }
+                )
+        }
+
+    @Test
+    fun selectedGridOption_shouldUpdateAfterApplyGridOption() =
         testScope.runTest {
             val selectedGridOption = collectLastValue(underTest.selectedGridOption)
-
-            underTest.applySelectedOption("circle", "practical")
-
+            underTest.applySelectedOption("practical")
             assertThat(selectedGridOption())
                 .isEqualTo(FakeShapeGridManager.DEFAULT_GRID_OPTION_LIST[1].copy(isCurrent = true))
         }
diff --git a/tests/robotests/src/com/android/customization/picker/grid/domain/interactor/ShapeGridInteractorTest.kt b/tests/robotests/src/com/android/customization/picker/grid/domain/interactor/ShapeGridInteractorTest.kt
index c0f519c..51642cb 100644
--- a/tests/robotests/src/com/android/customization/picker/grid/domain/interactor/ShapeGridInteractorTest.kt
+++ b/tests/robotests/src/com/android/customization/picker/grid/domain/interactor/ShapeGridInteractorTest.kt
@@ -61,87 +61,52 @@
     }
 
     @Test
-    fun shapeOptions_default() =
+    fun isGridOptionAvailable_false() =
         testScope.runTest {
-            val shapeOptions = collectLastValue(underTest.shapeOptions)
-
-            assertThat(shapeOptions()).isEqualTo(FakeShapeGridManager.DEFAULT_SHAPE_OPTION_LIST)
+            gridOptionsManager.isGridOptionAvailable = false
+            assertThat(underTest.isGridOptionAvailable()).isFalse()
         }
 
     @Test
-    fun shapeOptions_shouldUpdateAfterApplyGridOption() =
+    fun isGridOptionAvailable_true() =
         testScope.runTest {
-            val shapeOptions = collectLastValue(underTest.shapeOptions)
-
-            underTest.applySelectedOption("circle", "practical")
-
-            assertThat(shapeOptions())
-                .isEqualTo(
-                    FakeShapeGridManager.DEFAULT_SHAPE_OPTION_LIST.map {
-                        it.copy(isCurrent = (it.key == "circle"))
-                    }
-                )
-        }
-
-    @Test
-    fun selectedShapeOption_default() =
-        testScope.runTest {
-            val selectedShapeOption = collectLastValue(underTest.selectedShapeOption)
-
-            assertThat(selectedShapeOption())
-                .isEqualTo(FakeShapeGridManager.DEFAULT_SHAPE_OPTION_LIST[0])
-        }
-
-    @Test
-    fun selectedShapeOption_shouldUpdateAfterApplyGridOption() =
-        testScope.runTest {
-            val selectedShapeOption = collectLastValue(underTest.selectedShapeOption)
-
-            underTest.applySelectedOption("circle", "practical")
-
-            assertThat(selectedShapeOption())
-                .isEqualTo(FakeShapeGridManager.DEFAULT_SHAPE_OPTION_LIST[4].copy(isCurrent = true))
+            gridOptionsManager.isGridOptionAvailable = true
+            assertThat(underTest.isGridOptionAvailable()).isTrue()
         }
 
     @Test
     fun gridOptions_default() =
         testScope.runTest {
             val gridOptions = collectLastValue(underTest.gridOptions)
-
             assertThat(gridOptions()).isEqualTo(FakeShapeGridManager.DEFAULT_GRID_OPTION_LIST)
         }
 
     @Test
-    fun gridOptions_shouldUpdateAfterApplyGridOption() =
-        testScope.runTest {
-            val gridOptions = collectLastValue(underTest.gridOptions)
-
-            underTest.applySelectedOption("arch", "practical")
-
-            assertThat(gridOptions())
-                .isEqualTo(
-                    FakeShapeGridManager.DEFAULT_GRID_OPTION_LIST.map {
-                        it.copy(isCurrent = (it.key == "practical"))
-                    }
-                )
-        }
-
-    @Test
     fun selectedGridOption_default() =
         testScope.runTest {
             val selectedGridOption = collectLastValue(underTest.selectedGridOption)
-
             assertThat(selectedGridOption())
                 .isEqualTo(FakeShapeGridManager.DEFAULT_GRID_OPTION_LIST[0])
         }
 
     @Test
+    fun gridOptions_shouldUpdateAfterApplyGridOption() =
+        testScope.runTest {
+            val gridOptions = collectLastValue(underTest.gridOptions)
+            underTest.applySelectedOption("practical")
+            assertThat(gridOptions())
+                .isEqualTo(
+                    FakeShapeGridManager.DEFAULT_GRID_OPTION_LIST.map {
+                        it.copy(isCurrent = it.key == "practical")
+                    }
+                )
+        }
+
+    @Test
     fun selectedGridOption_shouldUpdateAfterApplyGridOption() =
         testScope.runTest {
             val selectedGridOption = collectLastValue(underTest.selectedGridOption)
-
-            underTest.applySelectedOption("arch", "practical")
-
+            underTest.applySelectedOption("practical")
             assertThat(selectedGridOption())
                 .isEqualTo(FakeShapeGridManager.DEFAULT_GRID_OPTION_LIST[1].copy(isCurrent = true))
         }
diff --git a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModelTest.kt b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModelTest.kt
index c99acca..dcf658e 100644
--- a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModelTest.kt
+++ b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModelTest.kt
@@ -109,16 +109,14 @@
         testScope.runTest {
             val shapeOptions = collectLastValue(underTest.shapeOptions)
 
-            for (i in 0 until FakeShapeGridManager.DEFAULT_SHAPE_OPTION_LIST.size) {
-                val (expectedKey, expectedPath, expectedTitle) =
-                    with(FakeShapeGridManager.DEFAULT_SHAPE_OPTION_LIST[i]) {
-                        arrayOf(key, path, title)
-                    }
+            for (i in 0 until underTest.shapePaths.size) {
+                val expectedKey = underTest.shapePaths[i].first
+                val expectedPath = underTest.shapePaths[i].second
                 assertShapeItem(
                     optionItem = shapeOptions()?.get(i),
-                    key = FakeShapeGridManager.DEFAULT_SHAPE_OPTION_LIST[i].key,
+                    key = expectedKey,
                     payload = ShapeIconViewModel(expectedKey, expectedPath),
-                    text = Text.Loaded(expectedTitle),
+                    text = Text.Loaded(expectedKey),
                     isTextUserVisible = true,
                     isSelected = expectedKey == "arch",
                     isEnabled = true,
@@ -138,15 +136,14 @@
             onCircleOptionClicked()?.invoke()
 
             assertThat(previewingShapeKey()).isEqualTo("circle")
-            for (i in 0 until FakeShapeGridManager.DEFAULT_SHAPE_OPTION_LIST.size) {
-                val expectedKey = FakeShapeGridManager.DEFAULT_SHAPE_OPTION_LIST[i].key
-                val expectedPath = FakeShapeGridManager.DEFAULT_SHAPE_OPTION_LIST[i].path
-                val expectedTitle = FakeShapeGridManager.DEFAULT_SHAPE_OPTION_LIST[i].title
+            for (i in 0 until underTest.shapePaths.size) {
+                val expectedKey = underTest.shapePaths[i].first
+                val expectedPath = underTest.shapePaths[i].second
                 assertShapeItem(
                     optionItem = shapeOptions()?.get(i),
                     key = expectedKey,
                     payload = ShapeIconViewModel(expectedKey, expectedPath),
-                    text = Text.Loaded(expectedTitle),
+                    text = Text.Loaded(expectedKey),
                     isTextUserVisible = true,
                     isSelected = expectedKey == "circle",
                     isEnabled = true,