Grid customiztion screen (1/2)

Test: Manually tested. See bug.
Bug: 362237825
Flag: com.android.systemui.shared.new_customization_picker_ui
Change-Id: I717112c9aa40a8deee4ddd1890d23738b8c75998
diff --git a/res/layout/customization_option_entry_app_grid.xml b/res/layout/customization_option_entry_app_grid.xml
index c31b2a2..ea6da46 100644
--- a/res/layout/customization_option_entry_app_grid.xml
+++ b/res/layout/customization_option_entry_app_grid.xml
@@ -31,7 +31,7 @@
         android:text="@string/grid_title"
         android:layout_marginEnd="@dimen/customization_option_entry_text_margin_end"
         app:layout_constraintStart_toStartOf="parent"
-        app:layout_constraintEnd_toStartOf="@+id/option_entry_app_grid_icon"
+        app:layout_constraintEnd_toStartOf="@+id/option_entry_app_grid_icon_container"
         app:layout_constraintBottom_toTopOf="@+id/option_entry_app_grid_description"
         app:layout_constraintTop_toTopOf="parent"
         app:layout_constraintVertical_chainStyle="packed" />
@@ -42,19 +42,25 @@
         android:layout_width="0dp"
         android:layout_height="wrap_content"
         android:layout_marginEnd="@dimen/customization_option_entry_text_margin_end"
-        android:text="4x4"
         app:layout_constraintBottom_toBottomOf="parent"
-        app:layout_constraintEnd_toStartOf="@+id/option_entry_app_grid_icon"
+        app:layout_constraintEnd_toStartOf="@+id/option_entry_app_grid_icon_container"
         app:layout_constraintStart_toStartOf="parent"
         app:layout_constraintTop_toBottomOf="@+id/option_entry_app_grid_title" />
 
     <FrameLayout
-        android:id="@+id/option_entry_app_grid_icon"
+        android:id="@+id/option_entry_app_grid_icon_container"
         android:layout_width="@dimen/customization_option_entry_icon_size"
         android:layout_height="@dimen/customization_option_entry_icon_size"
-        android:orientation="horizontal"
+        android:padding="@dimen/customization_option_entry_icon_padding"
         android:background="@drawable/customization_option_entry_icon_background"
         app:layout_constraintEnd_toEndOf="parent"
         app:layout_constraintTop_toTopOf="parent"
-        app:layout_constraintBottom_toBottomOf="parent" />
+        app:layout_constraintBottom_toBottomOf="parent">
+
+        <ImageView
+            android:id="@+id/option_entry_app_grid_icon"
+            android:layout_width="match_parent"
+            android:layout_height="match_parent"
+            android:contentDescription="@string/grid_preview_card_content_description" />
+    </FrameLayout>
 </androidx.constraintlayout.widget.ConstraintLayout>
diff --git a/res/layout/floating_sheet_shape_and_grid.xml b/res/layout/floating_sheet_shape_and_grid.xml
new file mode 100644
index 0000000..01a7a89
--- /dev/null
+++ b/res/layout/floating_sheet_shape_and_grid.xml
@@ -0,0 +1,54 @@
+<?xml version="1.0" encoding="utf-8"?><!--
+  ~ 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.
+  -->
+
+<LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="match_parent"
+    android:layout_height="wrap_content"
+    android:paddingHorizontal="@dimen/floating_sheet_horizontal_padding"
+    android:orientation="vertical">
+
+    <FrameLayout
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:paddingVertical="@dimen/floating_sheet_content_vertical_padding"
+        android:background="@drawable/floating_sheet_content_background"
+        android:clipToPadding="false"
+        android:clipChildren="false">
+
+        <!--
+        This is just an invisible placeholder put in place so that the parent keeps its height
+        stable as the RecyclerView updates from 0 items to N items. Keeping it stable allows the
+        layout logic to keep the size of the preview container stable as well, which bodes well
+        for setting up the SurfaceView for remote rendering without changing its size after the
+        content is loaded into the RecyclerView.
+
+        It's critical for any TextViews inside the included layout to have text.
+        -->
+        <include
+            layout="@layout/grid_option"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:visibility="invisible" />
+
+        <androidx.recyclerview.widget.RecyclerView
+            android:id="@id/options"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:layout_gravity="center_horizontal"
+            android:clipToPadding="false"
+            android:clipChildren="false" />
+    </FrameLayout>
+</LinearLayout>
\ No newline at end of file
diff --git a/src/com/android/customization/model/grid/DefaultGridOptionsManager.kt b/src/com/android/customization/model/grid/DefaultGridOptionsManager.kt
new file mode 100644
index 0000000..bc862fd
--- /dev/null
+++ b/src/com/android/customization/model/grid/DefaultGridOptionsManager.kt
@@ -0,0 +1,94 @@
+/*
+ * 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
+
+import android.content.ContentValues
+import android.content.Context
+import com.android.wallpaper.R
+import com.android.wallpaper.picker.di.modules.BackgroundDispatcher
+import com.android.wallpaper.util.PreviewUtils
+import dagger.hilt.android.qualifiers.ApplicationContext
+import javax.inject.Inject
+import javax.inject.Singleton
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.withContext
+
+@Singleton
+class DefaultGridOptionsManager
+@Inject
+constructor(
+    @ApplicationContext private val context: Context,
+    @BackgroundDispatcher private val bgDispatcher: CoroutineDispatcher,
+) : GridOptionsManager2 {
+
+    private val authorityMetadataKey: String =
+        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) {
+            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,
+                                )
+                            )
+                        }
+                    }
+                }
+        }
+
+    override fun applyGridOption(gridName: String): Int {
+        return context.contentResolver.update(
+            previewUtils.getUri(DEFAULT_GRID),
+            ContentValues().apply { put("name", gridName) },
+            null,
+            null,
+        )
+    }
+
+    companion object {
+        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"
+    }
+}
diff --git a/src/com/android/customization/model/grid/GridOptionModel.kt b/src/com/android/customization/model/grid/GridOptionModel.kt
new file mode 100644
index 0000000..3e10a01
--- /dev/null
+++ b/src/com/android/customization/model/grid/GridOptionModel.kt
@@ -0,0 +1,25 @@
+/*
+ * 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 GridOptionModel(
+    val key: String,
+    val title: String,
+    val isCurrent: Boolean,
+    val rows: Int,
+    val cols: Int,
+)
diff --git a/src/com/android/customization/model/grid/GridOptionsManager2.kt b/src/com/android/customization/model/grid/GridOptionsManager2.kt
new file mode 100644
index 0000000..ce8500a
--- /dev/null
+++ b/src/com/android/customization/model/grid/GridOptionsManager2.kt
@@ -0,0 +1,26 @@
+/*
+ * 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
+
+interface GridOptionsManager2 {
+
+    suspend fun isGridOptionAvailable(): Boolean
+
+    suspend fun getGridOptions(): List<GridOptionModel>?
+
+    fun applyGridOption(gridName: String): Int
+}
diff --git a/src/com/android/customization/picker/common/ui/view/SingleRowListItemSpacing.kt b/src/com/android/customization/picker/common/ui/view/SingleRowListItemSpacing.kt
new file mode 100644
index 0000000..5faf248
--- /dev/null
+++ b/src/com/android/customization/picker/common/ui/view/SingleRowListItemSpacing.kt
@@ -0,0 +1,51 @@
+/*
+ * 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.picker.common.ui.view
+
+import android.graphics.Rect
+import android.view.View
+import androidx.recyclerview.widget.RecyclerView
+
+/** Item spacing used by the horizontal RecyclerView with only 1 row. */
+class SingleRowListItemSpacing(
+    private val edgeItemSpacePx: Int,
+    private val itemHorizontalSpacePx: Int,
+) : RecyclerView.ItemDecoration() {
+    override fun getItemOffsets(
+        outRect: Rect,
+        view: View,
+        parent: RecyclerView,
+        state: RecyclerView.State,
+    ) {
+        val itemIndex = parent.getChildAdapterPosition(view)
+        val itemCount = parent.adapter?.itemCount ?: 0
+        val isRtl = parent.layoutManager?.layoutDirection == View.LAYOUT_DIRECTION_RTL
+        when (itemIndex) {
+            0 -> {
+                outRect.left = if (!isRtl) edgeItemSpacePx else itemHorizontalSpacePx
+                outRect.right = if (isRtl) edgeItemSpacePx else itemHorizontalSpacePx
+            }
+            itemCount - 1 -> {
+                outRect.right = if (!isRtl) edgeItemSpacePx else itemHorizontalSpacePx
+                outRect.left = if (isRtl) edgeItemSpacePx else itemHorizontalSpacePx
+            }
+            else -> {
+                outRect.left = itemHorizontalSpacePx
+                outRect.right = itemHorizontalSpacePx
+            }
+        }
+    }
+}
diff --git a/src/com/android/customization/picker/grid/data/repository/GridRepository2.kt b/src/com/android/customization/picker/grid/data/repository/GridRepository2.kt
new file mode 100644
index 0000000..8ce4374
--- /dev/null
+++ b/src/com/android/customization/picker/grid/data/repository/GridRepository2.kt
@@ -0,0 +1,67 @@
+/*
+ * 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.picker.grid.data.repository
+
+import com.android.customization.model.grid.GridOptionModel
+import com.android.customization.model.grid.GridOptionsManager2
+import com.android.wallpaper.picker.di.modules.BackgroundDispatcher
+import javax.inject.Inject
+import javax.inject.Singleton
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
+
+@Singleton
+class GridRepository2
+@Inject
+constructor(
+    private val manager: GridOptionsManager2,
+    @BackgroundDispatcher private val bgScope: CoroutineScope,
+    @BackgroundDispatcher private val bgDispatcher: CoroutineDispatcher,
+) {
+
+    suspend fun isGridOptionAvailable(): Boolean =
+        withContext(bgDispatcher) { manager.isGridOptionAvailable() }
+
+    private val _gridOptions = MutableStateFlow<List<GridOptionModel>?>(null)
+
+    init {
+        bgScope.launch {
+            val options = manager.getGridOptions()
+            _gridOptions.value = options
+        }
+    }
+
+    val gridOptions: StateFlow<List<GridOptionModel>?> = _gridOptions.asStateFlow()
+
+    val selectedGridOption: Flow<GridOptionModel?> =
+        gridOptions.map { gridOptions -> gridOptions?.firstOrNull { it.isCurrent } }
+
+    suspend fun applySelectedOption(key: String) =
+        withContext(bgDispatcher) {
+            manager.applyGridOption(key)
+            // After applying new grid option, we should query and update the grid options again.
+            _gridOptions.value = manager.getGridOptions()
+        }
+}
diff --git a/src/com/android/customization/picker/grid/domain/interactor/GridInteractor2.kt b/src/com/android/customization/picker/grid/domain/interactor/GridInteractor2.kt
new file mode 100644
index 0000000..30c87d8
--- /dev/null
+++ b/src/com/android/customization/picker/grid/domain/interactor/GridInteractor2.kt
@@ -0,0 +1,37 @@
+/*
+ * 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.picker.grid.domain.interactor
+
+import com.android.customization.picker.grid.data.repository.GridRepository2
+import javax.inject.Inject
+import javax.inject.Singleton
+
+@Singleton
+class GridInteractor2
+@Inject
+constructor(
+    private val repository: GridRepository2,
+) {
+    suspend fun isGridOptionAvailable(): Boolean = repository.isGridOptionAvailable()
+
+    val gridOptions = repository.gridOptions
+
+    val selectedGridOption = repository.selectedGridOption
+
+    suspend fun applySelectedOption(key: String) = repository.applySelectedOption(key)
+}
diff --git a/src/com/android/wallpaper/customization/ui/binder/ShapeAndGridFloatingSheetBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ShapeAndGridFloatingSheetBinder.kt
new file mode 100644
index 0000000..f03bc9b
--- /dev/null
+++ b/src/com/android/wallpaper/customization/ui/binder/ShapeAndGridFloatingSheetBinder.kt
@@ -0,0 +1,107 @@
+/*
+ * 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.content.Context
+import android.view.View
+import android.widget.ImageView
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.LifecycleOwner
+import androidx.lifecycle.lifecycleScope
+import androidx.lifecycle.repeatOnLifecycle
+import androidx.recyclerview.widget.LinearLayoutManager
+import androidx.recyclerview.widget.RecyclerView
+import com.android.customization.picker.common.ui.view.SingleRowListItemSpacing
+import com.android.customization.picker.grid.ui.binder.GridIconViewBinder
+import com.android.customization.picker.grid.ui.viewmodel.GridIconViewModel
+import com.android.wallpaper.R
+import com.android.wallpaper.customization.ui.viewmodel.ShapeAndGridPickerViewModel
+import com.android.wallpaper.picker.option.ui.adapter.OptionItemAdapter
+import com.android.wallpaper.picker.option.ui.binder.OptionItemBinder
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.launch
+
+object ShapeAndGridFloatingSheetBinder {
+
+    fun bind(
+        view: View,
+        viewModel: ShapeAndGridPickerViewModel,
+        lifecycleOwner: LifecycleOwner,
+        backgroundDispatcher: CoroutineDispatcher,
+    ) {
+        val adapter = createOptionItemAdapter(view.context, lifecycleOwner, backgroundDispatcher)
+        val gridOptionList =
+            view.requireViewById<RecyclerView>(R.id.options).also {
+                it.initGridOptionList(view.context, adapter)
+            }
+
+        lifecycleOwner.lifecycleScope.launch {
+            lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
+                launch { viewModel.optionItems.collect { options -> adapter.setItems(options) } }
+            }
+        }
+    }
+
+    private fun createOptionItemAdapter(
+        context: Context,
+        lifecycleOwner: LifecycleOwner,
+        backgroundDispatcher: CoroutineDispatcher,
+    ): OptionItemAdapter<GridIconViewModel> =
+        OptionItemAdapter(
+            layoutResourceId = com.android.themepicker.R.layout.grid_option,
+            lifecycleOwner = lifecycleOwner,
+            backgroundDispatcher = backgroundDispatcher,
+            foregroundTintSpec =
+                OptionItemBinder.TintSpec(
+                    selectedColor = context.getColor(R.color.system_on_surface),
+                    unselectedColor = context.getColor(R.color.system_on_surface),
+                ),
+            bindIcon = { foregroundView: View, gridIcon: GridIconViewModel ->
+                val imageView = foregroundView as? ImageView
+                imageView?.let { GridIconViewBinder.bind(imageView, gridIcon) }
+            }
+        )
+
+    private fun RecyclerView.initGridOptionList(
+        context: Context,
+        adapter: OptionItemAdapter<GridIconViewModel>,
+    ) {
+        apply {
+            this.layoutManager =
+                LinearLayoutManager(
+                    context,
+                    RecyclerView.HORIZONTAL,
+                    false,
+                )
+            addItemDecoration(
+                SingleRowListItemSpacing(
+                    edgeItemSpacePx =
+                        context.resources.getDimensionPixelSize(
+                            com.android.themepicker.R.dimen
+                                .floating_sheet_content_horizontal_padding
+                        ),
+                    itemHorizontalSpacePx =
+                        context.resources.getDimensionPixelSize(
+                            com.android.themepicker.R.dimen
+                                .floating_sheet_list_item_horizontal_space
+                        ),
+                )
+            )
+            this.adapter = adapter
+        }
+    }
+}
diff --git a/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt
index 02bcc03..d66e8a6 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt
@@ -19,11 +19,13 @@
 import android.view.View
 import android.widget.ImageView
 import android.widget.TextView
+import androidx.core.content.ContextCompat
 import androidx.core.view.isVisible
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
 import androidx.lifecycle.lifecycleScope
 import androidx.lifecycle.repeatOnLifecycle
+import com.android.customization.picker.grid.ui.binder.GridIconViewBinder
 import com.android.themepicker.R
 import com.android.wallpaper.customization.ui.util.ThemePickerCustomizationOptionUtil.ThemePickerHomeCustomizationOption
 import com.android.wallpaper.customization.ui.util.ThemePickerCustomizationOptionUtil.ThemePickerLockCustomizationOption
@@ -37,6 +39,7 @@
 import com.android.wallpaper.picker.customization.ui.viewmodel.CustomizationPickerViewModel2
 import javax.inject.Inject
 import javax.inject.Singleton
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
 
 @Singleton
@@ -91,6 +94,15 @@
                 .find { it.first == ThemePickerHomeCustomizationOption.COLORS }
                 ?.second
 
+        val optionShapeAndGrid =
+            homeScreenCustomizationOptionEntries
+                .find { it.first == ThemePickerHomeCustomizationOption.APP_GRID }
+                ?.second
+        val optionShapeAndGridDescription =
+            optionShapeAndGrid?.findViewById<TextView>(R.id.option_entry_app_grid_description)
+        val optionShapeAndGridIcon =
+            optionShapeAndGrid?.findViewById<ImageView>(R.id.option_entry_app_grid_icon)
+
         val optionsViewModel =
             viewModel.customizationOptionsViewModel as ThemePickerCustomizationOptionsViewModel
         lifecycleOwner.lifecycleScope.launch {
@@ -143,6 +155,36 @@
                         optionColors?.setOnClickListener { _ -> it?.invoke() }
                     }
                 }
+
+                launch {
+                    optionsViewModel.onCustomizeShapeAndGridClicked.collect {
+                        optionShapeAndGrid?.setOnClickListener { _ -> it?.invoke() }
+                    }
+                }
+
+                launch {
+                    optionsViewModel.shapeAndGridPickerViewModel.selectedGridOption.collect {
+                        gridOption ->
+                        optionShapeAndGridDescription?.let {
+                            TextViewBinder.bind(it, gridOption.text)
+                        }
+                        gridOption.payload?.let { gridIconViewModel ->
+                            optionShapeAndGridIcon?.let {
+                                GridIconViewBinder.bind(
+                                    view = it,
+                                    viewModel = gridIconViewModel,
+                                )
+                            }
+                            // TODO(b/363018910): Use ColorUpdateBinder to update color
+                            optionShapeAndGridIcon?.setColorFilter(
+                                ContextCompat.getColor(
+                                    view.context,
+                                    com.android.wallpaper.R.color.system_on_surface_variant
+                                )
+                            )
+                        }
+                    }
+                }
             }
         }
 
@@ -178,5 +220,16 @@
                     lifecycleOwner,
                 )
             }
+
+        customizationOptionFloatingSheetViewMap
+            ?.get(ThemePickerHomeCustomizationOption.APP_GRID)
+            ?.let {
+                ShapeAndGridFloatingSheetBinder.bind(
+                    it,
+                    optionsViewModel.shapeAndGridPickerViewModel,
+                    lifecycleOwner,
+                    Dispatchers.IO,
+                )
+            }
     }
 }
diff --git a/src/com/android/wallpaper/customization/ui/util/ThemePickerCustomizationOptionUtil.kt b/src/com/android/wallpaper/customization/ui/util/ThemePickerCustomizationOptionUtil.kt
index 5e7f568..2ff62c9 100644
--- a/src/com/android/wallpaper/customization/ui/util/ThemePickerCustomizationOptionUtil.kt
+++ b/src/com/android/wallpaper/customization/ui/util/ThemePickerCustomizationOptionUtil.kt
@@ -171,6 +171,15 @@
                     )
                     .also { bottomSheetContainer.addView(it) }
             )
+            put(
+                ThemePickerHomeCustomizationOption.APP_GRID,
+                inflateFloatingSheet(
+                        ThemePickerHomeCustomizationOption.APP_GRID,
+                        bottomSheetContainer,
+                        layoutInflater,
+                    )
+                    .also { bottomSheetContainer.addView(it) }
+            )
         }
     }
 
@@ -183,6 +192,7 @@
             ThemePickerLockCustomizationOption.CLOCK -> R.layout.floating_sheet_clock
             ThemePickerLockCustomizationOption.SHORTCUTS -> R.layout.floating_sheet_shortcut
             ThemePickerHomeCustomizationOption.COLORS -> R.layout.floating_sheet_colors
+            ThemePickerHomeCustomizationOption.APP_GRID -> R.layout.floating_sheet_shape_and_grid
             else ->
                 throw IllegalStateException(
                     "Customization option $option does not have a bottom sheet view"
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ShapeAndGridPickerViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ShapeAndGridPickerViewModel.kt
new file mode 100644
index 0000000..0fdc0b0
--- /dev/null
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ShapeAndGridPickerViewModel.kt
@@ -0,0 +1,129 @@
+/*
+ * 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 android.content.res.Resources
+import com.android.customization.model.ResourceConstants
+import com.android.customization.model.grid.GridOptionModel
+import com.android.customization.picker.grid.domain.interactor.GridInteractor2
+import com.android.customization.picker.grid.ui.viewmodel.GridIconViewModel
+import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
+import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
+import dagger.hilt.android.qualifiers.ApplicationContext
+import dagger.hilt.android.scopes.ViewModelScoped
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.filterNotNull
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.launch
+
+class ShapeAndGridPickerViewModel
+@AssistedInject
+constructor(
+    @ApplicationContext private val context: Context,
+    interactor: GridInteractor2,
+    @Assisted private val viewModelScope: CoroutineScope,
+) {
+    // The currently-set system grid option
+    val selectedGridOption =
+        interactor.selectedGridOption.filterNotNull().map { toOptionItemViewModel(it) }
+
+    private val overrideGridOptionKey = MutableStateFlow<String?>(null)
+
+    val optionItems: Flow<List<OptionItemViewModel<GridIconViewModel>>> =
+        interactor.gridOptions.filterNotNull().map { gridOptions ->
+            gridOptions.map { toOptionItemViewModel(it) }
+        }
+
+    val onApply: Flow<(() -> Unit)?> =
+        combine(selectedGridOption, overrideGridOptionKey) {
+            selectedGridOption,
+            overrideGridOptionKey ->
+            if (
+                overrideGridOptionKey == null ||
+                    overrideGridOptionKey == selectedGridOption.key.value
+            ) {
+                null
+            } else {
+                { viewModelScope.launch { interactor.applySelectedOption(overrideGridOptionKey) } }
+            }
+        }
+
+    val isOnApplyEnabled: Flow<Boolean> = onApply.map { it != null }
+
+    private fun toOptionItemViewModel(
+        option: GridOptionModel
+    ): OptionItemViewModel<GridIconViewModel> {
+        val iconShapePath =
+            context.resources.getString(
+                Resources.getSystem()
+                    .getIdentifier(
+                        ResourceConstants.CONFIG_ICON_MASK,
+                        "string",
+                        ResourceConstants.ANDROID_PACKAGE,
+                    )
+            )
+        val isSelected =
+            overrideGridOptionKey
+                .map {
+                    if (it == null) {
+                        option.isCurrent
+                    } else {
+                        it == option.key
+                    }
+                }
+                .stateIn(
+                    scope = viewModelScope,
+                    started = SharingStarted.Eagerly,
+                    initialValue = false,
+                )
+
+        return OptionItemViewModel(
+            key = MutableStateFlow(option.key),
+            payload =
+                GridIconViewModel(
+                    columns = option.cols,
+                    rows = option.rows,
+                    path = iconShapePath,
+                ),
+            text = Text.Loaded(option.title),
+            isSelected = isSelected,
+            onClicked =
+                isSelected.map {
+                    if (!it) {
+                        { overrideGridOptionKey.value = option.key }
+                    } else {
+                        null
+                    }
+                },
+        )
+    }
+
+    @ViewModelScoped
+    @AssistedFactory
+    interface Factory {
+        fun create(viewModelScope: CoroutineScope): ShapeAndGridPickerViewModel
+    }
+}
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
index 01c1c80..4576e0c 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
@@ -35,6 +35,7 @@
     keyguardQuickAffordancePickerViewModel2Factory: KeyguardQuickAffordancePickerViewModel2.Factory,
     colorPickerViewModel2Factory: ColorPickerViewModel2.Factory,
     clockPickerViewModelFactory: ClockPickerViewModel.Factory,
+    shapeAndGridPickerViewModelFactory: ShapeAndGridPickerViewModel.Factory,
     @Assisted private val viewModelScope: CoroutineScope,
 ) : CustomizationOptionsViewModel {
 
@@ -45,6 +46,8 @@
     val keyguardQuickAffordancePickerViewModel2 =
         keyguardQuickAffordancePickerViewModel2Factory.create(viewModelScope = viewModelScope)
     val colorPickerViewModel2 = colorPickerViewModel2Factory.create(viewModelScope = viewModelScope)
+    val shapeAndGridPickerViewModel =
+        shapeAndGridPickerViewModelFactory.create(viewModelScope = viewModelScope)
 
     override val selectedOption = defaultCustomizationOptionsViewModel.selectedOption
 
@@ -93,6 +96,20 @@
             }
         }
 
+    val onCustomizeShapeAndGridClicked: Flow<(() -> Unit)?> =
+        selectedOption.map {
+            if (it == null) {
+                {
+                    defaultCustomizationOptionsViewModel.selectOption(
+                        ThemePickerCustomizationOptionUtil.ThemePickerHomeCustomizationOption
+                            .APP_GRID
+                    )
+                }
+            } else {
+                null
+            }
+        }
+
     @ViewModelScoped
     @AssistedFactory
     interface Factory : CustomizationOptionsViewModelFactory {
diff --git a/src/com/android/wallpaper/picker/di/modules/ThemePickerSharedAppModule.kt b/src/com/android/wallpaper/picker/di/modules/ThemePickerSharedAppModule.kt
new file mode 100644
index 0000000..0b32196
--- /dev/null
+++ b/src/com/android/wallpaper/picker/di/modules/ThemePickerSharedAppModule.kt
@@ -0,0 +1,34 @@
+/*
+ * 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.di.modules
+
+import com.android.customization.model.grid.DefaultGridOptionsManager
+import com.android.customization.model.grid.GridOptionsManager2
+import dagger.Binds
+import dagger.Module
+import dagger.hilt.InstallIn
+import dagger.hilt.components.SingletonComponent
+import javax.inject.Singleton
+
+@Module
+@InstallIn(SingletonComponent::class)
+abstract class ThemePickerSharedAppModule {
+
+    @Binds
+    @Singleton
+    abstract fun bindGridOptionsManager2(impl: DefaultGridOptionsManager): GridOptionsManager2
+}
diff --git a/tests/common/src/com/android/customization/model/grid/FakeGridOptionsManager.kt b/tests/common/src/com/android/customization/model/grid/FakeGridOptionsManager.kt
new file mode 100644
index 0000000..cc23981
--- /dev/null
+++ b/tests/common/src/com/android/customization/model/grid/FakeGridOptionsManager.kt
@@ -0,0 +1,57 @@
+/*
+ * 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
+
+import javax.inject.Inject
+import javax.inject.Singleton
+
+@Singleton
+class FakeGridOptionsManager @Inject constructor() : GridOptionsManager2 {
+
+    var isGridOptionAvailable: Boolean = true
+
+    private var gridOptions: List<GridOptionModel>? = DEFAULT_GRID_OPTION_LIST
+
+    override suspend fun isGridOptionAvailable(): Boolean = isGridOptionAvailable
+
+    override suspend fun getGridOptions(): List<GridOptionModel>? = gridOptions
+
+    override fun applyGridOption(gridName: String): Int {
+        gridOptions = gridOptions?.map { it.copy(isCurrent = it.key == gridName) }
+        return 0
+    }
+
+    companion object {
+        val DEFAULT_GRID_OPTION_LIST =
+            listOf(
+                GridOptionModel(
+                    key = "normal",
+                    title = "5x5",
+                    isCurrent = true,
+                    rows = 5,
+                    cols = 5,
+                ),
+                GridOptionModel(
+                    key = "practical",
+                    title = "4x5",
+                    isCurrent = false,
+                    rows = 5,
+                    cols = 4,
+                ),
+            )
+    }
+}
diff --git a/tests/common/src/com/android/wallpaper/di/modules/ThemePickerSharedAppTestModule.kt b/tests/common/src/com/android/wallpaper/di/modules/ThemePickerSharedAppTestModule.kt
new file mode 100644
index 0000000..7781d4e
--- /dev/null
+++ b/tests/common/src/com/android/wallpaper/di/modules/ThemePickerSharedAppTestModule.kt
@@ -0,0 +1,38 @@
+/*
+ * 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.di.modules
+
+import com.android.customization.model.grid.FakeGridOptionsManager
+import com.android.customization.model.grid.GridOptionsManager2
+import com.android.wallpaper.picker.di.modules.ThemePickerSharedAppModule
+import dagger.Binds
+import dagger.Module
+import dagger.hilt.components.SingletonComponent
+import dagger.hilt.testing.TestInstallIn
+import javax.inject.Singleton
+
+@Module
+@TestInstallIn(
+    components = [SingletonComponent::class],
+    replaces = [ThemePickerSharedAppModule::class]
+)
+abstract class ThemePickerSharedAppTestModule {
+
+    @Binds
+    @Singleton
+    abstract fun bindGridOptionsManager2(impl: FakeGridOptionsManager): GridOptionsManager2
+}
diff --git a/tests/robotests/src/com/android/customization/picker/grid/data/repository/GridRepository2Test.kt b/tests/robotests/src/com/android/customization/picker/grid/data/repository/GridRepository2Test.kt
new file mode 100644
index 0000000..404f08b
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/picker/grid/data/repository/GridRepository2Test.kt
@@ -0,0 +1,123 @@
+/*
+ * 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.picker.grid.data.repository
+
+import androidx.test.filters.SmallTest
+import com.android.customization.model.grid.FakeGridOptionsManager
+import com.android.wallpaper.picker.di.modules.BackgroundDispatcher
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import dagger.hilt.android.testing.HiltAndroidRule
+import dagger.hilt.android.testing.HiltAndroidTest
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.resetMain
+import kotlinx.coroutines.test.runTest
+import org.junit.After
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
+
+@HiltAndroidTest
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(RobolectricTestRunner::class)
+class GridRepository2Test {
+
+    @get:Rule var hiltRule = HiltAndroidRule(this)
+    @Inject lateinit var gridOptionsManager: FakeGridOptionsManager
+    @Inject lateinit var testScope: TestScope
+    @BackgroundDispatcher @Inject lateinit var bgScope: CoroutineScope
+    @BackgroundDispatcher @Inject lateinit var bgDispatcher: CoroutineDispatcher
+
+    private lateinit var underTest: GridRepository2
+
+    @Before
+    fun setUp() {
+        hiltRule.inject()
+        underTest =
+            GridRepository2(
+                manager = gridOptionsManager,
+                bgScope = bgScope,
+                bgDispatcher = bgDispatcher,
+            )
+    }
+
+    @After
+    fun tearDown() {
+        Dispatchers.resetMain()
+    }
+
+    @Test
+    fun isGridOptionAvailable_false() =
+        testScope.runTest {
+            gridOptionsManager.isGridOptionAvailable = false
+            assertThat(underTest.isGridOptionAvailable()).isFalse()
+        }
+
+    @Test
+    fun isGridOptionAvailable_true() =
+        testScope.runTest {
+            gridOptionsManager.isGridOptionAvailable = true
+            assertThat(underTest.isGridOptionAvailable()).isTrue()
+        }
+
+    @Test
+    fun gridOptions_default() =
+        testScope.runTest {
+            val gridOptions = collectLastValue(underTest.gridOptions)
+            assertThat(gridOptions()).isEqualTo(FakeGridOptionsManager.DEFAULT_GRID_OPTION_LIST)
+        }
+
+    @Test
+    fun selectedGridOption_default() =
+        testScope.runTest {
+            val selectedGridOption = collectLastValue(underTest.selectedGridOption)
+            assertThat(selectedGridOption())
+                .isEqualTo(FakeGridOptionsManager.DEFAULT_GRID_OPTION_LIST[0])
+        }
+
+    @Test
+    fun gridOptions_shouldUpdateAfterApplyGridOption() =
+        testScope.runTest {
+            val gridOptions = collectLastValue(underTest.gridOptions)
+            underTest.applySelectedOption("practical")
+            assertThat(gridOptions())
+                .isEqualTo(
+                    FakeGridOptionsManager.DEFAULT_GRID_OPTION_LIST.map {
+                        it.copy(isCurrent = it.key == "practical")
+                    }
+                )
+        }
+
+    @Test
+    fun selectedGridOption_shouldUpdateAfterApplyGridOption() =
+        testScope.runTest {
+            val selectedGridOption = collectLastValue(underTest.selectedGridOption)
+            underTest.applySelectedOption("practical")
+            assertThat(selectedGridOption())
+                .isEqualTo(
+                    FakeGridOptionsManager.DEFAULT_GRID_OPTION_LIST[1].copy(isCurrent = true)
+                )
+        }
+}
diff --git a/tests/robotests/src/com/android/customization/picker/grid/domain/interactor/GridInteractor2Test.kt b/tests/robotests/src/com/android/customization/picker/grid/domain/interactor/GridInteractor2Test.kt
new file mode 100644
index 0000000..bfbe282
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/picker/grid/domain/interactor/GridInteractor2Test.kt
@@ -0,0 +1,115 @@
+/*
+ * 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.picker.grid.domain.interactor
+
+import androidx.test.filters.SmallTest
+import com.android.customization.model.grid.FakeGridOptionsManager
+import com.android.customization.picker.grid.data.repository.GridRepository2
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import dagger.hilt.android.testing.HiltAndroidRule
+import dagger.hilt.android.testing.HiltAndroidTest
+import javax.inject.Inject
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.resetMain
+import kotlinx.coroutines.test.runTest
+import org.junit.After
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
+
+@HiltAndroidTest
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(RobolectricTestRunner::class)
+class GridInteractor2Test {
+
+    @get:Rule var hiltRule = HiltAndroidRule(this)
+    @Inject lateinit var gridOptionsManager: FakeGridOptionsManager
+    @Inject lateinit var repository: GridRepository2
+    @Inject lateinit var testScope: TestScope
+
+    private lateinit var underTest: GridInteractor2
+
+    @Before
+    fun setUp() {
+        hiltRule.inject()
+        underTest = GridInteractor2(repository)
+    }
+
+    @After
+    fun tearDown() {
+        Dispatchers.resetMain()
+    }
+
+    @Test
+    fun isGridOptionAvailable_false() =
+        testScope.runTest {
+            gridOptionsManager.isGridOptionAvailable = false
+            assertThat(underTest.isGridOptionAvailable()).isFalse()
+        }
+
+    @Test
+    fun isGridOptionAvailable_true() =
+        testScope.runTest {
+            gridOptionsManager.isGridOptionAvailable = true
+            assertThat(underTest.isGridOptionAvailable()).isTrue()
+        }
+
+    @Test
+    fun gridOptions_default() =
+        testScope.runTest {
+            val gridOptions = collectLastValue(underTest.gridOptions)
+            assertThat(gridOptions()).isEqualTo(FakeGridOptionsManager.DEFAULT_GRID_OPTION_LIST)
+        }
+
+    @Test
+    fun selectedGridOption_default() =
+        testScope.runTest {
+            val selectedGridOption = collectLastValue(underTest.selectedGridOption)
+            assertThat(selectedGridOption())
+                .isEqualTo(FakeGridOptionsManager.DEFAULT_GRID_OPTION_LIST[0])
+        }
+
+    @Test
+    fun gridOptions_shouldUpdateAfterApplyGridOption() =
+        testScope.runTest {
+            val gridOptions = collectLastValue(underTest.gridOptions)
+            underTest.applySelectedOption("practical")
+            assertThat(gridOptions())
+                .isEqualTo(
+                    FakeGridOptionsManager.DEFAULT_GRID_OPTION_LIST.map {
+                        it.copy(isCurrent = it.key == "practical")
+                    }
+                )
+        }
+
+    @Test
+    fun selectedGridOption_shouldUpdateAfterApplyGridOption() =
+        testScope.runTest {
+            val selectedGridOption = collectLastValue(underTest.selectedGridOption)
+            underTest.applySelectedOption("practical")
+            assertThat(selectedGridOption())
+                .isEqualTo(
+                    FakeGridOptionsManager.DEFAULT_GRID_OPTION_LIST[1].copy(isCurrent = true)
+                )
+        }
+}
diff --git a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ShapeAndGridPickerViewModelTest.kt b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ShapeAndGridPickerViewModelTest.kt
new file mode 100644
index 0000000..02d3ce7
--- /dev/null
+++ b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ShapeAndGridPickerViewModelTest.kt
@@ -0,0 +1,196 @@
+/*
+ * 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 android.content.res.Resources
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.filters.SmallTest
+import com.android.customization.model.ResourceConstants
+import com.android.customization.model.grid.FakeGridOptionsManager
+import com.android.customization.picker.grid.domain.interactor.GridInteractor2
+import com.android.customization.picker.grid.ui.viewmodel.GridIconViewModel
+import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
+import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import dagger.hilt.android.qualifiers.ApplicationContext
+import dagger.hilt.android.testing.HiltAndroidRule
+import dagger.hilt.android.testing.HiltAndroidTest
+import javax.inject.Inject
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.resetMain
+import kotlinx.coroutines.test.runTest
+import org.junit.After
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
+
+@HiltAndroidTest
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(RobolectricTestRunner::class)
+class ShapeAndGridPickerViewModelTest {
+
+    @get:Rule var hiltRule = HiltAndroidRule(this)
+    @Inject lateinit var testScope: TestScope
+    @Inject lateinit var gridOptionsManager: FakeGridOptionsManager
+    @Inject lateinit var interactor: GridInteractor2
+    @Inject @ApplicationContext lateinit var appContext: Context
+
+    private val iconShapePath =
+        ApplicationProvider.getApplicationContext<Context>()
+            .resources
+            .getString(
+                Resources.getSystem()
+                    .getIdentifier(
+                        ResourceConstants.CONFIG_ICON_MASK,
+                        "string",
+                        ResourceConstants.ANDROID_PACKAGE,
+                    )
+            )
+
+    private lateinit var underTest: ShapeAndGridPickerViewModel
+
+    @Before
+    fun setUp() {
+        hiltRule.inject()
+        underTest = ShapeAndGridPickerViewModel(appContext, interactor, testScope.backgroundScope)
+    }
+
+    @After
+    fun tearDown() {
+        Dispatchers.resetMain()
+    }
+
+    @Test
+    fun selectedGridOption() =
+        testScope.runTest {
+            val selectedGridOption = collectLastValue(underTest.selectedGridOption)
+
+            assertOptionItem(
+                optionItem = selectedGridOption(),
+                key = "normal",
+                payload = GridIconViewModel(5, 5, iconShapePath),
+                text = Text.Loaded("5x5"),
+                isTextUserVisible = true,
+                isSelected = true,
+                isEnabled = true,
+            )
+        }
+
+    @Test
+    fun selectedGridOption_shouldUpdate_afterOnApply() =
+        testScope.runTest {
+            val selectedGridOption = collectLastValue(underTest.selectedGridOption)
+            val optionItems = collectLastValue(underTest.optionItems)
+            val onApply = collectLastValue(underTest.onApply)
+            val onPracticalOptionClick =
+                optionItems()?.get(1)?.onClicked?.let { collectLastValue(it) }
+            checkNotNull(onPracticalOptionClick)
+
+            onPracticalOptionClick()?.invoke()
+            onApply()?.invoke()
+
+            assertOptionItem(
+                optionItem = selectedGridOption(),
+                key = "practical",
+                payload = GridIconViewModel(4, 5, iconShapePath),
+                text = Text.Loaded("4x5"),
+                isTextUserVisible = true,
+                isSelected = true,
+                isEnabled = true,
+            )
+        }
+
+    @Test
+    fun optionItems() =
+        testScope.runTest {
+            val optionItems = collectLastValue(underTest.optionItems)
+
+            assertOptionItem(
+                optionItem = optionItems()?.get(0),
+                key = "normal",
+                payload = GridIconViewModel(5, 5, iconShapePath),
+                text = Text.Loaded("5x5"),
+                isTextUserVisible = true,
+                isSelected = true,
+                isEnabled = true,
+            )
+            assertOptionItem(
+                optionItem = optionItems()?.get(1),
+                key = "practical",
+                payload = GridIconViewModel(4, 5, iconShapePath),
+                text = Text.Loaded("4x5"),
+                isTextUserVisible = true,
+                isSelected = false,
+                isEnabled = true,
+            )
+        }
+
+    @Test
+    fun optionItems_whenClickOnPracticalOption() =
+        testScope.runTest {
+            val optionItems = collectLastValue(underTest.optionItems)
+            val onPracticalOptionClick =
+                optionItems()?.get(1)?.onClicked?.let { collectLastValue(it) }
+            checkNotNull(onPracticalOptionClick)
+
+            onPracticalOptionClick()?.invoke()
+
+            assertOptionItem(
+                optionItem = optionItems()?.get(0),
+                key = "normal",
+                payload = GridIconViewModel(5, 5, iconShapePath),
+                text = Text.Loaded("5x5"),
+                isTextUserVisible = true,
+                isSelected = false,
+                isEnabled = true,
+            )
+            assertOptionItem(
+                optionItem = optionItems()?.get(1),
+                key = "practical",
+                payload = GridIconViewModel(4, 5, iconShapePath),
+                text = Text.Loaded("4x5"),
+                isTextUserVisible = true,
+                isSelected = true,
+                isEnabled = true,
+            )
+        }
+
+    private fun assertOptionItem(
+        optionItem: OptionItemViewModel<GridIconViewModel>?,
+        key: String,
+        payload: GridIconViewModel?,
+        text: Text,
+        isTextUserVisible: Boolean,
+        isSelected: Boolean,
+        isEnabled: Boolean,
+    ) {
+        checkNotNull(optionItem)
+        assertThat(optionItem.key.value).isEqualTo(key)
+        assertThat(optionItem.text).isEqualTo(text)
+        assertThat(optionItem.payload).isEqualTo(payload)
+        assertThat(optionItem.isTextUserVisible).isEqualTo(isTextUserVisible)
+        assertThat(optionItem.isSelected.value).isEqualTo(isSelected)
+        assertThat(optionItem.isEnabled).isEqualTo(isEnabled)
+    }
+}