Merge "Reset preview only after transition completed (2/2)" into main
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
index 6bc6180..f29be47 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
@@ -75,19 +75,17 @@
             return true
         }
 
-        val isBackPressedHandled = defaultCustomizationOptionsViewModel.handleBackPressed()
+        return defaultCustomizationOptionsViewModel.handleBackPressed()
+    }
 
-        if (isBackPressedHandled) {
-            // If isBackPressedHandled is handled by DefaultCustomizationOptionsViewModel, it means
-            // we navigate back to the main screen from a secondary screen. Reset preview.
-            keyguardQuickAffordancePickerViewModel2.resetPreview()
-            shapeGridPickerViewModel.resetPreview()
-            clockPickerViewModel.resetPreview()
-            colorPickerViewModel2.resetPreview()
-            darkModeViewModel.resetPreview()
-        }
+    override fun resetPreview() {
+        defaultCustomizationOptionsViewModel.resetPreview()
 
-        return isBackPressedHandled
+        keyguardQuickAffordancePickerViewModel2.resetPreview()
+        shapeGridPickerViewModel.resetPreview()
+        clockPickerViewModel.resetPreview()
+        colorPickerViewModel2.resetPreview()
+        darkModeViewModel.resetPreview()
     }
 
     val onCustomizeClockClicked: Flow<(() -> Unit)?> =