1 files changed, 21 insertions, 11 deletions
diff --git a/src/com/android/wallpaper/module/DefaultCurrentWallpaperInfoFactory.java b/src/com/android/wallpaper/module/DefaultCurrentWallpaperInfoFactory.java
index 9e2787fd..c9f9b267 100755
--- a/src/com/android/wallpaper/module/DefaultCurrentWallpaperInfoFactory.java
+++ b/src/com/android/wallpaper/module/DefaultCurrentWallpaperInfoFactory.java
@@ -15,12 +15,14 @@
  */
 package com.android.wallpaper.module;
 
+import android.app.WallpaperManager;
 import android.content.Context;
 
 import androidx.annotation.Nullable;
 
 import com.android.wallpaper.compat.WallpaperManagerCompat;
 import com.android.wallpaper.model.CurrentWallpaperInfoVN;
+import com.android.wallpaper.model.LiveWallpaperMetadata;
 import com.android.wallpaper.model.WallpaperInfo;
 import com.android.wallpaper.module.WallpaperPreferences.PresentationMode;
 
@@ -33,6 +35,7 @@ public class DefaultCurrentWallpaperInfoFactory implements CurrentWallpaperInfoF
     private final Context mAppContext;
     private final WallpaperRefresher mWallpaperRefresher;
     private final LiveWallpaperInfoFactory mLiveWallpaperInfoFactory;
+    private final WallpaperManager mWallpaperManager;
 
     // Cached copies of the currently-set WallpaperInfo(s) and presentation mode.
     private WallpaperInfo mHomeWallpaper;
@@ -46,6 +49,7 @@ public class DefaultCurrentWallpaperInfoFactory implements CurrentWallpaperInfoF
         Injector injector = InjectorProvider.getInjector();
         mWallpaperRefresher = injector.getWallpaperRefresher(mAppContext);
         mLiveWallpaperInfoFactory = injector.getLiveWallpaperInfoFactory(mAppContext);
+        mWallpaperManager = WallpaperManager.getInstance(context);
     }
 
     @Override
@@ -68,7 +72,10 @@ public class DefaultCurrentWallpaperInfoFactory implements CurrentWallpaperInfoF
         mWallpaperRefresher.refresh(
                 (homeWallpaperMetadata, lockWallpaperMetadata, presentationMode) -> {
                     WallpaperInfo homeWallpaper;
-                    if (homeWallpaperMetadata.getWallpaperComponent() == null) {
+                    if (homeWallpaperMetadata instanceof LiveWallpaperMetadata) {
+                        homeWallpaper = mLiveWallpaperInfoFactory.getLiveWallpaperInfo(
+                                homeWallpaperMetadata.getWallpaperComponent());
+                    } else {
                         homeWallpaper = new CurrentWallpaperInfoVN(
                                 homeWallpaperMetadata.getAttributions(),
                                 homeWallpaperMetadata.getActionUrl(),
@@ -76,21 +83,24 @@ public class DefaultCurrentWallpaperInfoFactory implements CurrentWallpaperInfoF
                                 homeWallpaperMetadata.getActionIconRes(),
                                 homeWallpaperMetadata.getCollectionId(),
                                 WallpaperManagerCompat.FLAG_SYSTEM);
-                    } else {
-                        homeWallpaper = mLiveWallpaperInfoFactory.getLiveWallpaperInfo(
-                                homeWallpaperMetadata.getWallpaperComponent());
                     }
 
                     WallpaperInfo lockWallpaper = null;
 
                     if (lockWallpaperMetadata != null) {
-                        lockWallpaper = new CurrentWallpaperInfoVN(
-                                lockWallpaperMetadata.getAttributions(),
-                                lockWallpaperMetadata.getActionUrl(),
-                                lockWallpaperMetadata.getActionLabelRes(),
-                                lockWallpaperMetadata.getActionIconRes(),
-                                lockWallpaperMetadata.getCollectionId(),
-                                WallpaperManagerCompat.FLAG_LOCK);
+
+                        if (lockWallpaperMetadata instanceof LiveWallpaperMetadata) {
+                            lockWallpaper = mLiveWallpaperInfoFactory.getLiveWallpaperInfo(
+                                    lockWallpaperMetadata.getWallpaperComponent());
+                        } else {
+                            lockWallpaper = new CurrentWallpaperInfoVN(
+                                    lockWallpaperMetadata.getAttributions(),
+                                    lockWallpaperMetadata.getActionUrl(),
+                                    lockWallpaperMetadata.getActionLabelRes(),
+                                    lockWallpaperMetadata.getActionIconRes(),
+                                    lockWallpaperMetadata.getCollectionId(),
+                                    WallpaperManagerCompat.FLAG_LOCK);
+                        }
                     }
 
                     mHomeWallpaper = homeWallpaper;