Fixes preview when setting lock to same live wp as home

Fixes: b/282048281
Test: manually verified that lock screen preview updates
Change-Id: I29cadbbf1e6a92e5db152847dac924423fdf06a1
diff --git a/src/com/android/wallpaper/picker/customization/ui/binder/ScreenPreviewBinder.kt b/src/com/android/wallpaper/picker/customization/ui/binder/ScreenPreviewBinder.kt
index f4bce19..e00b2bd 100644
--- a/src/com/android/wallpaper/picker/customization/ui/binder/ScreenPreviewBinder.kt
+++ b/src/com/android/wallpaper/picker/customization/ui/binder/ScreenPreviewBinder.kt
@@ -19,6 +19,7 @@
 
 import android.app.Activity
 import android.app.WallpaperColors
+import android.app.WallpaperManager
 import android.content.Intent
 import android.os.Bundle
 import android.service.wallpaper.WallpaperService
@@ -38,6 +39,7 @@
 import com.android.wallpaper.asset.CurrentWallpaperAssetVN
 import com.android.wallpaper.model.LiveWallpaperInfo
 import com.android.wallpaper.model.WallpaperInfo
+import com.android.wallpaper.module.CustomizationSections
 import com.android.wallpaper.picker.WorkspaceSurfaceHolderCallback
 import com.android.wallpaper.picker.customization.ui.viewmodel.ScreenPreviewViewModel
 import com.android.wallpaper.util.ResourceUtils
@@ -81,6 +83,7 @@
         val workspaceSurface: SurfaceView = previewView.requireViewById(R.id.workspace_surface)
         val wallpaperSurface: SurfaceView = previewView.requireViewById(R.id.wallpaper_surface)
         wallpaperSurface.setZOrderOnTop(false)
+        val wallpaperManager = WallpaperManager.getInstance(activity)
 
         if (dimWallpaper) {
             previewView.requireViewById<View>(R.id.wallpaper_dimming_scrim).isVisible = true
@@ -160,7 +163,7 @@
                 launch {
                     lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
                         var initialWallpaperUpdate = true
-                        viewModel.wallpaperUpdateEvents()?.collect {
+                        viewModel.wallpaperUpdateEvents()?.collect { wallpaperModel ->
                             // Do not update screen preview on initial update,since the initial
                             // update results from starting or resuming the activity.
                             //
@@ -172,6 +175,20 @@
                                 initialWallpaperUpdate = false
                             } else if (viewModel.shouldHandleReload()) {
                                 onPreviewDirty()
+                            } else if (
+                                viewModel.screen == CustomizationSections.Screen.LOCK_SCREEN &&
+                                    wallpaperManager.getWallpaperInfo(
+                                        WallpaperManager.FLAG_SYSTEM
+                                    ) != null &&
+                                    wallpaperModel?.wallpaperId ==
+                                        wallpaperManager
+                                            .getWallpaperInfo(WallpaperManager.FLAG_SYSTEM)
+                                            .serviceName
+                            ) {
+                                // Setting the lock screen to the same live wp as the home screen
+                                // doesn't trigger a UI update, so fix that here for now.
+                                // TODO(b/281730113) Remove this once better solution is ready.
+                                onPreviewDirty()
                             }
                         }
                     }
diff --git a/src/com/android/wallpaper/picker/customization/ui/viewmodel/ScreenPreviewViewModel.kt b/src/com/android/wallpaper/picker/customization/ui/viewmodel/ScreenPreviewViewModel.kt
index 9e8de48..93e5223 100644
--- a/src/com/android/wallpaper/picker/customization/ui/viewmodel/ScreenPreviewViewModel.kt
+++ b/src/com/android/wallpaper/picker/customization/ui/viewmodel/ScreenPreviewViewModel.kt
@@ -33,7 +33,7 @@
     private val wallpaperInfoProvider: suspend (forceReload: Boolean) -> WallpaperInfo?,
     private val onWallpaperColorChanged: (WallpaperColors?) -> Unit = {},
     private val wallpaperInteractor: WallpaperInteractor,
-    private val screen: CustomizationSections.Screen,
+    val screen: CustomizationSections.Screen,
 ) {
     /** Returns whether wallpaper picker should handle reload */
     fun shouldHandleReload(): Boolean {