Loading generated preview on-demand instead of keeping everything in memory

Bug: 369906121
Test: atest GeneratedPreviewTest
Flag: EXEMPT bugfix
Change-Id: Idd7610e8a5c577d2c7b0a1d7d2a1f1efde40b11f
diff --git a/src/com/android/launcher3/model/WidgetItem.java b/src/com/android/launcher3/model/WidgetItem.java
index e757a68..0caeb8d 100644
--- a/src/com/android/launcher3/model/WidgetItem.java
+++ b/src/com/android/launcher3/model/WidgetItem.java
@@ -1,19 +1,8 @@
 package com.android.launcher3.model;
 
-import static android.appwidget.AppWidgetProviderInfo.WIDGET_CATEGORY_HOME_SCREEN;
-import static android.appwidget.AppWidgetProviderInfo.WIDGET_CATEGORY_KEYGUARD;
-import static android.appwidget.AppWidgetProviderInfo.WIDGET_CATEGORY_SEARCHBOX;
-
-import android.annotation.SuppressLint;
 import android.content.Context;
 import android.content.pm.ActivityInfo;
-import android.content.res.Resources;
-import android.util.SparseArray;
-import android.widget.RemoteViews;
 
-import androidx.core.os.BuildCompat;
-
-import com.android.launcher3.Flags;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.icons.BitmapInfo;
@@ -21,7 +10,6 @@
 import com.android.launcher3.pm.ShortcutConfigActivityInfo;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
-import com.android.launcher3.widget.WidgetManagerHelper;
 
 /**
  * An wrapper over various items displayed in a widget picker,
@@ -37,11 +25,9 @@
     public final String label;
     public final CharSequence description;
     public final int spanX, spanY;
-    public final SparseArray<RemoteViews> generatedPreviews;
 
     public WidgetItem(LauncherAppWidgetProviderInfo info,
-            InvariantDeviceProfile idp, IconCache iconCache, Context context,
-            WidgetManagerHelper helper) {
+            InvariantDeviceProfile idp, IconCache iconCache, Context context) {
         super(info.provider, info.getProfile());
 
         label = iconCache.getTitleNoCache(info);
@@ -51,27 +37,6 @@
 
         spanX = Math.min(info.spanX, idp.numColumns);
         spanY = Math.min(info.spanY, idp.numRows);
-
-        if (BuildCompat.isAtLeastV() && Flags.enableGeneratedPreviews()) {
-            generatedPreviews = new SparseArray<>(3);
-            for (int widgetCategory : new int[] {
-                    WIDGET_CATEGORY_HOME_SCREEN,
-                    WIDGET_CATEGORY_KEYGUARD,
-                    WIDGET_CATEGORY_SEARCHBOX,
-            }) {
-                if ((widgetCategory & widgetInfo.generatedPreviewCategories) != 0) {
-                    generatedPreviews.put(widgetCategory,
-                            helper.loadGeneratedPreview(widgetInfo, widgetCategory));
-                }
-            }
-        } else {
-            generatedPreviews = null;
-        }
-    }
-
-    public WidgetItem(LauncherAppWidgetProviderInfo info,
-            InvariantDeviceProfile idp, IconCache iconCache, Context context) {
-        this(info, idp, iconCache, context, new WidgetManagerHelper(context));
     }
 
     public WidgetItem(ShortcutConfigActivityInfo info, IconCache iconCache) {
@@ -82,7 +47,6 @@
         widgetInfo = null;
         activityInfo = info;
         spanX = spanY = 1;
-        generatedPreviews = null;
     }
 
     /**
@@ -101,26 +65,8 @@
         return false;
     }
 
-    /** Returns whether this {@link WidgetItem} has a preview layout that can be used. */
-    @SuppressLint("NewApi") // Already added API check.
-    public boolean hasPreviewLayout() {
-        return widgetInfo != null && widgetInfo.previewLayout != Resources.ID_NULL;
-    }
-
     /** Returns whether this {@link WidgetItem} is for a shortcut rather than an app widget. */
     public boolean isShortcut() {
         return activityInfo != null;
     }
-
-    /**
-     * Returns whether this {@link WidgetItem} has a generated preview for the given widget
-     * category.
-     */
-    public boolean hasGeneratedPreview(int widgetCategory) {
-        if (!Flags.enableGeneratedPreviews() || generatedPreviews == null) {
-            return false;
-        }
-        return generatedPreviews.contains(widgetCategory)
-                && generatedPreviews.get(widgetCategory) != null;
-    }
 }
diff --git a/src/com/android/launcher3/model/WidgetsModel.java b/src/com/android/launcher3/model/WidgetsModel.java
index 01d4996..a27d2f1 100644
--- a/src/com/android/launcher3/model/WidgetsModel.java
+++ b/src/com/android/launcher3/model/WidgetsModel.java
@@ -149,8 +149,7 @@
                         LauncherAppWidgetProviderInfo.fromProviderInfo(context, widgetInfo);
 
                 widgetsAndShortcuts.add(new WidgetItem(
-                        launcherWidgetInfo, idp, app.getIconCache(), app.getContext(),
-                        widgetManager));
+                        launcherWidgetInfo, idp, app.getIconCache(), app.getContext()));
                 updatedItems.add(launcherWidgetInfo);
             }
 
@@ -213,7 +212,6 @@
         if (!WIDGETS_ENABLED) {
             return;
         }
-        WidgetManagerHelper widgetManager = new WidgetManagerHelper(app.getContext());
         for (Entry<PackageItemInfo, List<WidgetItem>> entry : mWidgetsByPackageItem.entrySet()) {
             if (packageNames.contains(entry.getKey().packageName)) {
                 List<WidgetItem> items = entry.getValue();
@@ -226,7 +224,7 @@
                         } else {
                             items.set(i, new WidgetItem(item.widgetInfo,
                                     app.getInvariantDeviceProfile(), app.getIconCache(),
-                                    app.getContext(), widgetManager));
+                                    app.getContext()));
                         }
                     }
                 }
diff --git a/src/com/android/launcher3/widget/DatabaseWidgetPreviewLoader.java b/src/com/android/launcher3/widget/DatabaseWidgetPreviewLoader.java
index e100157..d4538dd 100644
--- a/src/com/android/launcher3/widget/DatabaseWidgetPreviewLoader.java
+++ b/src/com/android/launcher3/widget/DatabaseWidgetPreviewLoader.java
@@ -15,12 +15,15 @@
  */
 package com.android.launcher3.widget;
 
-import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
+import static android.appwidget.AppWidgetProviderInfo.WIDGET_CATEGORY_HOME_SCREEN;
 
+import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
+import static com.android.launcher3.widget.LauncherAppWidgetProviderInfo.fromProviderInfo;
+
+import android.appwidget.AppWidgetProviderInfo;
 import android.content.Context;
 import android.content.res.Resources;
 import android.graphics.Bitmap;
-import android.graphics.Canvas;
 import android.graphics.Color;
 import android.graphics.Paint;
 import android.graphics.PorterDuff;
@@ -30,17 +33,19 @@
 import android.os.Handler;
 import android.util.Log;
 import android.util.Size;
+import android.widget.RemoteViews;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.VisibleForTesting;
+import androidx.core.os.BuildCompat;
 
 import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.Flags;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.icons.BitmapRenderer;
 import com.android.launcher3.icons.LauncherIcons;
-import com.android.launcher3.icons.ShadowGenerator;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.pm.ShortcutConfigActivityInfo;
 import com.android.launcher3.util.CancellableTask;
@@ -52,20 +57,19 @@
 import java.util.concurrent.ExecutionException;
 import java.util.function.Consumer;
 
-/** Utility class to load widget previews */
+/**
+ * Utility class to generate widget previews
+ *
+ * Note that it no longer uses database, all previews are freshly generated
+ */
 public class DatabaseWidgetPreviewLoader {
 
     private static final String TAG = "WidgetPreviewLoader";
 
     private final Context mContext;
-    private final float mPreviewBoxCornerRadius;
 
     public DatabaseWidgetPreviewLoader(Context context) {
         mContext = context;
-        float previewCornerRadius = RoundedCornerEnforcement.computeEnforcedRadius(context);
-        mPreviewBoxCornerRadius = previewCornerRadius > 0
-                ? previewCornerRadius
-                : mContext.getResources().getDimension(R.dimen.widget_preview_corner_radius);
     }
 
     /**
@@ -77,10 +81,10 @@
     public CancellableTask loadPreview(
             @NonNull WidgetItem item,
             @NonNull Size previewSize,
-            @NonNull Consumer<Bitmap> callback) {
+            @NonNull Consumer<WidgetPreviewInfo> callback) {
         Handler handler = getLoaderExecutor().getHandler();
-        CancellableTask<Bitmap> request = new CancellableTask<>(
-                () -> generatePreview(item, previewSize.getWidth(), previewSize.getHeight()),
+        CancellableTask<WidgetPreviewInfo> request = new CancellableTask<>(
+                () -> generatePreviewInfoBg(item, previewSize.getWidth(), previewSize.getHeight()),
                 MAIN_EXECUTOR,
                 callback);
         Utilities.postAsyncCallback(handler, request);
@@ -93,6 +97,39 @@
         return Executors.UI_HELPER_EXECUTOR;
     }
 
+    /** Generated the preview object. This method must be called on a background thread */
+    @VisibleForTesting
+    @NonNull
+    public WidgetPreviewInfo generatePreviewInfoBg(
+            WidgetItem item, int previewWidth, int previewHeight) {
+        WidgetPreviewInfo result = new WidgetPreviewInfo();
+
+        AppWidgetProviderInfo widgetInfo = item.widgetInfo;
+        if (BuildCompat.isAtLeastV() && Flags.enableGeneratedPreviews() && widgetInfo != null
+                && ((widgetInfo.generatedPreviewCategories & WIDGET_CATEGORY_HOME_SCREEN) != 0)) {
+            result.remoteViews = new WidgetManagerHelper(mContext)
+                    .loadGeneratedPreview(widgetInfo, WIDGET_CATEGORY_HOME_SCREEN);
+            if (result.remoteViews != null) {
+                result.providerInfo = widgetInfo;
+            }
+        }
+
+        if (result.providerInfo == null && widgetInfo != null
+                && widgetInfo.previewLayout != Resources.ID_NULL) {
+            result.providerInfo = fromProviderInfo(mContext, widgetInfo.clone());
+            // A hack to force the initial layout to be the preview layout since there is no API for
+            // rendering a preview layout for work profile apps yet. For non-work profile layout, a
+            // proper solution is to use RemoteViews(PackageName, LayoutId).
+            result.providerInfo.initialLayout = item.widgetInfo.previewLayout;
+        }
+
+        if (result.providerInfo == null) {
+            // fallback to bitmap preview
+            result.previewBitmap = generatePreview(item, previewWidth, previewHeight);
+        }
+        return result;
+    }
+
     /**
      * Returns a generated preview for a widget and if the preview should be saved in persistent
      * storage.
@@ -232,21 +269,6 @@
         });
     }
 
-    private RectF drawBoxWithShadow(Canvas c, int width, int height) {
-        Resources res = mContext.getResources();
-
-        ShadowGenerator.Builder builder = new ShadowGenerator.Builder(Color.WHITE);
-        builder.shadowBlur = res.getDimension(R.dimen.widget_preview_shadow_blur);
-        builder.radius = mPreviewBoxCornerRadius;
-        builder.keyShadowDistance = res.getDimension(R.dimen.widget_preview_key_shadow_distance);
-
-        builder.bounds.set(builder.shadowBlur, builder.shadowBlur,
-                width - builder.shadowBlur,
-                height - builder.shadowBlur - builder.keyShadowDistance);
-        builder.drawShadow(c);
-        return builder.bounds;
-    }
-
     private Bitmap generateShortcutPreview(
             ShortcutConfigActivityInfo info, int maxWidth, int maxHeight) {
         int iconSize = ActivityContext.lookupContext(mContext).getDeviceProfile().allAppsIconSizePx;
@@ -280,4 +302,15 @@
             throw new RuntimeException(e);
         }
     }
+
+    /**
+     * Simple class to hold preview information
+     */
+    public static class WidgetPreviewInfo {
+
+        public AppWidgetProviderInfo providerInfo;
+        public RemoteViews remoteViews;
+
+        public Bitmap previewBitmap;
+    }
 }
diff --git a/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java b/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
index 0cf1f2e..b07d807 100644
--- a/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
+++ b/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
@@ -123,7 +123,7 @@
     @Override
     public void setAppWidget(int appWidgetId, AppWidgetProviderInfo info) {
         super.setAppWidget(appWidgetId, info);
-        if (!mTrackingWidgetUpdate) {
+        if (!mTrackingWidgetUpdate && appWidgetId != -1) {
             mTrackingWidgetUpdate = true;
             Trace.beginAsyncSection(TRACE_METHOD_NAME + info.provider, appWidgetId);
             Log.i(TAG, "App widget created with id: " + appWidgetId);
@@ -255,16 +255,6 @@
     }
 
     @Override
-    public AppWidgetProviderInfo getAppWidgetInfo() {
-        AppWidgetProviderInfo info = super.getAppWidgetInfo();
-        if (info != null && !(info instanceof LauncherAppWidgetProviderInfo)) {
-            throw new IllegalStateException("Launcher widget must have"
-                    + " LauncherAppWidgetProviderInfo");
-        }
-        return info;
-    }
-
-    @Override
     public void getFocusedRect(Rect r) {
         super.getFocusedRect(r);
         // Outset to a larger rect for drawing a padding between focus outline and widget
diff --git a/src/com/android/launcher3/widget/WidgetCell.java b/src/com/android/launcher3/widget/WidgetCell.java
index 9e635a3..4811a17 100644
--- a/src/com/android/launcher3/widget/WidgetCell.java
+++ b/src/com/android/launcher3/widget/WidgetCell.java
@@ -16,18 +16,17 @@
 
 package com.android.launcher3.widget;
 
-import static android.appwidget.AppWidgetProviderInfo.WIDGET_CATEGORY_HOME_SCREEN;
 import static android.view.accessibility.AccessibilityNodeInfo.ACTION_CLICK;
 
 import static com.android.launcher3.Flags.enableWidgetTapToAdd;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_TRAY;
-import static com.android.launcher3.widget.LauncherAppWidgetProviderInfo.fromProviderInfo;
 import static com.android.launcher3.widget.util.WidgetSizes.getWidgetItemSizePx;
 
 import android.animation.Animator;
 import android.animation.AnimatorSet;
 import android.animation.ObjectAnimator;
 import android.animation.TimeInterpolator;
+import android.appwidget.AppWidgetProviderInfo;
 import android.content.Context;
 import android.graphics.Bitmap;
 import android.graphics.Rect;
@@ -48,13 +47,10 @@
 import android.widget.RemoteViews;
 import android.widget.TextView;
 
-import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
 import com.android.app.animation.Interpolators;
 import com.android.launcher3.CheckLongPressHelper;
-import com.android.launcher3.Flags;
-import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.R;
 import com.android.launcher3.anim.AnimatedPropertySetter;
@@ -65,11 +61,10 @@
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.util.CancellableTask;
 import com.android.launcher3.views.ActivityContext;
+import com.android.launcher3.widget.DatabaseWidgetPreviewLoader.WidgetPreviewInfo;
 import com.android.launcher3.widget.picker.util.WidgetPreviewContainerSize;
 import com.android.launcher3.widget.util.WidgetSizes;
 
-import java.util.function.Consumer;
-
 /**
  * Represents the individual cell of the widget inside the widget tray. The preview is drawn
  * horizontally centered, and scaled down if needed.
@@ -242,17 +237,6 @@
      * Applies the item to this view
      */
     public void applyFromCellItem(WidgetItem item) {
-        applyFromCellItem(item, this::applyPreview, /*cachedPreview=*/null);
-    }
-
-    /**
-     * Applies the item to this view
-     * @param item item to apply
-     * @param callback callback when preview is loaded in case the preview is being loaded or cached
-     * @param cachedPreview previously cached preview bitmap is present
-     */
-    public void applyFromCellItem(WidgetItem item, @NonNull Consumer<Bitmap> callback,
-            @Nullable Bitmap cachedPreview) {
         Context context = getContext();
         mItem = item;
         mWidgetSize = getWidgetItemSizePx(getContext(), mActivity.getDeviceProfile(), mItem);
@@ -283,37 +267,28 @@
         }
 
         if (mRemoteViewsPreview != null) {
-            mAppWidgetHostViewPreview = createAppWidgetHostView(context);
-            setAppWidgetHostViewPreview(mAppWidgetHostViewPreview, item.widgetInfo,
-                    mRemoteViewsPreview);
-        } else if (Flags.enableGeneratedPreviews()
-                && item.hasGeneratedPreview(WIDGET_CATEGORY_HOME_SCREEN)) {
-            mAppWidgetHostViewPreview = createAppWidgetHostView(context);
-            setAppWidgetHostViewPreview(mAppWidgetHostViewPreview, item.widgetInfo,
-                    item.generatedPreviews.get(WIDGET_CATEGORY_HOME_SCREEN));
-        } else if (item.hasPreviewLayout()) {
-            // If the context is a Launcher activity, DragView will show mAppWidgetHostViewPreview
-            // as a preview during drag & drop. And thus, we should use LauncherAppWidgetHostView,
-            // which supports applying local color extraction during drag & drop.
-            mAppWidgetHostViewPreview = isLauncherContext(context)
-                    ? new LauncherAppWidgetHostView(context)
-                    : createAppWidgetHostView(context);
-            LauncherAppWidgetProviderInfo providerInfo =
-                    fromProviderInfo(context, item.widgetInfo.clone());
-            // A hack to force the initial layout to be the preview layout since there is no API for
-            // rendering a preview layout for work profile apps yet. For non-work profile layout, a
-            // proper solution is to use RemoteViews(PackageName, LayoutId).
-            providerInfo.initialLayout = item.widgetInfo.previewLayout;
-            setAppWidgetHostViewPreview(mAppWidgetHostViewPreview, providerInfo, null);
-        } else if (cachedPreview != null) {
-            applyPreview(cachedPreview);
+            WidgetPreviewInfo previewInfo = new WidgetPreviewInfo();
+            previewInfo.providerInfo = item.widgetInfo;
+            previewInfo.remoteViews = mRemoteViewsPreview;
+            applyPreview(previewInfo);
         } else {
             if (mActiveRequest == null) {
-                mActiveRequest = mWidgetPreviewLoader.loadPreview(mItem, mWidgetSize, callback);
+                mActiveRequest = mWidgetPreviewLoader.loadPreview(
+                        mItem, mWidgetSize, this::applyPreview);
             }
         }
     }
 
+    private void applyPreview(WidgetPreviewInfo previewInfo) {
+        if (previewInfo.providerInfo != null) {
+            mAppWidgetHostViewPreview = createAppWidgetHostView(getContext());
+            setAppWidgetHostViewPreview(mAppWidgetHostViewPreview, previewInfo.providerInfo,
+                    previewInfo.remoteViews);
+        } else {
+            applyBitmapPreview(previewInfo.previewBitmap);
+        }
+    }
+
     private void initPreviewContainerSizeAndScale() {
         WidgetPreviewContainerSize previewSize = WidgetPreviewContainerSize.Companion.forItem(mItem,
                 mActivity.getDeviceProfile());
@@ -337,7 +312,7 @@
 
     private void setAppWidgetHostViewPreview(
             NavigableAppWidgetHostView appWidgetHostViewPreview,
-            LauncherAppWidgetProviderInfo providerInfo,
+            AppWidgetProviderInfo providerInfo,
             @Nullable RemoteViews remoteViews) {
         appWidgetHostViewPreview.setImportantForAccessibility(IMPORTANT_FOR_ACCESSIBILITY_NO);
         appWidgetHostViewPreview.setAppWidget(/* appWidgetId= */ -1, providerInfo);
@@ -349,7 +324,7 @@
                 mWidgetSize.getWidth(), mWidgetSize.getHeight(), Gravity.CENTER);
         mWidgetImageContainer.addView(appWidgetHostViewPreview, /* index= */ 0, widgetHostLP);
         mWidgetImage.setVisibility(View.GONE);
-        applyPreview(null);
+        applyBitmapPreview(null);
 
         appWidgetHostViewPreview.addOnLayoutChangeListener(
                 (v, l, t, r, b, ol, ot, or, ob) ->
@@ -407,7 +382,7 @@
         mAnimatePreview = shouldAnimate;
     }
 
-    private void applyPreview(Bitmap bitmap) {
+    private void applyBitmapPreview(Bitmap bitmap) {
         if (bitmap != null) {
             Drawable drawable = new RoundDrawableWrapper(
                     new FastBitmapDrawable(bitmap), mEnforcedCornerRadius);
@@ -496,8 +471,8 @@
         mLongPressHelper.cancelLongPress();
     }
 
-    private static NavigableAppWidgetHostView createAppWidgetHostView(Context context) {
-        return new NavigableAppWidgetHostView(context) {
+    private static LauncherAppWidgetHostView createAppWidgetHostView(Context context) {
+        return new LauncherAppWidgetHostView(context) {
             @Override
             protected boolean shouldAllowDirectClick() {
                 return false;
@@ -505,10 +480,6 @@
         };
     }
 
-    private static boolean isLauncherContext(Context context) {
-        return ActivityContext.lookupContext(context) instanceof Launcher;
-    }
-
     @Override
     public CharSequence getAccessibilityClassName() {
         return WidgetCell.class.getName();
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/GeneratedPreviewTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/GeneratedPreviewTest.kt
index ac5fda2..b92582c 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/GeneratedPreviewTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/GeneratedPreviewTest.kt
@@ -1,9 +1,8 @@
 package com.android.launcher3.widget
 
+import android.appwidget.AppWidgetManager
 import android.appwidget.AppWidgetProviderInfo
 import android.appwidget.AppWidgetProviderInfo.WIDGET_CATEGORY_HOME_SCREEN
-import android.appwidget.AppWidgetProviderInfo.WIDGET_CATEGORY_KEYGUARD
-import android.appwidget.AppWidgetProviderInfo.WIDGET_CATEGORY_SEARCHBOX
 import android.content.ComponentName
 import android.content.Context
 import android.content.pm.ActivityInfo
@@ -14,23 +13,30 @@
 import android.view.ContextThemeWrapper
 import android.view.LayoutInflater
 import android.widget.RemoteViews
-import androidx.test.core.app.ApplicationProvider.getApplicationContext
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.Flags.FLAG_ENABLE_GENERATED_PREVIEWS
 import com.android.launcher3.InvariantDeviceProfile
+import com.android.launcher3.R
 import com.android.launcher3.icons.IconCache
-import com.android.launcher3.icons.IconProvider
 import com.android.launcher3.model.WidgetItem
 import com.android.launcher3.util.ActivityContextWrapper
 import com.android.launcher3.util.Executors
+import com.android.launcher3.util.Executors.MAIN_EXECUTOR
+import com.android.launcher3.util.LauncherModelHelper.SandboxModelContext
+import com.android.launcher3.util.TestUtil
 import com.google.common.truth.Truth.assertThat
 import org.junit.After
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doAnswer
+import org.mockito.kotlin.whenever
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
@@ -46,28 +52,25 @@
         getInstrumentation().context.run {
             resources.getIdentifier("test_layout_appwidget_blue", "layout", packageName)
         }
-    private lateinit var context: Context
+
+    private lateinit var context: SandboxModelContext
+    private lateinit var uiContext: Context
     private lateinit var generatedPreview: RemoteViews
     private lateinit var widgetCell: WidgetCell
-    private lateinit var helper: WidgetManagerHelper
     private lateinit var appWidgetProviderInfo: LauncherAppWidgetProviderInfo
     private lateinit var widgetItem: WidgetItem
-    private lateinit var iconCache: IconCache
+
+    @Mock lateinit var iconCache: IconCache
 
     @Before
     fun setup() {
-        context = getApplicationContext()
+        MockitoAnnotations.initMocks(this)
+        context = SandboxModelContext()
         generatedPreview = RemoteViews(context.packageName, generatedPreviewLayout)
+        uiContext =
+            ActivityContextWrapper(ContextThemeWrapper(context, R.style.WidgetContainerTheme))
         widgetCell =
-            LayoutInflater.from(
-                    ActivityContextWrapper(
-                        ContextThemeWrapper(
-                            context,
-                            com.android.launcher3.R.style.WidgetContainerTheme,
-                        )
-                    )
-                )
-                .inflate(com.android.launcher3.R.layout.widget_cell, null) as WidgetCell
+            LayoutInflater.from(uiContext).inflate(R.layout.widget_cell, null) as WidgetCell
         appWidgetProviderInfo =
             AppWidgetProviderInfo()
                 .apply {
@@ -76,72 +79,52 @@
                     providerInfo = ActivityInfo().apply { applicationInfo = ApplicationInfo() }
                 }
                 .let { LauncherAppWidgetProviderInfo.fromProviderInfo(context, it) }
-        helper =
-            object : WidgetManagerHelper(context) {
-                override fun loadGeneratedPreview(
-                    info: AppWidgetProviderInfo,
-                    widgetCategory: Int,
-                ) =
-                    generatedPreview.takeIf {
-                        info === appWidgetProviderInfo &&
-                            widgetCategory == WIDGET_CATEGORY_HOME_SCREEN
-                    }
+
+        val widgetManager = context.spyService(AppWidgetManager::class.java)
+        doAnswer { i ->
+                generatedPreview.takeIf {
+                    i.arguments[0] == appWidgetProviderInfo.provider &&
+                        i.arguments[1] == appWidgetProviderInfo.user &&
+                        i.arguments[2] == WIDGET_CATEGORY_HOME_SCREEN
+                }
             }
+            .whenever(widgetManager)
+            .getWidgetPreview(any(), any(), any())
         createWidgetItem()
     }
 
     @After
     fun tearDown() {
-        iconCache.close()
+        context.destroy()
     }
 
     private fun createWidgetItem() {
         Executors.MODEL_EXECUTOR.submit {
                 val idp = InvariantDeviceProfile()
-                if (::iconCache.isInitialized) iconCache.close()
-                iconCache = IconCache(context, idp, null, IconProvider(context))
-                widgetItem = WidgetItem(appWidgetProviderInfo, idp, iconCache, context, helper)
+                widgetItem = WidgetItem(appWidgetProviderInfo, idp, iconCache, context)
             }
             .get()
     }
 
     @Test
-    fun widgetItem_hasGeneratedPreview() {
-        assertThat(widgetItem.hasGeneratedPreview(WIDGET_CATEGORY_HOME_SCREEN)).isTrue()
-        assertThat(widgetItem.hasGeneratedPreview(WIDGET_CATEGORY_KEYGUARD)).isFalse()
-        assertThat(widgetItem.hasGeneratedPreview(WIDGET_CATEGORY_SEARCHBOX)).isFalse()
-    }
-
-    @Test
     fun widgetItem_hasGeneratedPreview_noPreview() {
         appWidgetProviderInfo.generatedPreviewCategories = 0
         createWidgetItem()
-        assertThat(widgetItem.hasGeneratedPreview(WIDGET_CATEGORY_HOME_SCREEN)).isFalse()
-        assertThat(widgetItem.hasGeneratedPreview(WIDGET_CATEGORY_KEYGUARD)).isFalse()
-        assertThat(widgetItem.hasGeneratedPreview(WIDGET_CATEGORY_SEARCHBOX)).isFalse()
-    }
-
-    @Test
-    fun widgetItem_hasGeneratedPreview_nullPreview() {
-        appWidgetProviderInfo.generatedPreviewCategories =
-            WIDGET_CATEGORY_HOME_SCREEN or WIDGET_CATEGORY_KEYGUARD
-        createWidgetItem()
-        assertThat(widgetItem.hasGeneratedPreview(WIDGET_CATEGORY_HOME_SCREEN)).isTrue()
-        // loadGeneratedPreview returns null for KEYGUARD, so this should still be false.
-        assertThat(widgetItem.hasGeneratedPreview(WIDGET_CATEGORY_KEYGUARD)).isFalse()
-        assertThat(widgetItem.hasGeneratedPreview(WIDGET_CATEGORY_SEARCHBOX)).isFalse()
+        val preview = DatabaseWidgetPreviewLoader(uiContext).generatePreviewInfoBg(widgetItem, 1, 1)
+        assertThat(preview.remoteViews).isNull()
     }
 
     @Test
     fun widgetItem_getGeneratedPreview() {
-        val preview = widgetItem.generatedPreviews.get(WIDGET_CATEGORY_HOME_SCREEN)
-        assertThat(preview).isEqualTo(generatedPreview)
+        val preview = DatabaseWidgetPreviewLoader(uiContext).generatePreviewInfoBg(widgetItem, 1, 1)
+        assertThat(preview.remoteViews).isEqualTo(generatedPreview)
     }
 
     @Test
     fun widgetCell_showGeneratedPreview() {
         widgetCell.applyFromCellItem(widgetItem)
-        DatabaseWidgetPreviewLoader.getLoaderExecutor().submit {}.get()
+        TestUtil.runOnExecutorSync(DatabaseWidgetPreviewLoader.getLoaderExecutor()) {}
+        TestUtil.runOnExecutorSync(MAIN_EXECUTOR) {}
         assertThat(widgetCell.appWidgetHostViewPreview).isNotNull()
         assertThat(widgetCell.appWidgetHostViewPreview?.appWidgetInfo)
             .isEqualTo(appWidgetProviderInfo)
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java b/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
index 0d9464a..86bbcc1 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
@@ -54,7 +54,6 @@
 import com.android.launcher3.widget.DatabaseWidgetPreviewLoader;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.WidgetCell;
-import com.android.launcher3.widget.WidgetManagerHelper;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
 
 import org.junit.Before;
@@ -143,7 +142,6 @@
     }
 
     private List<WidgetItem> generateWidgetItems(String packageName, int numOfWidgets) {
-        WidgetManagerHelper widgetManager = new WidgetManagerHelper(mContext);
         ArrayList<WidgetItem> widgetItems = new ArrayList<>();
         for (int i = 0; i < numOfWidgets; i++) {
             ComponentName cn = ComponentName.createRelative(packageName, ".SampleWidget" + i);
@@ -151,7 +149,7 @@
 
             widgetItems.add(new WidgetItem(
                     LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, widgetInfo),
-                    mTestProfile, mIconCache, mContext, widgetManager));
+                    mTestProfile, mIconCache, mContext));
         }
         return widgetItems;
     }