diff --git a/src/com/android/customization/model/color/ColorProvider.kt b/src/com/android/customization/model/color/ColorProvider.kt
index ef5e3d6..5c2b891 100644
--- a/src/com/android/customization/model/color/ColorProvider.kt
+++ b/src/com/android/customization/model/color/ColorProvider.kt
@@ -117,16 +117,19 @@
         }
 
         scope.launch {
+            // Wait for the previous preset color loading job to finish before evaluating whether to
+            // start a new one
             loaderJob?.join()
             if (presetColorBundles == null || reload) {
-                try {
-                    loaderJob = launch { loadPreset(isNewPickerUi) }
-                } catch (e: Throwable) {
-                    colorsAvailable = false
-                    callback?.onError(e)
-                    return@launch
+                loaderJob = launch {
+                    try {
+                        loadPreset(isNewPickerUi)
+                        callback?.onOptionsLoaded(buildFinalList(isNewPickerUi))
+                    } catch (e: Throwable) {
+                        colorsAvailable = false
+                        callback?.onError(e)
+                    }
                 }
-                callback?.onOptionsLoaded(buildFinalList(isNewPickerUi))
             } else {
                 callback?.onOptionsLoaded(buildFinalList(isNewPickerUi))
             }
diff --git a/src/com/android/customization/model/grid/DefaultShapeGridManager.kt b/src/com/android/customization/model/grid/DefaultShapeGridManager.kt
index 6f20dd3..f8b0c3a 100644
--- a/src/com/android/customization/model/grid/DefaultShapeGridManager.kt
+++ b/src/com/android/customization/model/grid/DefaultShapeGridManager.kt
@@ -18,7 +18,9 @@
 
 import android.content.ContentValues
 import android.content.Context
+import android.content.res.Resources
 import android.graphics.drawable.Drawable
+import android.util.Log
 import androidx.core.content.res.ResourcesCompat
 import com.android.wallpaper.R
 import com.android.wallpaper.picker.di.modules.BackgroundDispatcher
@@ -156,15 +158,26 @@
         )
     }
 
-    override fun getGridOptionDrawble(iconId: Int): Drawable? {
-        return ResourcesCompat.getDrawable(
-            context.packageManager.getResourcesForApplication(APP_RESOURCES_PACKAGE_NAME),
-            iconId,
-            /* theme = */ null,
-        )
+    override fun getGridOptionDrawable(iconId: Int): Drawable? {
+        try {
+            val drawable =
+                ResourcesCompat.getDrawable(
+                    context.packageManager.getResourcesForApplication(APP_RESOURCES_PACKAGE_NAME),
+                    iconId,
+                    /* theme = */ null,
+                )
+            return drawable
+        } catch (exception: Resources.NotFoundException) {
+            Log.w(
+                TAG,
+                "Unable to find drawable resource from package $APP_RESOURCES_PACKAGE_NAME with resource ID $iconId",
+            )
+            return null
+        }
     }
 
     companion object {
+        const val TAG = "DefaultShapeGridManager"
         const val SHAPE_OPTIONS: String = "shape_options"
         const val GRID_OPTIONS: String = "list_options"
         const val SHAPE_GRID: String = "default_grid"
diff --git a/src/com/android/customization/model/grid/ShapeGridManager.kt b/src/com/android/customization/model/grid/ShapeGridManager.kt
index 7a74b8d..7699783 100644
--- a/src/com/android/customization/model/grid/ShapeGridManager.kt
+++ b/src/com/android/customization/model/grid/ShapeGridManager.kt
@@ -26,5 +26,5 @@
 
     fun applyShapeGridOption(shapeKey: String, gridKey: String): Int
 
-    fun getGridOptionDrawble(iconId: Int): Drawable?
+    fun getGridOptionDrawable(iconId: Int): Drawable?
 }
diff --git a/src/com/android/customization/picker/clock/ui/view/ClockCarouselView.kt b/src/com/android/customization/picker/clock/ui/view/ClockCarouselView.kt
index 27bc42c..dccef6e 100644
--- a/src/com/android/customization/picker/clock/ui/view/ClockCarouselView.kt
+++ b/src/com/android/customization/picker/clock/ui/view/ClockCarouselView.kt
@@ -467,15 +467,16 @@
                 it.pivotY = it.height / 2F
             }
 
-            val controller = clockViewFactory.getController(clockId)
-            if (isMiddleView) {
-                clockScaleView.scaleX = 1f
-                clockScaleView.scaleY = 1f
-                controller.largeClock.animations.onPickerCarouselSwiping(1F)
-            } else {
-                clockScaleView.scaleX = clockViewScale
-                clockScaleView.scaleY = clockViewScale
-                controller.largeClock.animations.onPickerCarouselSwiping(0F)
+            clockViewFactory.getController(clockId)?.let { controller ->
+                if (isMiddleView) {
+                    clockScaleView.scaleX = 1f
+                    clockScaleView.scaleY = 1f
+                    controller.largeClock.animations.onPickerCarouselSwiping(1F)
+                } else {
+                    clockScaleView.scaleX = clockViewScale
+                    clockScaleView.scaleY = clockViewScale
+                    controller.largeClock.animations.onPickerCarouselSwiping(0F)
+                }
             }
         }
 
diff --git a/src/com/android/customization/picker/clock/ui/view/ThemePickerClockViewFactory.kt b/src/com/android/customization/picker/clock/ui/view/ThemePickerClockViewFactory.kt
index 73ebb0f..c48d810 100644
--- a/src/com/android/customization/picker/clock/ui/view/ThemePickerClockViewFactory.kt
+++ b/src/com/android/customization/picker/clock/ui/view/ThemePickerClockViewFactory.kt
@@ -43,7 +43,7 @@
 class ThemePickerClockViewFactory
 @Inject
 constructor(
-    activity: Activity,
+    private val activity: Activity,
     private val wallpaperManager: WallpaperManager,
     private val registry: ClockRegistry,
 ) : ClockViewFactory {
@@ -55,9 +55,9 @@
     private val clockControllers: ConcurrentHashMap<String, ClockController> = ConcurrentHashMap()
     private val smallClockFrames: HashMap<String, FrameLayout> = HashMap()
 
-    override fun getController(clockId: String): ClockController {
+    override fun getController(clockId: String): ClockController? {
         return clockControllers[clockId]
-            ?: initClockController(clockId).also { clockControllers[clockId] = it }
+            ?: initClockController(clockId)?.also { clockControllers[clockId] = it }
     }
 
     /**
@@ -66,10 +66,10 @@
      */
     override fun getLargeView(clockId: String): View {
         assert(!Flags.newCustomizationPickerUi())
-        return getController(clockId).largeClock.let {
+        return getController(clockId)?.largeClock?.let {
             it.animations.onPickerCarouselSwiping(1F)
             it.view
-        }
+        } ?: FrameLayout(activity)
     }
 
     /**
@@ -83,9 +83,9 @@
                 (layoutParams as FrameLayout.LayoutParams).topMargin = getSmallClockTopMargin()
                 (layoutParams as FrameLayout.LayoutParams).marginStart = getSmallClockStartPadding()
             }
-                ?: createSmallClockFrame().also {
-                    it.addView(getController(clockId).smallClock.view)
-                    smallClockFrames[clockId] = it
+                ?: createSmallClockFrame().also { frame ->
+                    getController(clockId)?.let { frame.addView(it.smallClock.view) }
+                    smallClockFrames[clockId] = frame
                 }
         smallClockFrame.translationX = 0F
         smallClockFrame.translationY = 0F
@@ -130,14 +130,14 @@
     }
 
     override fun updateColor(clockId: String, @ColorInt seedColor: Int?) {
-        getController(clockId).let {
+        getController(clockId)?.let {
             it.largeClock.run { events.onThemeChanged(theme.copy(seedColor = seedColor)) }
             it.smallClock.run { events.onThemeChanged(theme.copy(seedColor = seedColor)) }
         }
     }
 
     override fun updateFontAxes(clockId: String, settings: List<ClockFontAxisSetting>) {
-        getController(clockId).let { it.events.onFontAxesChanged(settings) }
+        getController(clockId)?.let { it.events.onFontAxesChanged(settings) }
     }
 
     override fun updateRegionDarkness() {
@@ -155,8 +155,8 @@
 
     override fun updateTimeFormat(clockId: String) {
         getController(clockId)
-            .events
-            .onTimeFormatChanged(android.text.format.DateFormat.is24HourFormat(appContext))
+            ?.events
+            ?.onTimeFormatChanged(android.text.format.DateFormat.is24HourFormat(appContext))
     }
 
     override fun registerTimeTicker(owner: LifecycleOwner) {
@@ -190,33 +190,32 @@
         }
     }
 
-    private fun initClockController(clockId: String): ClockController {
+    private fun initClockController(clockId: String): ClockController? {
         val isWallpaperDark = isLockscreenWallpaperDark()
-        val controller =
-            registry.createExampleClock(clockId).also { it?.initialize(isWallpaperDark, 0f, 0f) }
-        checkNotNull(controller)
+        return registry.createExampleClock(clockId)?.also { controller ->
+            controller.initialize(isWallpaperDark, 0f, 0f)
 
-        // Initialize large clock
-        controller.largeClock.events.onFontSettingChanged(
-            resources
-                .getDimensionPixelSize(
-                    com.android.systemui.customization.R.dimen.large_clock_text_size
-                )
-                .toFloat()
-        )
-        controller.largeClock.events.onTargetRegionChanged(getLargeClockRegion())
+            // Initialize large clock
+            controller.largeClock.events.onFontSettingChanged(
+                resources
+                    .getDimensionPixelSize(
+                        com.android.systemui.customization.R.dimen.large_clock_text_size
+                    )
+                    .toFloat()
+            )
+            controller.largeClock.events.onTargetRegionChanged(getLargeClockRegion())
 
-        // Initialize small clock
-        controller.smallClock.events.onFontSettingChanged(
-            resources
-                .getDimensionPixelSize(
-                    com.android.systemui.customization.R.dimen.small_clock_text_size
-                )
-                .toFloat()
-        )
-        controller.smallClock.events.onTargetRegionChanged(getSmallClockRegion())
-        controller.events.onWeatherDataChanged(WeatherData.getPlaceholderWeatherData())
-        return controller
+            // Initialize small clock
+            controller.smallClock.events.onFontSettingChanged(
+                resources
+                    .getDimensionPixelSize(
+                        com.android.systemui.customization.R.dimen.small_clock_text_size
+                    )
+                    .toFloat()
+            )
+            controller.smallClock.events.onTargetRegionChanged(getSmallClockRegion())
+            controller.events.onWeatherDataChanged(WeatherData.getPlaceholderWeatherData())
+        }
     }
 
     /**
diff --git a/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModel.kt b/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModel.kt
index 28f58c4..874b49d 100644
--- a/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModel.kt
+++ b/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModel.kt
@@ -32,6 +32,7 @@
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
@@ -57,14 +58,22 @@
             .mapLatest { allClocks ->
                 // Delay to avoid the case that the full list of clocks is not initiated.
                 delay(CLOCKS_EVENT_UPDATE_DELAY_MILLIS)
-                allClocks.map {
-                    val contentDescription =
-                        resources.getString(
-                            R.string.select_clock_action_description,
-                            clockViewFactory.getController(it.clockId).config.description
-                        )
-                    ClockCarouselItemViewModel(it.clockId, it.isSelected, contentDescription)
-                }
+                allClocks
+                    .map { model ->
+                        clockViewFactory.getController(model.clockId)?.let { clock ->
+                            val contentDescription =
+                                resources.getString(
+                                    R.string.select_clock_action_description,
+                                    clock.config.description,
+                                )
+                            ClockCarouselItemViewModel(
+                                model.clockId,
+                                model.isSelected,
+                                contentDescription,
+                            )
+                        }
+                    }
+                    .filterNotNull()
             }
             // makes sure that the operations above this statement are executed on I/O dispatcher
             // while parallelism limits the number of threads this can run on which makes sure that
@@ -126,6 +135,7 @@
             .mapNotNull { it }
 
     private var setSelectedClockJob: Job? = null
+
     fun setSelectedClock(clockId: String) {
         setSelectedClockJob?.cancel()
         setSelectedClockJob =
diff --git a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
index 43d510c..7e63834 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
@@ -33,7 +33,6 @@
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.combine
-import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.suspendCancellableCoroutine
 
 // TODO (b/262924623): refactor to remove dependency on ColorCustomizationManager & ColorOption
@@ -58,68 +57,61 @@
 
     override val colorOptions: Flow<Map<ColorType, List<ColorOptionModel>>> =
         combine(homeWallpaperColors, lockWallpaperColors) { homeColors, lockColors ->
-                homeColors to lockColors
-            }
-            .map { (homeColors, lockColors) ->
-                suspendCancellableCoroutine { continuation ->
-                    if (
-                        homeColors is WallpaperColorsModel.Loading ||
-                            lockColors is WallpaperColorsModel.Loading
-                    ) {
-                        continuation.resumeWith(
-                            Result.success(
-                                mapOf(
-                                    ColorType.WALLPAPER_COLOR to listOf(),
-                                    ColorType.PRESET_COLOR to listOf(),
-                                )
+            suspendCancellableCoroutine { continuation ->
+                if (
+                    homeColors is WallpaperColorsModel.Loading ||
+                        lockColors is WallpaperColorsModel.Loading
+                ) {
+                    continuation.resumeWith(
+                        Result.success(
+                            mapOf(
+                                ColorType.WALLPAPER_COLOR to listOf(),
+                                ColorType.PRESET_COLOR to listOf(),
                             )
                         )
-                        return@suspendCancellableCoroutine
-                    }
-                    val homeColorsLoaded = homeColors as WallpaperColorsModel.Loaded
-                    val lockColorsLoaded = lockColors as WallpaperColorsModel.Loaded
-                    colorManager.setWallpaperColors(
-                        homeColorsLoaded.colors,
-                        lockColorsLoaded.colors,
                     )
-                    colorManager.fetchOptions(
-                        object : CustomizationManager.OptionsFetchedListener<ColorOption?> {
-                            override fun onOptionsLoaded(options: MutableList<ColorOption?>?) {
-                                val wallpaperColorOptions: MutableList<ColorOptionModel> =
-                                    mutableListOf()
-                                val presetColorOptions: MutableList<ColorOptionModel> =
-                                    mutableListOf()
-                                options?.forEach { option ->
-                                    when ((option as ColorOptionImpl).type) {
-                                        ColorType.WALLPAPER_COLOR ->
-                                            wallpaperColorOptions.add(option.toModel())
-                                        ColorType.PRESET_COLOR ->
-                                            presetColorOptions.add(option.toModel())
-                                    }
-                                }
-                                continuation.resumeWith(
-                                    Result.success(
-                                        mapOf(
-                                            ColorType.WALLPAPER_COLOR to wallpaperColorOptions,
-                                            ColorType.PRESET_COLOR to presetColorOptions,
-                                        )
-                                    )
-                                )
-                            }
-
-                            override fun onError(throwable: Throwable?) {
-                                Log.e(TAG, "Error loading theme bundles", throwable)
-                                continuation.resumeWith(
-                                    Result.failure(
-                                        throwable ?: Throwable("Error loading theme bundles")
-                                    )
-                                )
-                            }
-                        },
-                        /* reload= */ false,
-                    )
+                    return@suspendCancellableCoroutine
                 }
+                val homeColorsLoaded = homeColors as WallpaperColorsModel.Loaded
+                val lockColorsLoaded = lockColors as WallpaperColorsModel.Loaded
+                colorManager.setWallpaperColors(homeColorsLoaded.colors, lockColorsLoaded.colors)
+                colorManager.fetchOptions(
+                    object : CustomizationManager.OptionsFetchedListener<ColorOption?> {
+                        override fun onOptionsLoaded(options: MutableList<ColorOption?>?) {
+                            val wallpaperColorOptions: MutableList<ColorOptionModel> =
+                                mutableListOf()
+                            val presetColorOptions: MutableList<ColorOptionModel> = mutableListOf()
+                            options?.forEach { option ->
+                                when ((option as ColorOptionImpl).type) {
+                                    ColorType.WALLPAPER_COLOR ->
+                                        wallpaperColorOptions.add(option.toModel())
+                                    ColorType.PRESET_COLOR ->
+                                        presetColorOptions.add(option.toModel())
+                                }
+                            }
+                            continuation.resumeWith(
+                                Result.success(
+                                    mapOf(
+                                        ColorType.WALLPAPER_COLOR to wallpaperColorOptions,
+                                        ColorType.PRESET_COLOR to presetColorOptions,
+                                    )
+                                )
+                            )
+                        }
+
+                        override fun onError(throwable: Throwable?) {
+                            Log.e(TAG, "Error loading theme bundles", throwable)
+                            continuation.resumeWith(
+                                Result.failure(
+                                    throwable ?: Throwable("Error loading theme bundles")
+                                )
+                            )
+                        }
+                    },
+                    /* reload= */ false,
+                )
             }
+        }
 
     override suspend fun select(colorOptionModel: ColorOptionModel) {
         _isApplyingSystemColor.value = true
diff --git a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl2.kt b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl2.kt
index 96f9368..c925748 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl2.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl2.kt
@@ -78,8 +78,7 @@
             .map { (_, colors) -> colors }
 
     override val colorOptions: Flow<Map<ColorType, List<ColorOption>>> =
-        combine(homeWallpaperColors, lockWallpaperColors, ::Pair)
-            .map { (homeColors, lockColors) ->
+        combine(homeWallpaperColors, lockWallpaperColors) { homeColors, lockColors ->
                 suspendCancellableCoroutine { continuation ->
                     colorManager.setWallpaperColors(homeColors, lockColors)
                     colorManager.fetchOptions(
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 84462ad..93a2d01 100644
--- a/src/com/android/customization/picker/grid/data/repository/ShapeGridRepository.kt
+++ b/src/com/android/customization/picker/grid/data/repository/ShapeGridRepository.kt
@@ -72,6 +72,6 @@
         }
 
     fun getGridOptionDrawable(iconId: Int): Drawable? {
-        return manager.getGridOptionDrawble(iconId)
+        return manager.getGridOptionDrawable(iconId)
     }
 }
diff --git a/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt
index b44b152..9ce7b61 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt
@@ -418,7 +418,7 @@
                             clockHostView.removeAllViews()
                             // For new customization picker, we should get views from clocklayout
                             if (Flags.newCustomizationPickerUi()) {
-                                clockViewFactory.getController(clock.clockId).let { clockController
+                                clockViewFactory.getController(clock.clockId)?.let { clockController
                                     ->
                                     val udfpsTop =
                                         clockPickerViewModel.getUdfpsLocation()?.let {
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModel.kt
index ce97154..fe1c6f8 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModel.kt
@@ -17,11 +17,14 @@
 package com.android.wallpaper.customization.ui.viewmodel
 
 import android.content.Context
+import android.content.res.Resources
 import android.graphics.drawable.Drawable
+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.ShapeIconViewModel
+import com.android.customization.widget.GridTileDrawable
 import com.android.themepicker.R
 import com.android.wallpaper.picker.common.icon.ui.viewmodel.Icon
 import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
@@ -186,6 +189,21 @@
     }
 
     private fun toGridOptionItemViewModel(option: GridOptionModel): OptionItemViewModel2<Drawable> {
+        // Fallback to use GridTileDrawable when no resource found for the icon ID
+        val drawable =
+            interactor.getGridOptionDrawable(option.iconId)
+                ?: GridTileDrawable(
+                    option.cols,
+                    option.rows,
+                    context.resources.getString(
+                        Resources.getSystem()
+                            .getIdentifier(
+                                ResourceConstants.CONFIG_ICON_MASK,
+                                "string",
+                                ResourceConstants.ANDROID_PACKAGE,
+                            )
+                    ),
+                )
         val isSelected =
             previewingGridKey
                 .map { it == option.key }
@@ -194,10 +212,9 @@
                     started = SharingStarted.Lazily,
                     initialValue = false,
                 )
-
         return OptionItemViewModel2(
             key = MutableStateFlow(option.key),
-            payload = interactor.getGridOptionDrawable(option.iconId),
+            payload = drawable,
             text = Text.Loaded(option.title),
             isSelected = isSelected,
             onClicked =
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 9fc59c7..1f2a0f3 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
@@ -146,14 +146,14 @@
                             )
                             .collect { (previewingClock, previewingClockSize) ->
                                 val hideSmartspace =
-                                    clockViewFactory.getController(previewingClock.clockId).let {
+                                    clockViewFactory.getController(previewingClock.clockId)?.let {
                                         when (previewingClockSize) {
                                             ClockSize.DYNAMIC ->
                                                 it.largeClock.config.hasCustomWeatherDataDisplay
                                             ClockSize.SMALL ->
                                                 it.smallClock.config.hasCustomWeatherDataDisplay
                                         }
-                                    }
+                                    } ?: false
                                 workspaceCallback.sendMessage(
                                     MESSAGE_ID_HIDE_SMART_SPACE,
                                     Bundle().apply {
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..b39988a 100644
--- a/tests/common/src/com/android/customization/model/grid/FakeShapeGridManager.kt
+++ b/tests/common/src/com/android/customization/model/grid/FakeShapeGridManager.kt
@@ -42,7 +42,7 @@
         return 0
     }
 
-    override fun getGridOptionDrawble(iconId: Int): Drawable? {
+    override fun getGridOptionDrawable(iconId: Int): Drawable? {
         return when (iconId) {
             0 -> gridOptionDrawable0
             1 -> gridOptionDrawable1
@@ -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/clock/ui/FakeClockViewFactory.kt b/tests/robotests/src/com/android/customization/picker/clock/ui/FakeClockViewFactory.kt
index 418b439..d6286a4 100644
--- a/tests/robotests/src/com/android/customization/picker/clock/ui/FakeClockViewFactory.kt
+++ b/tests/robotests/src/com/android/customization/picker/clock/ui/FakeClockViewFactory.kt
@@ -35,7 +35,7 @@
         override fun dump(pw: PrintWriter) = TODO("Not yet implemented")
     }
 
-    override fun getController(clockId: String): ClockController = clockControllers[clockId]!!
+    override fun getController(clockId: String): ClockController? = clockControllers[clockId]
 
     override fun getLargeView(clockId: String): View {
         TODO("Not yet implemented")
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()