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 6269fee..0166cc0 100644
--- a/tests/common/src/com/android/customization/model/grid/FakeShapeGridManager.kt
+++ b/tests/common/src/com/android/customization/model/grid/FakeShapeGridManager.kt
@@ -95,13 +95,6 @@
                     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 =
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..5fa1545 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
@@ -104,11 +104,14 @@
     fun selectedShapeOption_shouldUpdateAfterApplyShapeGridOption() =
         testScope.runTest {
             val selectedShapeOption = collectLastValue(underTest.selectedShapeOption)
+            val expectedShapeKey = "circle"
+            val expectedShapeOption =
+                FakeShapeGridManager.DEFAULT_SHAPE_OPTION_LIST.first { it.key == expectedShapeKey }
+                    .copy(isCurrent = true)
 
-            underTest.applySelectedOption("circle", "practical")
+            underTest.applySelectedOption(expectedShapeKey, "practical")
 
-            assertThat(selectedShapeOption())
-                .isEqualTo(FakeShapeGridManager.DEFAULT_SHAPE_OPTION_LIST[4].copy(isCurrent = true))
+            assertThat(selectedShapeOption()).isEqualTo(expectedShapeOption)
         }
 
     @Test
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..3886690 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
@@ -96,11 +96,14 @@
     fun selectedShapeOption_shouldUpdateAfterApplyGridOption() =
         testScope.runTest {
             val selectedShapeOption = collectLastValue(underTest.selectedShapeOption)
+            val expectedShapeKey = "circle"
+            val expectedShapeOption =
+                FakeShapeGridManager.DEFAULT_SHAPE_OPTION_LIST.first { it.key == expectedShapeKey }
+                    .copy(isCurrent = true)
 
-            underTest.applySelectedOption("circle", "practical")
+            underTest.applySelectedOption(expectedShapeKey, "practical")
 
-            assertThat(selectedShapeOption())
-                .isEqualTo(FakeShapeGridManager.DEFAULT_SHAPE_OPTION_LIST[4].copy(isCurrent = true))
+            assertThat(selectedShapeOption()).isEqualTo(expectedShapeOption)
         }
 
     @Test
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 0bf5125..15d1487 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
@@ -116,8 +116,8 @@
         testScope.runTest {
             val shapeOptions = collectLastValue(underTest.shapeOptions)
             val previewingShapeKey = collectLastValue(underTest.previewingShapeKey)
-            val onCircleOptionClicked =
-                shapeOptions()?.get(4)?.onClicked?.let { collectLastValue(it) }
+            val circleOption = shapeOptions()?.firstOrNull { it.key.value == "circle" }
+            val onCircleOptionClicked = circleOption?.onClicked?.let { collectLastValue(it) }
             checkNotNull(onCircleOptionClicked)
 
             onCircleOptionClicked()?.invoke()