From bb7553017c8f272040d99bbe09a312c1178b2bcb Mon Sep 17 00:00:00 2001
From: Jon Spivack <spivack@google.com>
Date: Tue, 10 Aug 2021 18:19:24 -0700
Subject: [PATCH 01/10] Add callback to ImageActionUtils for unresolved intents

Adding an optional callback to ImageActionUtils.persistBitmapAndStartActivity allows for external handling of the case in which the given intent cannot be resolved. The main use case is for NIU Actions when the user has chosen an assistant that does not support the feature. TaskOverlayFactoryGo provides a callback that shows a dialog in this case.

Various parts of TaskOverlayFactoryGo were also made public for better testability.

Bug: 196125662
Bug: 192406446
Bug: 195681795
Test: m -j RunLauncherGoGoogleRoboTests ROBOTEST_FILTER=TaskOverlayFactoryGoTest
Change-Id: I64f3a1274bc942a64e964dca20bd4245e336ad9d
---
 .../quickstep/TaskOverlayFactoryGo.java       | 29 ++++++++++++-------
 .../android/quickstep/ImageActionsApi.java    | 13 +++++----
 .../quickstep/util/ImageActionUtils.java      | 15 ++++++++++
 3 files changed, 42 insertions(+), 15 deletions(-)

diff --git a/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java b/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
index 6502526dcc..bc38739d57 100644
--- a/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
+++ b/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
@@ -50,7 +50,6 @@ import androidx.annotation.IntDef;
 import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.BaseActivity;
-import com.android.launcher3.BaseDraggingActivity;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.views.ArrowTipView;
@@ -76,7 +75,7 @@ public final class TaskOverlayFactoryGo extends TaskOverlayFactory {
     public static final String ACTIONS_ERROR_CODE = "niu_actions_app_error_code";
     public static final int ERROR_PERMISSIONS_STRUCTURE = 1;
     public static final int ERROR_PERMISSIONS_SCREENSHOT = 2;
-    private static final String NIU_ACTIONS_CONFIRMED = "launcher_go.niu_actions_confirmed";
+    public static final String NIU_ACTIONS_CONFIRMED = "launcher_go.niu_actions_confirmed";
     private static final String ASSIST_SETTINGS_ARGS_BUNDLE = ":settings:show_fragment_args";
     private static final String ASSIST_SETTINGS_ARGS_KEY = ":settings:fragment_args_key";
     private static final String ASSIST_SETTINGS_PREFERENCE_KEY = "default_assist";
@@ -87,10 +86,11 @@ public final class TaskOverlayFactoryGo extends TaskOverlayFactory {
 
     @Retention(SOURCE)
     @IntDef({PRIVACY_CONFIRMATION, ASSISTANT_NOT_SELECTED, ASSISTANT_NOT_SUPPORTED})
-    private @interface DialogType{}
-    private static final int PRIVACY_CONFIRMATION = 0;
-    private static final int ASSISTANT_NOT_SELECTED = 1;
-    private static final int ASSISTANT_NOT_SUPPORTED = 2;
+    @VisibleForTesting
+    public @interface DialogType{}
+    public static final int PRIVACY_CONFIRMATION = 0;
+    public static final int ASSISTANT_NOT_SELECTED = 1;
+    public static final int ASSISTANT_NOT_SUPPORTED = 2;
 
     private AssistContentRequester mContentRequester;
 
@@ -211,7 +211,8 @@ public final class TaskOverlayFactoryGo extends TaskOverlayFactory {
             Intent intent = createNIUIntent(actionType);
             // Only add and send the image if the appropriate permissions are held
             if (mAssistStructurePermitted && mAssistScreenshotPermitted) {
-                mImageApi.shareAsDataWithExplicitIntent(/* crop */ null, intent);
+                mImageApi.shareAsDataWithExplicitIntent(/* crop */ null, intent,
+                        () -> showDialog(actionType, ASSISTANT_NOT_SUPPORTED));
             } else {
                 // If both permissions are disabled, the structure error code takes priority
                 // The user must enable that one before they can enable screenshots
@@ -301,7 +302,6 @@ public final class TaskOverlayFactoryGo extends TaskOverlayFactory {
             mImageApi = imageActionsApi;
         }
 
-        // TODO (b/192406446): Test that these dialogs are shown at the appropriate times
         private void showDialog(String action, @DialogType int type) {
             switch (type) {
                 case PRIVACY_CONFIRMATION:
@@ -334,7 +334,7 @@ public final class TaskOverlayFactoryGo extends TaskOverlayFactory {
                                 int bodyTextID, int button1TextID,
                                 View.OnClickListener button1Callback, int button2TextID,
                                 View.OnClickListener button2Callback) {
-            BaseDraggingActivity activity = BaseActivity.fromContext(getActionsView().getContext());
+            BaseActivity activity = BaseActivity.fromContext(getActionsView().getContext());
             LayoutInflater inflater = LayoutInflater.from(activity);
             View view = inflater.inflate(R.layout.niu_actions_dialog, /* root */ null);
 
@@ -368,6 +368,11 @@ public final class TaskOverlayFactoryGo extends TaskOverlayFactory {
             mDialog.cancel();
         }
 
+        @VisibleForTesting
+        public OverlayDialogGo getDialog() {
+            return mDialog;
+        }
+
         private void onDialogClickSettings(View v) {
             mDialog.dismiss();
 
@@ -401,7 +406,11 @@ public final class TaskOverlayFactoryGo extends TaskOverlayFactory {
         }
     }
 
-    private static final class OverlayDialogGo extends AlertDialog {
+    /**
+     * Basic modal dialog for various user prompts
+     */
+    @VisibleForTesting
+    public static final class OverlayDialogGo extends AlertDialog {
         private final String mAction;
         private final @DialogType int mType;
 
diff --git a/quickstep/src/com/android/quickstep/ImageActionsApi.java b/quickstep/src/com/android/quickstep/ImageActionsApi.java
index 8cb64c24b8..154848d224 100644
--- a/quickstep/src/com/android/quickstep/ImageActionsApi.java
+++ b/quickstep/src/com/android/quickstep/ImageActionsApi.java
@@ -64,20 +64,23 @@ public class ImageActionsApi {
      */
     @UiThread
     public void shareWithExplicitIntent(@Nullable Rect crop, Intent intent) {
-        addImageAndSendIntent(crop, intent, false);
+        addImageAndSendIntent(crop, intent, false, null /* exceptionCallback */);
     }
 
     /**
      * Share the image this api was constructed with using the provided intent. The implementation
      * should set the intent's data field to the URI pointing to the image.
+     * @param exceptionCallback An optional callback to be called when the intent can't be resolved
      */
     @UiThread
-    public void shareAsDataWithExplicitIntent(@Nullable Rect crop, Intent intent) {
-        addImageAndSendIntent(crop, intent, true);
+    public void shareAsDataWithExplicitIntent(@Nullable Rect crop, Intent intent,
+            @Nullable Runnable exceptionCallback) {
+        addImageAndSendIntent(crop, intent, true, exceptionCallback);
     }
 
     @UiThread
-    private void addImageAndSendIntent(@Nullable Rect crop, Intent intent, boolean setData) {
+    private void addImageAndSendIntent(@Nullable Rect crop, Intent intent, boolean setData,
+            @Nullable Runnable exceptionCallback) {
         if (mBitmapSupplier.get() == null) {
             Log.e(TAG, "No snapshot available, not starting share.");
             return;
@@ -92,7 +95,7 @@ public class ImageActionsApi {
                         intentForUri.putExtra(EXTRA_STREAM, uri);
                     }
                     return new Intent[]{intentForUri};
-                }, TAG));
+                }, TAG, exceptionCallback));
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/util/ImageActionUtils.java b/quickstep/src/com/android/quickstep/util/ImageActionUtils.java
index de7dbd64f5..51a9915523 100644
--- a/quickstep/src/com/android/quickstep/util/ImageActionUtils.java
+++ b/quickstep/src/com/android/quickstep/util/ImageActionUtils.java
@@ -154,6 +154,18 @@ public class ImageActionUtils {
     @WorkerThread
     public static void persistBitmapAndStartActivity(Context context, Bitmap bitmap, Rect crop,
             Intent intent, BiFunction<Uri, Intent, Intent[]> uriToIntentMap, String tag) {
+        persistBitmapAndStartActivity(context, bitmap, crop, intent, uriToIntentMap, tag,
+                (Runnable) null);
+    }
+
+    /**
+     * Starts activity based on given intent created from image uri.
+     * @param exceptionCallback An optional callback to be called when the intent can't be resolved
+     */
+    @WorkerThread
+    public static void persistBitmapAndStartActivity(Context context, Bitmap bitmap, Rect crop,
+            Intent intent, BiFunction<Uri, Intent, Intent[]> uriToIntentMap, String tag,
+            Runnable exceptionCallback) {
         Intent[] intents = uriToIntentMap.apply(getImageUri(bitmap, crop, context, tag), intent);
 
         try {
@@ -165,6 +177,9 @@ public class ImageActionUtils {
             }
         } catch (ActivityNotFoundException e) {
             Log.e(TAG, "No activity found to receive image intent");
+            if (exceptionCallback != null) {
+                exceptionCallback.run();
+            }
         }
     }
 

From 272d622bc4c572fbfc8e2b9362f810cf411def9d Mon Sep 17 00:00:00 2001
From: vadimt <vadimt@google.com>
Date: Tue, 17 Aug 2021 18:47:12 -0700
Subject: [PATCH 02/10] Improving diags for containers disappearing while
 getting children

Bug: 184609576
Test: presubmit
Change-Id: I7f527a9991a58fd148b17783078b5ac80979a5b3
---
 .../android/launcher3/tapl/LauncherInstrumentation.java  | 9 +++++++++
 tests/tapl/com/android/launcher3/tapl/Widgets.java       | 4 ++--
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 2ca40d8229..efd8a551d4 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -1016,6 +1016,15 @@ public final class LauncherInstrumentation {
         }
     }
 
+    List<UiObject2> getChildren(UiObject2 container) {
+        try {
+            return container.getChildren();
+        } catch (StaleObjectException e) {
+            fail("The container disappeared from screen");
+            return null;
+        }
+    }
+
     private boolean hasLauncherObject(String resId) {
         return mDevice.hasObject(getLauncherObjectSelector(resId));
     }
diff --git a/tests/tapl/com/android/launcher3/tapl/Widgets.java b/tests/tapl/com/android/launcher3/tapl/Widgets.java
index 99d988926a..6e7264ac0b 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widgets.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widgets.java
@@ -116,9 +116,9 @@ public final class Widgets extends LauncherInstrumentation.VisibleContainer {
                     "widget_preview");
             int i = 0;
             for (; ; ) {
-                final Collection<UiObject2> tableRows = widgetsContainer.getChildren();
+                final Collection<UiObject2> tableRows = mLauncher.getChildren(widgetsContainer);
                 for (UiObject2 row : tableRows) {
-                    final Collection<UiObject2> widgetCells = row.getChildren();
+                    final Collection<UiObject2> widgetCells = mLauncher.getChildren(row);
                     for (UiObject2 widget : widgetCells) {
                         final UiObject2 label = mLauncher.findObjectInContainer(widget,
                                 labelSelector);

From 39aa2e0f1da9012da0e0080d8448664eaecbf2fd Mon Sep 17 00:00:00 2001
From: Sunny Goyal <sunnygoyal@google.com>
Date: Thu, 19 Aug 2021 12:02:57 -0700
Subject: [PATCH 03/10] Moving OrientationTouchTransformerTest to
 instrumentation tests

Bug: 196825541
Test: Presubmit
Change-Id: Ifb0b00f789214a8dde246ab13703211d536364af
---
 .../quickstep/OrientationTouchTransformer.java       |  4 ++--
 .../quickstep/OrientationTouchTransformerTest.java   | 12 ++++++++----
 2 files changed, 10 insertions(+), 6 deletions(-)
 rename quickstep/{robolectric_tests => tests}/src/com/android/quickstep/OrientationTouchTransformerTest.java (98%)

diff --git a/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java b/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java
index 35a851ab65..81e69176c2 100644
--- a/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java
+++ b/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java
@@ -462,7 +462,7 @@ class OrientationTouchTransformer {
                             + "mRotation: " + mRotation
                             + " this: " + this);
                 }
-                event.transform(mTmpMatrix);
+                event.applyTransform(mTmpMatrix);
                 return true;
             }
             mTmpPoint[0] = event.getX();
@@ -478,7 +478,7 @@ class OrientationTouchTransformer {
             }
 
             if (contains(mTmpPoint[0], mTmpPoint[1])) {
-                event.transform(mTmpMatrix);
+                event.applyTransform(mTmpMatrix);
                 return true;
             }
             return false;
diff --git a/quickstep/robolectric_tests/src/com/android/quickstep/OrientationTouchTransformerTest.java b/quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
similarity index 98%
rename from quickstep/robolectric_tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
rename to quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
index 2b1b57c450..159a51f733 100644
--- a/quickstep/robolectric_tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
+++ b/quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
@@ -19,6 +19,8 @@ package com.android.quickstep;
 
 import static android.view.Display.DEFAULT_DISPLAY;
 
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
 import static com.android.quickstep.SysUINavigationMode.Mode.NO_BUTTON;
 
 import static org.junit.Assert.assertFalse;
@@ -40,6 +42,9 @@ import android.view.Display;
 import android.view.MotionEvent;
 import android.view.Surface;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
 import com.android.launcher3.ResourceUtils;
 import com.android.launcher3.util.DisplayController;
 
@@ -47,10 +52,9 @@ import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.MockitoAnnotations;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
 
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class OrientationTouchTransformerTest {
     static class ScreenSize {
         int mHeight;
@@ -293,7 +297,7 @@ public class OrientationTouchTransformerTest {
     }
 
     private DisplayController.Info createDisplayInfo(ScreenSize screenSize, int rotation) {
-        Context context = RuntimeEnvironment.application;
+        Context context = getApplicationContext();
         Display display = spy(context.getSystemService(DisplayManager.class)
                 .getDisplay(DEFAULT_DISPLAY));
 

From bea7e2ee6e7a7c457f83260015d439c0e1631b4f Mon Sep 17 00:00:00 2001
From: vadimt <vadimt@google.com>
Date: Mon, 23 Aug 2021 13:57:51 -0700
Subject: [PATCH 04/10] More logging for for Nexus home activity appearing
 after to-home gesture in L3

Bug: 192018189
Test: presubmit
Change-Id: I92cfa9e40f71773e636c700ab7f84df1d9feac19
---
 .../android/quickstep/AbsSwipeUpHandler.java    | 17 +++++++++++++++++
 .../OtherActivityInputConsumer.java             |  7 +++++++
 2 files changed, 24 insertions(+)

diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index fe16b33069..46d58577cd 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -65,6 +65,7 @@ import android.graphics.RectF;
 import android.os.Build;
 import android.os.IBinder;
 import android.os.SystemClock;
+import android.util.Log;
 import android.view.MotionEvent;
 import android.view.View;
 import android.view.View.OnApplyWindowInsetsListener;
@@ -87,6 +88,7 @@ import com.android.launcher3.logging.StatsLogManager;
 import com.android.launcher3.logging.StatsLogManager.StatsLogger;
 import com.android.launcher3.statemanager.BaseState;
 import com.android.launcher3.statemanager.StatefulActivity;
+import com.android.launcher3.testing.TestProtocol;
 import com.android.launcher3.tracing.InputConsumerProto;
 import com.android.launcher3.tracing.SwipeHandlerProto;
 import com.android.launcher3.util.TraceHelper;
@@ -876,6 +878,9 @@ public abstract class AbsSwipeUpHandler<T extends StatefulActivity<S>,
      */
     @UiThread
     public void onGestureEnded(float endVelocity, PointF velocity, PointF downPos) {
+        if (Utilities.IS_RUNNING_IN_TEST_HARNESS) {
+            Log.d(TestProtocol.L3_SWIPE_TO_HOME, "3");
+        }
         float flingThreshold = mContext.getResources()
                 .getDimension(R.dimen.quickstep_fling_threshold_speed);
         boolean isFling = mGestureStarted && !mIsMotionPaused
@@ -1037,6 +1042,9 @@ public abstract class AbsSwipeUpHandler<T extends StatefulActivity<S>,
     @UiThread
     private void handleNormalGestureEnd(float endVelocity, boolean isFling, PointF velocity,
             boolean isCancel) {
+        if (Utilities.IS_RUNNING_IN_TEST_HARNESS) {
+            Log.d(TestProtocol.L3_SWIPE_TO_HOME, "4");
+        }
         long duration = MAX_SWIPE_DURATION;
         float currentShift = mCurrentShift.value;
         final GestureEndTarget endTarget = calculateEndTarget(velocity, endVelocity,
@@ -1155,6 +1163,9 @@ public abstract class AbsSwipeUpHandler<T extends StatefulActivity<S>,
     @UiThread
     private void animateToProgress(float start, float end, long duration, Interpolator interpolator,
             GestureEndTarget target, PointF velocityPxPerMs) {
+        if (Utilities.IS_RUNNING_IN_TEST_HARNESS) {
+            Log.d(TestProtocol.L3_SWIPE_TO_HOME, "5");
+        }
         runOnRecentsAnimationStart(() -> animateToProgressInternal(start, end, duration,
                 interpolator, target, velocityPxPerMs));
     }
@@ -1183,6 +1194,9 @@ public abstract class AbsSwipeUpHandler<T extends StatefulActivity<S>,
     @UiThread
     private void animateToProgressInternal(float start, float end, long duration,
             Interpolator interpolator, GestureEndTarget target, PointF velocityPxPerMs) {
+        if (Utilities.IS_RUNNING_IN_TEST_HARNESS) {
+            Log.d(TestProtocol.L3_SWIPE_TO_HOME, "7");
+        }
         maybeUpdateRecentsAttachedState();
 
         // If we are transitioning to launcher, then listen for the activity to be restarted while
@@ -1812,6 +1826,9 @@ public abstract class AbsSwipeUpHandler<T extends StatefulActivity<S>,
      * be run when it is next started.
      */
     protected void runOnRecentsAnimationStart(Runnable action) {
+        if (Utilities.IS_RUNNING_IN_TEST_HARNESS) {
+            Log.d(TestProtocol.L3_SWIPE_TO_HOME, "6");
+        }
         if (mRecentsAnimationTargets == null) {
             mRecentsAnimationStartCallbacks.add(action);
         } else {
diff --git a/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java b/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
index 725c7c45a5..e6285f2902 100644
--- a/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
+++ b/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
@@ -50,6 +50,7 @@ import android.view.ViewConfiguration;
 import androidx.annotation.UiThread;
 
 import com.android.launcher3.R;
+import com.android.launcher3.Utilities;
 import com.android.launcher3.testing.TestLogging;
 import com.android.launcher3.testing.TestProtocol;
 import com.android.launcher3.tracing.InputConsumerProto;
@@ -356,6 +357,9 @@ public class OtherActivityInputConsumer extends ContextWrapper implements InputC
             }
             case ACTION_CANCEL:
             case ACTION_UP: {
+                if (Utilities.IS_RUNNING_IN_TEST_HARNESS) {
+                    Log.d(TestProtocol.L3_SWIPE_TO_HOME, "1");
+                }
                 if (DEBUG_FAILED_QUICKSWITCH && !mPassedWindowMoveSlop) {
                     float displacementX = mLastPos.x - mDownPos.x;
                     float displacementY = mLastPos.y - mDownPos.y;
@@ -409,6 +413,9 @@ public class OtherActivityInputConsumer extends ContextWrapper implements InputC
      * the animation can still be running.
      */
     private void finishTouchTracking(MotionEvent ev) {
+        if (Utilities.IS_RUNNING_IN_TEST_HARNESS) {
+            Log.d(TestProtocol.L3_SWIPE_TO_HOME, "2");
+        }
         Object traceToken = TraceHelper.INSTANCE.beginSection(UP_EVT,
                 FLAG_CHECK_FOR_RACE_CONDITIONS);
 

From 98314d0d4be52c93975ac6c8c2eb4e14a0ea6d09 Mon Sep 17 00:00:00 2001
From: Vinit Nayak <peanutbutter@google.com>
Date: Mon, 23 Aug 2021 16:29:42 -0700
Subject: [PATCH 05/10] Rely on presense of divider target to determine split
 screen state

* Don't rely on the number of leashes, since an app with
assistant invoked returns multiple remote app targets
even though the device isn't in split screen

Fixes: 197293347
Test: Repro steps in bug don't cause crash.
Less fatal bugs need to be addressed,
TODO(b/197568823)

Change-Id: I3432e3d8c54a433dd38d297db73ea3d46bad7ba9
---
 .../com/android/quickstep/SwipeUpAnimationLogic.java |  2 ++
 .../src/com/android/quickstep/views/RecentsView.java | 12 ++++++------
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java b/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java
index ec51599fbf..1f57e9962a 100644
--- a/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java
+++ b/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java
@@ -278,6 +278,8 @@ public abstract class SwipeUpAnimationLogic implements
         RemoteAnimationTargetCompat primaryTaskTarget;
         RemoteAnimationTargetCompat secondaryTaskTarget;
 
+        // TODO(b/197568823) Determine if we need to exclude assistant as one of the targets we
+        //  animate
         if (!mIsSwipeForStagedSplit) {
             primaryTaskTarget = targets.findTask(mGestureState.getRunningTaskId());
             mRemoteTargetHandles[0].mTransformParams.setTargetSet(targets);
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index abeadfd29b..6094c2f760 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -4019,12 +4019,14 @@ public abstract class RecentsView<ACTIVITY_TYPE extends StatefulActivity<STATE_T
             recentsAnimationTargets.addReleaseCheck(mSyncTransactionApplier);
         }
 
-        // TODO Consolidate this shared code with SwipeUpAnimationLogic (or mabe just reuse
-        //  what that class has and pass it into here
-        mRemoteTargetHandles = new RemoteTargetHandle[recentsAnimationTargets.apps.length];
+        RemoteAnimationTargetCompat dividerTarget =
+                recentsAnimationTargets.getNonAppTargetOfType(TYPE_DOCK_DIVIDER);
+        // TODO Consolidate this shared code with SwipeUpAnimationLogic (or maybe just reuse
+        //  what that class has and pass it into here)
+        mRemoteTargetHandles = new RemoteTargetHandle[dividerTarget == null ? 1 : 2];
         TaskViewSimulator primaryTvs = createTaskViewSimulator();
         mRemoteTargetHandles[0] = new RemoteTargetHandle(primaryTvs, new TransformParams());
-        if (recentsAnimationTargets.apps.length == 1) {
+        if (dividerTarget == null) {
             mRemoteTargetHandles[0].mTaskViewSimulator
                     .setPreview(recentsAnimationTargets.apps[0], null);
             mRemoteTargetHandles[0].mTransformParams.setTargetSet(recentsAnimationTargets);
@@ -4034,8 +4036,6 @@ public abstract class RecentsView<ACTIVITY_TYPE extends StatefulActivity<STATE_T
             secondaryTvs.recentsViewScale.value = 1;
 
             mRemoteTargetHandles[1] = new RemoteTargetHandle(secondaryTvs, new TransformParams());
-            RemoteAnimationTargetCompat dividerTarget =
-                    recentsAnimationTargets.getNonAppTargetOfType(TYPE_DOCK_DIVIDER);
             RemoteAnimationTargetCompat primaryTaskTarget = recentsAnimationTargets.apps[0];
             RemoteAnimationTargetCompat secondaryTaskTarget = recentsAnimationTargets.apps[1];
             SplitConfigurationOptions.StagedSplitBounds

From a0cc0903e8295a8febebd771503b74c9d302b14e Mon Sep 17 00:00:00 2001
From: Andy Wickham <awickham@google.com>
Date: Wed, 18 Aug 2021 15:17:54 -1000
Subject: [PATCH 06/10] Adds additional pages to Taskbar Edu.

Currently there are pages for the following:
 - Splitscreen
 - Customize (add apps/predicted apps)
 - Docking (long press to hide)

At the moment they all use the same placeholder
image from before.

Button states:
 Page 1: Close and Next
 Page 2: Back and Next
 Page 3: Back and Done

You can also swipe left and right between the pages.

Demo: https://drive.google.com/file/d/1_D3i-jZxCRRVHV92p6hG5cm3VGcJ_bhK/view?usp=sharing&resourcekey=0-KHLHTTx67JlmVv-UZoAUAw

Bug: 180605356
Test: Manual
Change-Id: Ibbb81610a611f6ca412e53ed90dc1c67764f417e
---
 quickstep/res/layout/taskbar_edu.xml          | 120 ++++++++++++++----
 quickstep/res/values/strings.xml              |  24 +++-
 quickstep/res/values/styles.xml               |   2 +
 .../taskbar/TaskbarEduController.java         |  28 +++-
 .../taskbar/TaskbarEduPagedView.java          |  79 ++++++++++++
 .../launcher3/taskbar/TaskbarEduView.java     |  38 +++++-
 6 files changed, 260 insertions(+), 31 deletions(-)
 create mode 100644 quickstep/src/com/android/launcher3/taskbar/TaskbarEduPagedView.java

diff --git a/quickstep/res/layout/taskbar_edu.xml b/quickstep/res/layout/taskbar_edu.xml
index b7717b71be..ef57a53690 100644
--- a/quickstep/res/layout/taskbar_edu.xml
+++ b/quickstep/res/layout/taskbar_edu.xml
@@ -17,6 +17,7 @@
 <com.android.launcher3.taskbar.TaskbarEduView xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:androidprv="http://schemas.android.com/apk/prv/res/android"
     xmlns:app="http://schemas.android.com/apk/res-auto"
+    xmlns:launcher="http://schemas.android.com/apk/res-auto"
     android:layout_width="match_parent"
     android:layout_height="wrap_content"
     android:layout_gravity="bottom"
@@ -33,46 +34,121 @@
         android:paddingHorizontal="36dp"
         android:paddingTop="64dp">
 
-        <TextView
-            android:id="@+id/edu_header"
+        <com.android.launcher3.taskbar.TaskbarEduPagedView
+            android:id="@+id/content"
+            android:clipToPadding="false"
             android:layout_width="match_parent"
-            android:layout_height="wrap_content"
-            android:layout_marginHorizontal="16dp"
+            android:layout_height="378dp"
             app:layout_constraintTop_toTopOf="parent"
-            android:gravity="center_horizontal"
-            style="@style/TextHeadline"
-            android:text="@string/taskbar_edu_header_1"
-            android:fontFamily="google-sans"
-            android:textColor="?android:attr/textColorPrimary"
-            android:textSize="24sp"
-            android:maxLines="2"/>
+            launcher:pageIndicator="@+id/content_page_indicator">
 
-        <ImageView
-            android:id="@+id/edu_illustration"
-            android:layout_width="match_parent"
-            android:layout_height="wrap_content"
-            android:layout_marginTop="16dp"
-            app:layout_constraintTop_toBottomOf="@id/edu_header"
-            android:src="@drawable/taskbar_edu_splitscreen"/>
+            <LinearLayout
+                android:id="@+id/page_splitscreen"
+                android:layout_width="match_parent"
+                android:layout_height="wrap_content"
+                android:orientation="vertical"
+                android:gravity="center_horizontal">
+
+                <TextView
+                    android:layout_width="match_parent"
+                    android:layout_height="wrap_content"
+                    android:layout_marginHorizontal="16dp"
+                    android:gravity="center_horizontal"
+                    style="@style/TextHeadline"
+                    android:text="@string/taskbar_edu_splitscreen"
+                    android:fontFamily="google-sans"
+                    android:textColor="?android:attr/textColorPrimary"
+                    android:textSize="24sp"
+                    android:maxLines="2"/>
+
+                <ImageView
+                    android:layout_width="322dp"
+                    android:layout_height="282dp"
+                    android:layout_marginTop="16dp"
+                    android:src="@drawable/taskbar_edu_splitscreen"/>
+            </LinearLayout>
+
+            <LinearLayout
+                android:id="@+id/page_customize"
+                android:layout_width="match_parent"
+                android:layout_height="wrap_content"
+                android:orientation="vertical"
+                android:gravity="center_horizontal">
+
+                <TextView
+                    android:layout_width="match_parent"
+                    android:layout_height="wrap_content"
+                    android:layout_marginHorizontal="16dp"
+                    android:gravity="center_horizontal"
+                    style="@style/TextHeadline"
+                    android:text="@string/taskbar_edu_customize"
+                    android:fontFamily="google-sans"
+                    android:textColor="?android:attr/textColorPrimary"
+                    android:textSize="24sp"
+                    android:maxLines="2"/>
+
+                <ImageView
+                    android:layout_width="322dp"
+                    android:layout_height="282dp"
+                    android:layout_marginTop="16dp"
+                    android:src="@drawable/taskbar_edu_splitscreen"/>
+            </LinearLayout>
+
+            <LinearLayout
+                android:id="@+id/page_dock"
+                android:layout_width="match_parent"
+                android:layout_height="wrap_content"
+                android:orientation="vertical"
+                android:gravity="center_horizontal">
+
+                <TextView
+                    android:layout_width="match_parent"
+                    android:layout_height="wrap_content"
+                    android:layout_marginHorizontal="16dp"
+                    android:gravity="center_horizontal"
+                    style="@style/TextHeadline"
+                    android:text="@string/taskbar_edu_dock"
+                    android:fontFamily="google-sans"
+                    android:textColor="?android:attr/textColorPrimary"
+                    android:textSize="24sp"
+                    android:maxLines="2"/>
+
+                <ImageView
+                    android:layout_width="322dp"
+                    android:layout_height="282dp"
+                    android:layout_marginTop="16dp"
+                    android:src="@drawable/taskbar_edu_splitscreen"/>
+            </LinearLayout>
+        </com.android.launcher3.taskbar.TaskbarEduPagedView>
 
         <Button
-            android:id="@+id/edu_close_button"
+            android:id="@+id/edu_start_button"
             android:layout_width="wrap_content"
             android:layout_height="36dp"
             android:layout_marginBottom="92dp"
-            app:layout_constraintTop_toBottomOf="@id/edu_illustration"
+            app:layout_constraintTop_toBottomOf="@id/content"
             app:layout_constraintBottom_toBottomOf="parent"
             app:layout_constraintStart_toStartOf="parent"
             android:text="@string/taskbar_edu_close"
             style="@style/TaskbarEdu.Button.Close"
             android:textColor="?android:attr/textColorPrimary"/>
 
+        <com.android.launcher3.pageindicators.PageIndicatorDots
+            android:id="@+id/content_page_indicator"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            app:layout_constraintTop_toTopOf="@id/edu_start_button"
+            app:layout_constraintBottom_toBottomOf="@id/edu_start_button"
+            app:layout_constraintStart_toStartOf="parent"
+            app:layout_constraintEnd_toEndOf="parent"
+            android:elevation="1dp" />
+
         <Button
-            android:id="@+id/edu_next_button"
+            android:id="@+id/edu_end_button"
             android:layout_width="wrap_content"
             android:layout_height="36dp"
             android:layout_marginBottom="92dp"
-            app:layout_constraintTop_toBottomOf="@id/edu_illustration"
+            app:layout_constraintTop_toBottomOf="@id/content"
             app:layout_constraintBottom_toBottomOf="parent"
             app:layout_constraintEnd_toEndOf="parent"
             android:text="@string/taskbar_edu_next"
diff --git a/quickstep/res/values/strings.xml b/quickstep/res/values/strings.xml
index da19c95596..2916121c83 100644
--- a/quickstep/res/values/strings.xml
+++ b/quickstep/res/values/strings.xml
@@ -205,11 +205,25 @@
     <string name="gesture_tutorial_action_button_label_skip">Skip</string>
 
     <!-- ******* Taskbar Edu ******* -->
-    <!-- Text in dialog that lets a user know how they can use the taskbar on their device.
-         [CHAR_LIMIT=NONE] -->
-    <string name="taskbar_edu_header_1">Use 2 apps at once and switch apps with the taskbar</string>
-    <!-- Text on button to go to the next screen of a tutorial [CHAR_LIMIT=30] -->
+    <!-- Accessibility text spoken when the taskbar education panel appears [CHAR_LIMIT=NONE] -->
+    <string name="taskbar_edu_opened">Taskbar education appeared</string>
+    <!-- Accessibility text spoken when the taskbar education panel disappears [CHAR_LIMIT=NONE] -->
+    <string name="taskbar_edu_closed">Taskbar education closed</string>
+    <!-- Text in dialog that lets a user know how they can use the taskbar to use multiple apps on their device.
+         [CHAR_LIMIT=60] -->
+    <string name="taskbar_edu_splitscreen" translatable="false">Use 2 apps at once and switch apps with the taskbar</string>
+    <!-- Text in dialog that lets a user know how they can customize the taskbar on their device.
+         [CHAR_LIMIT=60] -->
+    <string name="taskbar_edu_customize" translatable="false">Add your favorite apps and get automatic suggestions</string>
+    <!-- Text in dialog that lets a user know how they can hide the taskbar on their device.
+         [CHAR_LIMIT=60] -->
+    <string name="taskbar_edu_dock">Touch &amp; hold to hide the taskbar anytime</string>
+    <!-- Text on button to go to the next screen of a tutorial [CHAR_LIMIT=16] -->
     <string name="taskbar_edu_next">Next</string>
-    <!-- Text on button to exit a tutorial [CHAR_LIMIT=30] -->
+    <!-- Text on button to go to the previous screen of a tutorial [CHAR_LIMIT=16] -->
+    <string name="taskbar_edu_previous">Back</string>
+    <!-- Text on button to exit a tutorial [CHAR_LIMIT=16] -->
     <string name="taskbar_edu_close">Close</string>
+    <!-- Text on button to finish a tutorial [CHAR_LIMIT=16] -->
+    <string name="taskbar_edu_done">Done</string>
 </resources>
diff --git a/quickstep/res/values/styles.xml b/quickstep/res/values/styles.xml
index c23a3a01a6..ac702794b2 100644
--- a/quickstep/res/values/styles.xml
+++ b/quickstep/res/values/styles.xml
@@ -144,11 +144,13 @@
         <item name="android:background">@drawable/button_taskbar_edu_bordered</item>
         <item name="android:stateListAnimator">@null</item>
         <item name="android:textSize">16sp</item>
+        <item name="android:padding">4dp</item>
     </style>
 
     <style name="TaskbarEdu.Button.Next" parent="@android:style/Widget.Material.Button">
         <item name="android:background">@drawable/button_taskbar_edu_colored</item>
         <item name="android:stateListAnimator">@null</item>
         <item name="android:textSize">16sp</item>
+        <item name="android:padding">4dp</item>
     </style>
 </resources>
\ No newline at end of file
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarEduController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarEduController.java
index c5a0fc1902..ae9592d779 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarEduController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarEduController.java
@@ -32,9 +32,8 @@ public class TaskbarEduController {
         mActivity.getDragLayer().post(() -> {
             mTaskbarEduView = (TaskbarEduView) mActivity.getLayoutInflater().inflate(
                     R.layout.taskbar_edu, mActivity.getDragLayer(), false);
-            mTaskbarEduView.addOnCloseListener(() -> {
-                mTaskbarEduView = null;
-            });
+            mTaskbarEduView.init(new TaskbarEduCallbacks());
+            mTaskbarEduView.addOnCloseListener(() -> mTaskbarEduView = null);
             mTaskbarEduView.show();
         });
     }
@@ -44,4 +43,27 @@ public class TaskbarEduController {
             mTaskbarEduView.close(true /* animate */);
         }
     }
+
+
+    /**
+     * Callbacks for {@link TaskbarEduView} to interact with its controller.
+     */
+    class TaskbarEduCallbacks {
+        void onPageChanged(int currentPage, int pageCount) {
+            if (currentPage == 0) {
+                mTaskbarEduView.updateStartButton(R.string.taskbar_edu_close,
+                        v -> mTaskbarEduView.close(true /* animate */));
+            } else {
+                mTaskbarEduView.updateStartButton(R.string.taskbar_edu_previous,
+                        v -> mTaskbarEduView.snapToPage(currentPage - 1));
+            }
+            if (currentPage == pageCount - 1) {
+                mTaskbarEduView.updateEndButton(R.string.taskbar_edu_done,
+                        v -> mTaskbarEduView.close(true /* animate */));
+            } else {
+                mTaskbarEduView.updateEndButton(R.string.taskbar_edu_next,
+                        v -> mTaskbarEduView.snapToPage(currentPage + 1));
+            }
+        }
+    }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarEduPagedView.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarEduPagedView.java
new file mode 100644
index 0000000000..5efcc4df4c
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarEduPagedView.java
@@ -0,0 +1,79 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.taskbar;
+
+import static com.android.launcher3.AbstractFloatingView.TYPE_ALL;
+import static com.android.launcher3.AbstractFloatingView.TYPE_TASKBAR_EDUCATION_DIALOG;
+
+import android.content.Context;
+import android.util.AttributeSet;
+
+import com.android.launcher3.AbstractFloatingView;
+import com.android.launcher3.PagedView;
+import com.android.launcher3.R;
+import com.android.launcher3.pageindicators.PageIndicatorDots;
+import com.android.launcher3.taskbar.TaskbarEduController.TaskbarEduCallbacks;
+import com.android.launcher3.views.ActivityContext;
+
+/** Horizontal carousel of tutorial screens for Taskbar Edu. */
+public class TaskbarEduPagedView extends PagedView<PageIndicatorDots> {
+
+    private TaskbarEduView mTaskbarEduView;
+    private TaskbarEduCallbacks mControllerCallbacks;
+
+    public TaskbarEduPagedView(Context context, AttributeSet attrs) {
+        super(context, attrs);
+        setImportantForAccessibility(IMPORTANT_FOR_ACCESSIBILITY_YES);
+    }
+
+    void setTaskbarEduView(TaskbarEduView taskbarEduView) {
+        mTaskbarEduView = taskbarEduView;
+        mPageIndicator = taskbarEduView.findViewById(R.id.content_page_indicator);
+        initParentViews(taskbarEduView);
+    }
+
+    void setControllerCallbacks(TaskbarEduCallbacks controllerCallbacks) {
+        mControllerCallbacks = controllerCallbacks;
+        mControllerCallbacks.onPageChanged(getCurrentPage(), getPageCount());
+    }
+
+    @Override
+    protected int getChildGap() {
+        return mTaskbarEduView.getPaddingLeft() + mTaskbarEduView.getPaddingRight();
+    }
+
+    @Override
+    protected void onScrollChanged(int l, int t, int oldl, int oldt) {
+        super.onScrollChanged(l, t, oldl, oldt);
+        if (mMaxScroll > 0) {
+            mPageIndicator.setScroll(l, mMaxScroll);
+        }
+    }
+
+    @Override
+    protected void notifyPageSwitchListener(int prevPage) {
+        super.notifyPageSwitchListener(prevPage);
+        mControllerCallbacks.onPageChanged(getCurrentPage(), getPageCount());
+    }
+
+    @Override
+    protected boolean canScroll(float absVScroll, float absHScroll) {
+        return AbstractFloatingView.getTopOpenViewWithType(
+                ActivityContext.lookupContext(getContext()),
+                TYPE_ALL & ~TYPE_TASKBAR_EDUCATION_DIALOG) == null;
+    }
+}
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarEduView.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarEduView.java
index b5dab7e770..9c4e844fa8 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarEduView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarEduView.java
@@ -21,6 +21,9 @@ import android.animation.PropertyValuesHolder;
 import android.content.Context;
 import android.graphics.Rect;
 import android.util.AttributeSet;
+import android.util.Pair;
+import android.view.View;
+import android.widget.Button;
 
 import com.android.launcher3.Insettable;
 import com.android.launcher3.R;
@@ -34,6 +37,10 @@ public class TaskbarEduView extends AbstractSlideInView<TaskbarActivityContext>
 
     private final Rect mInsets = new Rect();
 
+    private Button mStartButton;
+    private Button mEndButton;
+    private TaskbarEduPagedView mPagedView;
+
     public TaskbarEduView(Context context, AttributeSet attr) {
         this(context, attr, 0);
     }
@@ -43,6 +50,12 @@ public class TaskbarEduView extends AbstractSlideInView<TaskbarActivityContext>
         super(context, attrs, defStyleAttr);
     }
 
+    protected void init(TaskbarEduController.TaskbarEduCallbacks callbacks) {
+        if (mPagedView != null) {
+            mPagedView.setControllerCallbacks(callbacks);
+        }
+    }
+
     @Override
     protected void handleClose(boolean animate) {
         handleClose(animate, DEFAULT_CLOSE_DURATION);
@@ -57,7 +70,10 @@ public class TaskbarEduView extends AbstractSlideInView<TaskbarActivityContext>
     protected void onFinishInflate() {
         super.onFinishInflate();
         mContent = findViewById(R.id.edu_view);
-        findViewById(R.id.edu_close_button).setOnClickListener(v -> close(true /* animate */));
+        mStartButton = findViewById(R.id.edu_start_button);
+        mEndButton = findViewById(R.id.edu_end_button);
+        mPagedView = findViewById(R.id.content);
+        mPagedView.setTaskbarEduView(this);
     }
 
     @Override
@@ -81,6 +97,12 @@ public class TaskbarEduView extends AbstractSlideInView<TaskbarActivityContext>
         animateOpen();
     }
 
+    @Override
+    protected Pair<View, String> getAccessibilityTarget() {
+        return Pair.create(mContent, mIsOpen ? getContext().getString(R.string.taskbar_edu_opened)
+                : getContext().getString(R.string.taskbar_edu_closed));
+    }
+
     @Override
     protected int getScrimColor(Context context) {
         return context.getResources().getColor(R.color.widgets_picker_scrim);
@@ -110,4 +132,18 @@ public class TaskbarEduView extends AbstractSlideInView<TaskbarActivityContext>
         mOpenCloseAnimator.setInterpolator(FAST_OUT_SLOW_IN);
         mOpenCloseAnimator.start();
     }
+
+    void snapToPage(int page) {
+        mPagedView.snapToPage(page);
+    }
+
+    void updateStartButton(int textResId, OnClickListener onClickListener) {
+        mStartButton.setText(textResId);
+        mStartButton.setOnClickListener(onClickListener);
+    }
+
+    void updateEndButton(int textResId, OnClickListener onClickListener) {
+        mEndButton.setText(textResId);
+        mEndButton.setOnClickListener(onClickListener);
+    }
 }

From 6ca6de89c0aaf3da2c255b5e1a48458103f626b6 Mon Sep 17 00:00:00 2001
From: Jerry Chang <chenghsiuchang@google.com>
Date: Tue, 24 Aug 2021 04:48:55 +0000
Subject: [PATCH 07/10] Revert "Move/Launch task on the TaskView's display."

This reverts commit 2cdc359b84c0a9c5fae396220cf8e2b46f8c6042.

Reason for revert: this brokes split screen and result to device reboot.
Bug: 197293995

Change-Id: I154287230fa3849cad1c0e337e3bf3fd826257a4
---
 quickstep/src/com/android/quickstep/views/TaskView.java | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index dbda9490ca..a3b9c6f453 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -654,7 +654,6 @@ public class TaskView extends FrameLayout implements Reusable {
             TestLogging.recordEvent(
                     TestProtocol.SEQUENCE_MAIN, "startActivityFromRecentsAsync", mTask);
             ActivityOptionsWrapper opts =  mActivity.getActivityLaunchOptions(this, null);
-            opts.options.setLaunchDisplayId(getRootViewDisplayId());
             if (ActivityManagerWrapper.getInstance()
                     .startActivityFromRecents(mTask.key, opts.options)) {
                 RecentsView recentsView = getRecentsView();
@@ -695,7 +694,6 @@ public class TaskView extends FrameLayout implements Reusable {
             // Indicate success once the system has indicated that the transition has started
             ActivityOptions opts = ActivityOptionsCompat.makeCustomAnimation(
                     getContext(), 0, 0, () -> callback.accept(true), MAIN_EXECUTOR.getHandler());
-            opts.setLaunchDisplayId(getRootViewDisplayId());
             if (freezeTaskList) {
                 ActivityOptionsCompat.setFreezeRecentTasksList(opts);
             }
@@ -1507,10 +1505,6 @@ public class TaskView extends FrameLayout implements Reusable {
         mDigitalWellBeingToast.setBannerColorTint(tintColor, amount);
     }
 
-    private int getRootViewDisplayId() {
-        return getRootView().getDisplay().getDisplayId();
-    }
-
     /**
      * We update and subsequently draw these in {@link #setFullscreenProgress(float)}.
      */

From 96ffcbcd2d39eb69e8e8db4cd8bed83c4448042b Mon Sep 17 00:00:00 2001
From: Alex Chau <alexchau@google.com>
Date: Tue, 24 Aug 2021 12:36:27 +0100
Subject: [PATCH 08/10] Avoid unnecessary onLayout if gesture is going to state
 without overview panel

Fixes: 186508379
Test: swipe up to carousel, then to home
Test: quick switch, immediately swipe up to overview, then to home
Change-Id: I0b28e163cbb289a636a89bc52c43a1977a59d86f
---
 quickstep/src/com/android/quickstep/views/RecentsView.java | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index abeadfd29b..dc7347e9b0 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -3314,9 +3314,9 @@ public abstract class RecentsView<ACTIVITY_TYPE extends StatefulActivity<STATE_T
 
     @Override
     protected void onLayout(boolean changed, int left, int top, int right, int bottom) {
-        // If we're going to HOME, avoid unnecessary onLayout that cause TaskViews to re-arrange
-        // during animation to HOME.
-        if (mCurrentGestureEndTarget == GestureState.GestureEndTarget.HOME) {
+        // If we're going to a state without overview panel, avoid unnecessary onLayout that
+        // cause TaskViews to re-arrange during animation to that state.
+        if (!mOverviewStateEnabled && !mFirstLayout) {
             return;
         }
 

From 51ae38304bf11ed8596e7b4aaec7ca3506c5cc99 Mon Sep 17 00:00:00 2001
From: vadimt <vadimt@google.com>
Date: Fri, 20 Aug 2021 14:02:45 -0700
Subject: [PATCH 09/10] Not checking for some anomalies when switching nav
 modes

Test: presubmit
Bug: 196820244
Change-Id: I3e9554d3deb453625cde5003ccf132f21253f58a
---
 .../quickstep/NavigationModeSwitchRule.java   |  8 +----
 .../src/com/android/launcher3/util/Wait.java  |  2 +-
 .../tapl/LauncherInstrumentation.java         | 34 +++++++++++--------
 3 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
index bc8602c940..416c193cfc 100644
--- a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
+++ b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
@@ -30,7 +30,6 @@ import android.content.Context;
 import android.content.pm.PackageManager;
 import android.util.Log;
 
-import androidx.test.uiautomator.By;
 import androidx.test.uiautomator.UiDevice;
 
 import com.android.launcher3.tapl.LauncherInstrumentation;
@@ -217,12 +216,7 @@ public class NavigationModeSwitchRule implements TestRule {
 
     private static void assertTrue(LauncherInstrumentation launcher, String message,
             boolean condition, Description description) {
-        if (launcher.getDevice().hasObject(By.textStartsWith(""))) {
-            // The condition above is "screen is not empty". We are not treating
-            // "Screen is empty" as an anomaly here. It's an acceptable state when
-            // Launcher just starts under instrumentation.
-            launcher.checkForAnomaly();
-        }
+        launcher.checkForAnomaly(true);
         if (!condition) {
             final AssertionError assertionError = new AssertionError(message);
             if (description != null) {
diff --git a/tests/src/com/android/launcher3/util/Wait.java b/tests/src/com/android/launcher3/util/Wait.java
index fe6143c681..8b5e1970de 100644
--- a/tests/src/com/android/launcher3/util/Wait.java
+++ b/tests/src/com/android/launcher3/util/Wait.java
@@ -52,7 +52,7 @@ public class Wait {
             throw new RuntimeException(t);
         }
         Log.d("Wait", "atMost: timed out: " + SystemClock.uptimeMillis());
-        launcher.checkForAnomaly();
+        launcher.checkForAnomaly(false);
         Assert.fail(message.get());
     }
 
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 2ca40d8229..f273d6f34b 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -305,7 +305,7 @@ public final class LauncherInstrumentation {
 
     public boolean isTwoPanels() {
         return getTestInfo(TestProtocol.REQUEST_IS_TWO_PANELS)
-            .getBoolean(TestProtocol.TEST_INFO_RESPONSE_FIELD);
+                .getBoolean(TestProtocol.TEST_INFO_RESPONSE_FIELD);
     }
 
     private void setForcePauseTimeout(long timeout) {
@@ -379,7 +379,7 @@ public final class LauncherInstrumentation {
         }
     }
 
-    private String getSystemAnomalyMessage() {
+    private String getSystemAnomalyMessage(boolean ignoreNavmodeChangeStates) {
         try {
             {
                 final StringBuilder sb = new StringBuilder();
@@ -402,16 +402,18 @@ public final class LauncherInstrumentation {
 
             if (hasSystemUiObject("keyguard_status_view")) return "Phone is locked";
 
-            final String visibleApps = mDevice.findObjects(getAnyObjectSelector())
-                    .stream()
-                    .map(LauncherInstrumentation::getApplicationPackageSafe)
-                    .distinct()
-                    .filter(pkg -> pkg != null)
-                    .collect(Collectors.joining(","));
-            if (SYSTEMUI_PACKAGE.equals(visibleApps)) return "Only System UI views are visible";
+            if (!ignoreNavmodeChangeStates) {
+                final String visibleApps = mDevice.findObjects(getAnyObjectSelector())
+                        .stream()
+                        .map(LauncherInstrumentation::getApplicationPackageSafe)
+                        .distinct()
+                        .filter(pkg -> pkg != null)
+                        .collect(Collectors.joining(","));
+                if (SYSTEMUI_PACKAGE.equals(visibleApps)) return "Only System UI views are visible";
 
-            if (!mDevice.wait(Until.hasObject(getAnyObjectSelector()), WAIT_TIME_MS)) {
-                return "Screen is empty";
+                if (!mDevice.wait(Until.hasObject(getAnyObjectSelector()), WAIT_TIME_MS)) {
+                    return "Screen is empty";
+                }
             }
 
             final String navigationModeError = getNavigationModeMismatchError(true);
@@ -423,8 +425,12 @@ public final class LauncherInstrumentation {
         return null;
     }
 
-    public void checkForAnomaly() {
-        final String systemAnomalyMessage = getSystemAnomalyMessage();
+    private void checkForAnomaly() {
+        checkForAnomaly(false);
+    }
+
+    public void checkForAnomaly(boolean ignoreNavmodeChangeStates) {
+        final String systemAnomalyMessage = getSystemAnomalyMessage(ignoreNavmodeChangeStates);
         if (systemAnomalyMessage != null) {
             Assert.fail(formatSystemHealthMessage(formatErrorWithEvents(
                     "http://go/tapl : Tests are broken by a non-Launcher system error: "
@@ -756,7 +762,7 @@ public final class LauncherInstrumentation {
 
                 if (hasLauncherObject(CONTEXT_MENU_RES_ID)) {
                     GestureScope gestureScope = gestureStartFromLauncher
-                            ? (isTablet()? GestureScope.INSIDE : GestureScope.INSIDE_TO_OUTSIDE)
+                            ? (isTablet() ? GestureScope.INSIDE : GestureScope.INSIDE_TO_OUTSIDE)
                             : GestureScope.OUTSIDE_WITH_PILFER;
                     linearGesture(
                             displaySize.x / 2, displaySize.y - 1,

From 7a4a30d86d471e6c45adc2a9907efb27e9b1799b Mon Sep 17 00:00:00 2001
From: Alex Chau <alexchau@google.com>
Date: Tue, 24 Aug 2021 13:46:47 +0000
Subject: [PATCH 10/10] Revert "Migrating all model tests to Instrumentation
 tests"

This reverts commit 455ba62b5db7a8415e0f250d62c30e9c24ebbb6b.

Reason for revert: Breaks tests on all targets

Bug: 196825541
Bug: 197570389
Change-Id: Ida6792a09c54dee46105b6b914a0dd1545bdf48d
---
 .../WidgetsPredicationUpdateTaskTest.java     | 128 ++++----
 .../cache_data_updated_task_data.txt          |   0
 .../resources}/db_schema_v10.json             |   0
 ...package_install_state_change_task_data.txt |   0
 .../widgets_predication_update_task_data.txt  |   0
 .../model/AddWorkspaceItemsTaskTest.java      |  24 +-
 .../launcher3/model/BackupRestoreTest.java    | 111 +++----
 .../model/CacheDataUpdatedTaskTest.java       |  20 +-
 .../model/DbDowngradeHelperTest.java          |  20 +-
 .../model/DefaultLayoutProviderTest.java      |  45 +--
 .../model/GridSizeMigrationTaskV2Test.java    |  69 ++--
 .../launcher3/model/LoaderCursorTest.java     |  32 +-
 .../model/ModelMultiCallbacksTest.java        | 123 +++----
 .../PackageInstallStateChangedTaskTest.java   |  20 +-
 .../launcher3/provider/RestoreDbTaskTest.java |  11 +-
 .../secondarydisplay/SDWorkModeTest.java      | 125 ++++++++
 .../launcher3/testing/TestActivity.java       |  46 +++
 .../launcher3/util/LauncherLayoutBuilder.java |   0
 .../launcher3/util/LauncherModelHelper.java   | 299 +++++-------------
 .../launcher3/util/LauncherUIHelper.java      | 101 ++++++
 .../android/launcher3/LauncherAppState.java   |   5 +-
 src/com/android/launcher3/LauncherModel.java  |   7 +-
 .../graphics/LauncherPreviewRenderer.java     |  50 ++-
 .../android/launcher3/icons/IconCache.java    |   3 -
 .../launcher3/util/DisplayController.java     |  17 +-
 .../util/MainThreadInitializedObject.java     |  87 +----
 .../android/launcher3/util/SettingsCache.java |   8 +-
 .../launcher3/util/UiThreadHelper.java        |   4 +-
 .../widget/custom/CustomWidgetManager.java    |   6 +-
 .../secondarydisplay/SDLauncherTest.java      |  53 ----
 .../launcher3/util/ReflectionHelpers.java     |  58 ----
 31 files changed, 686 insertions(+), 786 deletions(-)
 rename quickstep/{tests => robolectric_tests}/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java (65%)
 rename {tests/res/raw => robolectric_tests/resources}/cache_data_updated_task_data.txt (100%)
 rename {tests/res/raw => robolectric_tests/resources}/db_schema_v10.json (100%)
 rename {tests/res/raw => robolectric_tests/resources}/package_install_state_change_task_data.txt (100%)
 rename {tests/res/raw => robolectric_tests/resources}/widgets_predication_update_task_data.txt (100%)
 rename {tests => robolectric_tests}/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java (93%)
 rename {tests => robolectric_tests}/src/com/android/launcher3/model/BackupRestoreTest.java (62%)
 rename {tests => robolectric_tests}/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java (92%)
 rename {tests => robolectric_tests}/src/com/android/launcher3/model/DbDowngradeHelperTest.java (91%)
 rename {tests => robolectric_tests}/src/com/android/launcher3/model/DefaultLayoutProviderTest.java (77%)
 rename {tests => robolectric_tests}/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java (89%)
 rename {tests => robolectric_tests}/src/com/android/launcher3/model/LoaderCursorTest.java (93%)
 rename {tests => robolectric_tests}/src/com/android/launcher3/model/ModelMultiCallbacksTest.java (65%)
 rename {tests => robolectric_tests}/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java (87%)
 rename {tests => robolectric_tests}/src/com/android/launcher3/provider/RestoreDbTaskTest.java (93%)
 create mode 100644 robolectric_tests/src/com/android/launcher3/secondarydisplay/SDWorkModeTest.java
 create mode 100644 robolectric_tests/src/com/android/launcher3/testing/TestActivity.java
 rename {tests => robolectric_tests}/src/com/android/launcher3/util/LauncherLayoutBuilder.java (100%)
 rename {tests => robolectric_tests}/src/com/android/launcher3/util/LauncherModelHelper.java (60%)
 create mode 100644 robolectric_tests/src/com/android/launcher3/util/LauncherUIHelper.java
 delete mode 100644 tests/src/com/android/launcher3/secondarydisplay/SDLauncherTest.java
 delete mode 100644 tests/src/com/android/launcher3/util/ReflectionHelpers.java

diff --git a/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java b/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
similarity index 65%
rename from quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
rename to quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
index c1b3beb475..f82fbcc15b 100644
--- a/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
+++ b/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
@@ -15,31 +15,24 @@
  */
 package com.android.launcher3.model;
 
-import static android.os.Process.myUserHandle;
-
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
-import static com.android.launcher3.util.WidgetUtils.createAppWidgetProviderInfo;
 
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
-import static org.mockito.Mockito.doReturn;
+import static org.robolectric.Shadows.shadowOf;
 
 import android.app.prediction.AppTarget;
 import android.app.prediction.AppTargetId;
 import android.appwidget.AppWidgetManager;
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ComponentName;
+import android.content.Context;
+import android.os.Process;
 import android.os.UserHandle;
-import android.text.TextUtils;
-import android.util.Log;
-
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
 
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.config.FeatureFlags;
@@ -47,34 +40,38 @@ import com.android.launcher3.icons.ComponentWithLabel;
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
 import com.android.launcher3.model.QuickstepModelDelegate.PredictorState;
+import com.android.launcher3.shadows.ShadowDeviceFlag;
 import com.android.launcher3.util.LauncherModelHelper;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.PendingAddWidgetInfo;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.shadow.api.Shadow;
+import org.robolectric.shadows.ShadowAppWidgetManager;
+import org.robolectric.shadows.ShadowPackageManager;
+import org.robolectric.util.ReflectionHelpers;
 
-import java.util.Arrays;
 import java.util.List;
 import java.util.stream.Collectors;
 
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
 public final class WidgetsPredicationUpdateTaskTest {
 
-    private AppWidgetProviderInfo mApp1Provider1;
-    private AppWidgetProviderInfo mApp1Provider2;
-    private AppWidgetProviderInfo mApp2Provider1;
-    private AppWidgetProviderInfo mApp4Provider1;
-    private AppWidgetProviderInfo mApp4Provider2;
-    private AppWidgetProviderInfo mApp5Provider1;
-    private List<AppWidgetProviderInfo> allWidgets;
+    private AppWidgetProviderInfo mApp1Provider1 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp1Provider2 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp2Provider1 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp4Provider1 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp4Provider2 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp5Provider1 = new AppWidgetProviderInfo();
 
     private FakeBgDataModelCallback mCallback = new FakeBgDataModelCallback();
+    private Context mContext;
     private LauncherModelHelper mModelHelper;
     private UserHandle mUserHandle;
 
@@ -83,56 +80,54 @@ public final class WidgetsPredicationUpdateTaskTest {
 
     @Before
     public void setup() throws Exception {
-        mModelHelper = new LauncherModelHelper();
         MockitoAnnotations.initMocks(this);
         doAnswer(invocation -> {
             ComponentWithLabel componentWithLabel = invocation.getArgument(0);
             return componentWithLabel.getComponent().getShortClassName();
         }).when(mIconCache).getTitleNoCache(any());
 
-        mUserHandle = myUserHandle();
-        mApp1Provider1 = createAppWidgetProviderInfo(
-                ComponentName.createRelative("app1", "provider1"));
-        mApp1Provider2 = createAppWidgetProviderInfo(
-                ComponentName.createRelative("app1", "provider2"));
-        mApp2Provider1 = createAppWidgetProviderInfo(
-                ComponentName.createRelative("app2", "provider1"));
-        mApp4Provider1 = createAppWidgetProviderInfo(
-                ComponentName.createRelative("app4", "provider1"));
-        mApp4Provider2 = createAppWidgetProviderInfo(
-                ComponentName.createRelative("app4", ".provider2"));
-        mApp5Provider1 = createAppWidgetProviderInfo(
-                ComponentName.createRelative("app5", "provider1"));
-        allWidgets = Arrays.asList(mApp1Provider1, mApp1Provider2, mApp2Provider1,
-                mApp4Provider1, mApp4Provider2, mApp5Provider1);
-
-        AppWidgetManager manager = mModelHelper.sandboxContext.spyService(AppWidgetManager.class);
-        doReturn(allWidgets).when(manager).getInstalledProviders();
-        doReturn(allWidgets).when(manager).getInstalledProvidersForProfile(eq(myUserHandle()));
-        doAnswer(i -> {
-            String pkg = i.getArgument(0);
-            Log.e("Hello", "Getting v " + pkg);
-            return TextUtils.isEmpty(pkg) ? allWidgets : allWidgets.stream()
-                    .filter(a -> pkg.equals(a.provider.getPackageName()))
-                    .collect(Collectors.toList());
-        }).when(manager).getInstalledProvidersForPackage(any(), eq(myUserHandle()));
-
+        mContext = RuntimeEnvironment.application;
+        mModelHelper = new LauncherModelHelper();
+        mUserHandle = Process.myUserHandle();
         // 2 widgets, app4/provider1 & app5/provider1, have already been added to the workspace.
-        mModelHelper.initializeData("widgets_predication_update_task_data");
+        mModelHelper.initializeData("/widgets_predication_update_task_data.txt");
+
+        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
+        mApp1Provider1.provider = ComponentName.createRelative("app1", "provider1");
+        ReflectionHelpers.setField(mApp1Provider1, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp1Provider1.provider));
+        mApp1Provider2.provider = ComponentName.createRelative("app1", "provider2");
+        ReflectionHelpers.setField(mApp1Provider2, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp1Provider2.provider));
+        mApp2Provider1.provider = ComponentName.createRelative("app2", "provider1");
+        ReflectionHelpers.setField(mApp2Provider1, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp2Provider1.provider));
+        mApp4Provider1.provider = ComponentName.createRelative("app4", "provider1");
+        ReflectionHelpers.setField(mApp4Provider1, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp4Provider1.provider));
+        mApp4Provider2.provider = ComponentName.createRelative("app4", ".provider2");
+        ReflectionHelpers.setField(mApp4Provider2, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp4Provider2.provider));
+        mApp5Provider1.provider = ComponentName.createRelative("app5", "provider1");
+        ReflectionHelpers.setField(mApp5Provider1, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp5Provider1.provider));
+
+        ShadowAppWidgetManager shadowAppWidgetManager =
+                shadowOf(mContext.getSystemService(AppWidgetManager.class));
+        shadowAppWidgetManager.addInstalledProvider(mApp1Provider1);
+        shadowAppWidgetManager.addInstalledProvider(mApp1Provider2);
+        shadowAppWidgetManager.addInstalledProvider(mApp2Provider1);
+        shadowAppWidgetManager.addInstalledProvider(mApp4Provider1);
+        shadowAppWidgetManager.addInstalledProvider(mApp4Provider2);
+        shadowAppWidgetManager.addInstalledProvider(mApp5Provider1);
+
+        mModelHelper.getModel().addCallbacks(mCallback);
 
-        MAIN_EXECUTOR.submit(() -> mModelHelper.getModel().addCallbacks(mCallback)).get();
         MODEL_EXECUTOR.post(() -> mModelHelper.getBgDataModel().widgetsModel.update(
-                LauncherAppState.getInstance(mModelHelper.sandboxContext),
-                /* packageUser= */ null));
-
-        MODEL_EXECUTOR.submit(() -> { }).get();
-        MAIN_EXECUTOR.submit(() -> { }).get();
+                LauncherAppState.getInstance(mContext), /* packageUser= */ null));
+        waitUntilIdle();
     }
 
-    @After
-    public void tearDown() {
-        mModelHelper.destroy();
-    }
 
     @Test
     public void widgetsRecommendationRan_shouldOnlyReturnNotAddedWidgetsInAppPredictionOrder()
@@ -170,9 +165,9 @@ public final class WidgetsPredicationUpdateTaskTest {
     @Test
     public void widgetsRecommendationRan_localFilterDisabled_shouldReturnWidgetsInPredicationOrder()
             throws Exception {
-        if (FeatureFlags.ENABLE_LOCAL_RECOMMENDED_WIDGETS_FILTER.get()) {
-            return;
-        }
+        ShadowDeviceFlag shadowDeviceFlag = Shadow.extract(
+                FeatureFlags.ENABLE_LOCAL_RECOMMENDED_WIDGETS_FILTER);
+        shadowDeviceFlag.setValue(false);
 
         // WHEN newPredicationTask is executed with 5 predicated widgets.
         AppTarget widget1 = new AppTarget(new AppTargetId("app1"), "app1", "provider1",
@@ -208,8 +203,13 @@ public final class WidgetsPredicationUpdateTaskTest {
         assertThat(actual.getUser()).isEqualTo(expected.getProfile());
     }
 
+    private void waitUntilIdle() {
+        shadowOf(MODEL_EXECUTOR.getLooper()).idle();
+        shadowOf(MAIN_EXECUTOR.getLooper()).idle();
+    }
+
     private WidgetsPredictionUpdateTask newWidgetsPredicationTask(List<AppTarget> appTargets) {
-       return new WidgetsPredictionUpdateTask(
+        return new WidgetsPredictionUpdateTask(
                 new PredictorState(CONTAINER_WIDGETS_PREDICTION, "test_widgets_prediction"),
                 appTargets);
     }
diff --git a/tests/res/raw/cache_data_updated_task_data.txt b/robolectric_tests/resources/cache_data_updated_task_data.txt
similarity index 100%
rename from tests/res/raw/cache_data_updated_task_data.txt
rename to robolectric_tests/resources/cache_data_updated_task_data.txt
diff --git a/tests/res/raw/db_schema_v10.json b/robolectric_tests/resources/db_schema_v10.json
similarity index 100%
rename from tests/res/raw/db_schema_v10.json
rename to robolectric_tests/resources/db_schema_v10.json
diff --git a/tests/res/raw/package_install_state_change_task_data.txt b/robolectric_tests/resources/package_install_state_change_task_data.txt
similarity index 100%
rename from tests/res/raw/package_install_state_change_task_data.txt
rename to robolectric_tests/resources/package_install_state_change_task_data.txt
diff --git a/tests/res/raw/widgets_predication_update_task_data.txt b/robolectric_tests/resources/widgets_predication_update_task_data.txt
similarity index 100%
rename from tests/res/raw/widgets_predication_update_task_data.txt
rename to robolectric_tests/resources/widgets_predication_update_task_data.txt
diff --git a/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java b/robolectric_tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
similarity index 93%
rename from tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
rename to robolectric_tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
index 8a4590a388..8aa6f374aa 100644
--- a/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
@@ -13,9 +13,6 @@ import android.content.Intent;
 import android.graphics.Rect;
 import android.util.Pair;
 
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherSettings;
@@ -24,17 +21,19 @@ import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.util.ContentWriter;
-import com.android.launcher3.util.Executors;
 import com.android.launcher3.util.GridOccupancy;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSparseArrayMap;
 import com.android.launcher3.util.LauncherModelHelper;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.annotation.LooperMode.Mode;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -42,8 +41,8 @@ import java.util.List;
 /**
  * Tests for {@link AddWorkspaceItemsTask}
  */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
+@LooperMode(Mode.PAUSED)
 public class AddWorkspaceItemsTaskTest {
 
     private final ComponentName mComponent1 = new ComponentName("a", "b");
@@ -61,7 +60,7 @@ public class AddWorkspaceItemsTaskTest {
     @Before
     public void setup() {
         mModelHelper = new LauncherModelHelper();
-        mTargetContext = mModelHelper.sandboxContext;
+        mTargetContext = RuntimeEnvironment.application;
         mIdp = InvariantDeviceProfile.INSTANCE.get(mTargetContext);
         mIdp.numColumns = mIdp.numRows = 5;
         mAppState = LauncherAppState.getInstance(mTargetContext);
@@ -71,11 +70,6 @@ public class AddWorkspaceItemsTaskTest {
         mNewScreens = new IntArray();
     }
 
-    @After
-    public void tearDown() {
-        mModelHelper.destroy();
-    }
-
     private AddWorkspaceItemsTask newTask(ItemInfo... items) {
         List<Pair<ItemInfo, Object>> list = new ArrayList<>();
         for (ItemInfo item : items) {
@@ -94,7 +88,7 @@ public class AddWorkspaceItemsTaskTest {
 
         int[] spaceFound = newTask().findSpaceForItem(
                 mAppState, mModelHelper.getBgDataModel(), mExistingScreens, mNewScreens, 1, 1);
-        assertEquals(1, spaceFound[0]);
+        assertEquals(2, spaceFound[0]);
         assertTrue(mScreenOccupancy.get(spaceFound[0])
                 .isRegionVacant(spaceFound[1], spaceFound[2], 1, 1));
 
@@ -133,7 +127,7 @@ public class AddWorkspaceItemsTaskTest {
     @Test
     public void testAddItem_some_items_added() throws Exception {
         Callbacks callbacks = mock(Callbacks.class);
-        Executors.MAIN_EXECUTOR.submit(() -> mModelHelper.getModel().addCallbacks(callbacks)).get();
+        mModelHelper.getModel().addCallbacks(callbacks);
 
         WorkspaceItemInfo info = new WorkspaceItemInfo();
         info.intent = new Intent().setComponent(mComponent1);
diff --git a/tests/src/com/android/launcher3/model/BackupRestoreTest.java b/robolectric_tests/src/com/android/launcher3/model/BackupRestoreTest.java
similarity index 62%
rename from tests/src/com/android/launcher3/model/BackupRestoreTest.java
rename to robolectric_tests/src/com/android/launcher3/model/BackupRestoreTest.java
index 41914de1ab..a397db565d 100644
--- a/tests/src/com/android/launcher3/model/BackupRestoreTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/BackupRestoreTest.java
@@ -17,7 +17,6 @@
 package com.android.launcher3.model;
 
 import static android.content.pm.PackageManager.INSTALL_REASON_DEVICE_RESTORE;
-import static android.os.Process.myUserHandle;
 
 import static com.android.launcher3.LauncherSettings.Favorites.BACKUP_TABLE_NAME;
 import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
@@ -27,110 +26,74 @@ import static com.android.launcher3.provider.LauncherDbUtils.tableExists;
 import static com.android.launcher3.util.LauncherModelHelper.APP_ICON;
 import static com.android.launcher3.util.LauncherModelHelper.NO__ICON;
 import static com.android.launcher3.util.LauncherModelHelper.SHORTCUT;
-import static com.android.launcher3.util.ReflectionHelpers.getField;
-import static com.android.launcher3.util.ReflectionHelpers.setField;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.doReturn;
-import static org.mockito.Mockito.spy;
+import static org.robolectric.util.ReflectionHelpers.setField;
 
 import android.app.backup.BackupManager;
 import android.content.pm.PackageInstaller;
 import android.database.Cursor;
 import android.database.sqlite.SQLiteDatabase;
+import android.os.Process;
 import android.os.UserHandle;
-import android.util.ArrayMap;
-import android.util.LongSparseArray;
-
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
+import android.os.UserManager;
 
 import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.provider.RestoreDbTask;
+import com.android.launcher3.shadows.LShadowBackupManager;
 import com.android.launcher3.util.LauncherModelHelper;
-import com.android.launcher3.util.SafeCloseable;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.shadow.api.Shadow;
+import org.robolectric.shadows.ShadowUserManager;
 
 /**
  * Tests to verify backup and restore flow.
  */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
+@LooperMode(LooperMode.Mode.PAUSED)
 public class BackupRestoreTest {
 
-    private static final int PER_USER_RANGE = 200000;
-
-
-    private long mCurrentMyProfileId;
-    private long mOldMyProfileId;
-
-    private long mCurrentWorkProfileId;
-    private long mOldWorkProfileId;
+    private static final long MY_OLD_PROFILE_ID = 1;
+    private static final long MY_PROFILE_ID = 0;
+    private static final long OLD_WORK_PROFILE_ID = 11;
+    private static final int WORK_PROFILE_ID = 10;
 
+    private ShadowUserManager mUserManager;
     private BackupManager mBackupManager;
     private LauncherModelHelper mModelHelper;
     private SQLiteDatabase mDb;
     private InvariantDeviceProfile mIdp;
 
-    private UserHandle mWorkUserHandle;
-
-    private SafeCloseable mUserChangeListener;
-
     @Before
     public void setUp() {
-        mModelHelper = new LauncherModelHelper();
-
-        mCurrentMyProfileId = mModelHelper.defaultProfileId;
-        mOldMyProfileId = mCurrentMyProfileId + 1;
-        mCurrentWorkProfileId = mOldMyProfileId + 1;
-        mOldWorkProfileId = mCurrentWorkProfileId + 1;
-
-        mWorkUserHandle = UserHandle.getUserHandleForUid(PER_USER_RANGE);
-        mUserChangeListener = UserCache.INSTANCE.get(mModelHelper.sandboxContext)
-                .addUserChangeListener(() -> { });
-
         setupUserManager();
         setupBackupManager();
-        RestoreDbTask.setPending(mModelHelper.sandboxContext);
+        mModelHelper = new LauncherModelHelper();
+        RestoreDbTask.setPending(RuntimeEnvironment.application);
         mDb = mModelHelper.provider.getDb();
-        mIdp = InvariantDeviceProfile.INSTANCE.get(mModelHelper.sandboxContext);
-
-    }
-
-    @After
-    public void tearDown() {
-        mUserChangeListener.close();
-        mModelHelper.destroy();
+        mIdp = InvariantDeviceProfile.INSTANCE.get(RuntimeEnvironment.application);
     }
 
     private void setupUserManager() {
-        UserCache cache = UserCache.INSTANCE.get(mModelHelper.sandboxContext);
-        synchronized (cache) {
-            LongSparseArray<UserHandle> users = getField(cache, "mUsers");
-            users.clear();
-            users.put(mCurrentMyProfileId, myUserHandle());
-            users.put(mCurrentWorkProfileId, mWorkUserHandle);
-
-            ArrayMap<UserHandle, Long> userMap = getField(cache, "mUserToSerialMap");
-            userMap.clear();
-            userMap.put(myUserHandle(), mCurrentMyProfileId);
-            userMap.put(mWorkUserHandle, mCurrentWorkProfileId);
-        }
+        final UserManager userManager = RuntimeEnvironment.application.getSystemService(
+                UserManager.class);
+        mUserManager = Shadow.extract(userManager);
+        // sign in to work profile
+        mUserManager.addUser(WORK_PROFILE_ID, "work", ShadowUserManager.FLAG_MANAGED_PROFILE);
     }
 
     private void setupBackupManager() {
-        mBackupManager = spy(new BackupManager(mModelHelper.sandboxContext));
-        doReturn(myUserHandle()).when(mBackupManager)
-                .getUserForAncestralSerialNumber(eq(mOldMyProfileId));
-        doReturn(mWorkUserHandle).when(mBackupManager)
-                .getUserForAncestralSerialNumber(eq(mOldWorkProfileId));
+        mBackupManager = new BackupManager(RuntimeEnvironment.application);
+        final LShadowBackupManager bm = Shadow.extract(mBackupManager);
+        bm.addProfile(MY_OLD_PROFILE_ID, Process.myUserHandle());
+        bm.addProfile(OLD_WORK_PROFILE_ID, UserHandle.of(WORK_PROFILE_ID));
     }
 
     @Test
@@ -155,18 +118,18 @@ public class BackupRestoreTest {
                 { SHORTCUT, SHORTCUT, NO__ICON, NO__ICON},
                 { NO__ICON, NO__ICON, SHORTCUT, SHORTCUT},
                 { APP_ICON, SHORTCUT, SHORTCUT, APP_ICON},
-            }}, 1, mOldMyProfileId);
+            }}, 1, MY_OLD_PROFILE_ID);
         // setup grid for work profile on second screen
         mModelHelper.createGrid(new int[][][]{{
                 { NO__ICON, APP_ICON, SHORTCUT, SHORTCUT},
                 { SHORTCUT, SHORTCUT, NO__ICON, NO__ICON},
                 { NO__ICON, NO__ICON, SHORTCUT, SHORTCUT},
                 { APP_ICON, SHORTCUT, SHORTCUT, NO__ICON},
-            }}, 2, mOldWorkProfileId);
+            }}, 2, OLD_WORK_PROFILE_ID);
         // simulates the creation of backup upon restore
-        new GridBackupTable(mModelHelper.sandboxContext, mDb, mIdp.numDatabaseHotseatIcons,
+        new GridBackupTable(RuntimeEnvironment.application, mDb, mIdp.numDatabaseHotseatIcons,
                 mIdp.numColumns, mIdp.numRows).doBackup(
-                mOldMyProfileId, GridBackupTable.OPTION_REQUIRES_SANITIZATION);
+                        MY_OLD_PROFILE_ID, GridBackupTable.OPTION_REQUIRES_SANITIZATION);
         // reset favorites table
         createTableUsingOldProfileId();
     }
@@ -178,28 +141,28 @@ public class BackupRestoreTest {
     private void verifyTableIsFilled(String tableName, boolean sanitized) {
         assertEquals(sanitized ? 12 : 13, getCount(mDb,
                 "SELECT * FROM " + tableName + " WHERE profileId = "
-                        + (sanitized ? mCurrentMyProfileId : mOldMyProfileId)));
+                        + (sanitized ? MY_PROFILE_ID : MY_OLD_PROFILE_ID)));
         assertEquals(10, getCount(mDb, "SELECT * FROM " + tableName + " WHERE profileId = "
-                + (sanitized ? mCurrentWorkProfileId : mOldWorkProfileId)));
+                + (sanitized ? WORK_PROFILE_ID : OLD_WORK_PROFILE_ID)));
     }
 
     private void createTableUsingOldProfileId() {
         // simulates the creation of favorites table on old device
         dropTable(mDb, TABLE_NAME);
-        addTableToDb(mDb, mOldMyProfileId, false);
+        addTableToDb(mDb, MY_OLD_PROFILE_ID, false);
     }
 
     private void createRestoreSession() throws Exception {
         final PackageInstaller.SessionParams params = new PackageInstaller.SessionParams(
                 PackageInstaller.SessionParams.MODE_FULL_INSTALL);
-        final PackageInstaller installer = mModelHelper.sandboxContext.getPackageManager()
+        final PackageInstaller installer = RuntimeEnvironment.application.getPackageManager()
                 .getPackageInstaller();
         final int sessionId = installer.createSession(params);
         final PackageInstaller.SessionInfo info = installer.getSessionInfo(sessionId);
         setField(info, "installReason", INSTALL_REASON_DEVICE_RESTORE);
         // TODO: (b/148410677) we should verify the following call instead
         //  InstallSessionHelper.INSTANCE.get(getContext()).restoreDbIfApplicable(info);
-        RestoreDbTask.restoreIfPossible(mModelHelper.sandboxContext,
+        RestoreDbTask.restoreIfPossible(RuntimeEnvironment.application,
                 mModelHelper.provider.getHelper(), mBackupManager);
     }
 
diff --git a/tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java b/robolectric_tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
similarity index 92%
rename from tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
rename to robolectric_tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
index dba0a4063f..9ac3fe761e 100644
--- a/tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
@@ -16,8 +16,6 @@ import android.os.UserHandle;
 import android.os.UserManager;
 
 import androidx.annotation.NonNull;
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
 
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.icons.BitmapInfo;
@@ -28,10 +26,13 @@ import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.util.LauncherModelHelper;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.annotation.LooperMode.Mode;
 
 import java.util.Arrays;
 import java.util.HashSet;
@@ -39,8 +40,8 @@ import java.util.HashSet;
 /**
  * Tests for {@link CacheDataUpdatedTask}
  */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
+@LooperMode(Mode.PAUSED)
 public class CacheDataUpdatedTaskTest {
 
     private static final String NEW_LABEL_PREFIX = "new-label-";
@@ -50,10 +51,10 @@ public class CacheDataUpdatedTaskTest {
     @Before
     public void setup() throws Exception {
         mModelHelper = new LauncherModelHelper();
-        mModelHelper.initializeData("cache_data_updated_task_data");
+        mModelHelper.initializeData("/cache_data_updated_task_data.txt");
 
         // Add placeholder entries in the cache to simulate update
-        Context context = mModelHelper.sandboxContext;
+        Context context = RuntimeEnvironment.application;
         IconCache iconCache = LauncherAppState.getInstance(context).getIconCache();
         CachingLogic<ItemInfo> placeholderLogic = new CachingLogic<ItemInfo>() {
             @Override
@@ -85,11 +86,6 @@ public class CacheDataUpdatedTaskTest {
         }
     }
 
-    @After
-    public void tearDown() {
-        mModelHelper.destroy();
-    }
-
     private CacheDataUpdatedTask newTask(int op, String... pkg) {
         return new CacheDataUpdatedTask(op, Process.myUserHandle(),
                 new HashSet<>(Arrays.asList(pkg)));
diff --git a/tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java b/robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
similarity index 91%
rename from tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
rename to robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
index d849c8fff9..be03c7dd4b 100644
--- a/tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
@@ -15,13 +15,12 @@
  */
 package com.android.launcher3.model;
 
-import static androidx.test.InstrumentationRegistry.getContext;
-
 import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertFalse;
 import static junit.framework.Assert.assertNotSame;
 import static junit.framework.Assert.assertTrue;
 
+import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.when;
@@ -33,10 +32,6 @@ import android.database.Cursor;
 import android.database.sqlite.SQLiteDatabase;
 import android.database.sqlite.SQLiteOpenHelper;
 
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-import androidx.test.platform.app.InstrumentationRegistry;
-
 import com.android.launcher3.LauncherProvider;
 import com.android.launcher3.LauncherProvider.DatabaseHelper;
 import com.android.launcher3.LauncherSettings.Favorites;
@@ -45,14 +40,15 @@ import com.android.launcher3.R;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
 
 import java.io.File;
 
 /**
  * Tests for {@link DbDowngradeHelper}
  */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
 public class DbDowngradeHelperTest {
 
     private static final String SCHEMA_FILE = "test_schema.json";
@@ -64,7 +60,7 @@ public class DbDowngradeHelperTest {
 
     @Before
     public void setup() {
-        mContext = InstrumentationRegistry.getInstrumentation().getTargetContext();
+        mContext = RuntimeEnvironment.application;
         mSchemaFile = mContext.getFileStreamPath(SCHEMA_FILE);
         mDbFile = mContext.getDatabasePath(DB_FILE);
     }
@@ -81,10 +77,8 @@ public class DbDowngradeHelperTest {
     public void testUpdateSchemaFile() throws Exception {
         // Setup mock resources
         Resources res = spy(mContext.getResources());
-        Resources myRes = getContext().getResources();
-        doAnswer(i -> myRes.openRawResource(
-                myRes.getIdentifier("db_schema_v10", "raw", getContext().getPackageName())))
-                .when(res).openRawResource(R.raw.downgrade_schema);
+        doAnswer(i ->this.getClass().getResourceAsStream("/db_schema_v10.json"))
+                .when(res).openRawResource(eq(R.raw.downgrade_schema));
         Context context = spy(mContext);
         when(context.getResources()).thenReturn(res);
 
diff --git a/tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java b/robolectric_tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
similarity index 77%
rename from tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
rename to robolectric_tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
index 004ed06b32..655237d08c 100644
--- a/tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
@@ -16,18 +16,18 @@
 
 package com.android.launcher3.model;
 
-import static com.android.launcher3.util.LauncherModelHelper.TEST_ACTIVITY;
 import static com.android.launcher3.util.LauncherModelHelper.TEST_PACKAGE;
 
 import static org.junit.Assert.assertEquals;
+import static org.robolectric.Shadows.shadowOf;
+import static org.robolectric.util.ReflectionHelpers.setField;
 
+import android.content.ComponentName;
 import android.content.Context;
 import android.content.pm.PackageInstaller;
+import android.content.pm.PackageInstaller.SessionInfo;
 import android.content.pm.PackageInstaller.SessionParams;
 
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.icons.BitmapInfo;
 import com.android.launcher3.model.data.FolderInfo;
@@ -35,16 +35,19 @@ import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.util.LauncherLayoutBuilder;
 import com.android.launcher3.util.LauncherModelHelper;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.annotation.LooperMode.Mode;
 
 /**
  * Tests for layout parser for remote layout
  */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
+@LooperMode(Mode.PAUSED)
 public class DefaultLayoutProviderTest {
 
     private LauncherModelHelper mModelHelper;
@@ -53,18 +56,16 @@ public class DefaultLayoutProviderTest {
     @Before
     public void setUp() {
         mModelHelper = new LauncherModelHelper();
-        mTargetContext = mModelHelper.sandboxContext;
-    }
+        mTargetContext = RuntimeEnvironment.application;
 
-    @After
-    public void tearDown() {
-        mModelHelper.destroy();
+        shadowOf(mTargetContext.getPackageManager())
+                .addActivityIfNotPresent(new ComponentName(TEST_PACKAGE, TEST_PACKAGE));
     }
 
     @Test
     public void testCustomProfileLoaded_with_icon_on_hotseat() throws Exception {
         writeLayoutAndLoad(new LauncherLayoutBuilder().atHotseat(0)
-                .putApp(TEST_PACKAGE, TEST_ACTIVITY));
+                .putApp(TEST_PACKAGE, TEST_PACKAGE));
 
         // Verify one item in hotseat
         assertEquals(1, mModelHelper.getBgDataModel().workspaceItems.size());
@@ -76,9 +77,9 @@ public class DefaultLayoutProviderTest {
     @Test
     public void testCustomProfileLoaded_with_folder() throws Exception {
         writeLayoutAndLoad(new LauncherLayoutBuilder().atHotseat(0).putFolder(android.R.string.copy)
-                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
-                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
-                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
+                .addApp(TEST_PACKAGE, TEST_PACKAGE)
+                .addApp(TEST_PACKAGE, TEST_PACKAGE)
+                .addApp(TEST_PACKAGE, TEST_PACKAGE)
                 .build());
 
         // Verify folder
@@ -91,9 +92,9 @@ public class DefaultLayoutProviderTest {
     @Test
     public void testCustomProfileLoaded_with_folder_custom_title() throws Exception {
         writeLayoutAndLoad(new LauncherLayoutBuilder().atHotseat(0).putFolder("CustomFolder")
-                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
-                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
-                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
+                .addApp(TEST_PACKAGE, TEST_PACKAGE)
+                .addApp(TEST_PACKAGE, TEST_PACKAGE)
+                .addApp(TEST_PACKAGE, TEST_PACKAGE)
                 .build());
 
         // Verify folder
@@ -111,10 +112,12 @@ public class DefaultLayoutProviderTest {
         // Add a placeholder session info so that the widget exists
         SessionParams params = new SessionParams(SessionParams.MODE_FULL_INSTALL);
         params.setAppPackageName(pendingAppPkg);
-        params.setAppIcon(BitmapInfo.LOW_RES_ICON);
 
         PackageInstaller installer = mTargetContext.getPackageManager().getPackageInstaller();
-        installer.createSession(params);
+        int sessionId = installer.createSession(params);
+        SessionInfo sessionInfo = installer.getSessionInfo(sessionId);
+        setField(sessionInfo, "installerPackageName", "com.test");
+        setField(sessionInfo, "appIcon", BitmapInfo.LOW_RES_ICON);
 
         writeLayoutAndLoad(new LauncherLayoutBuilder().atWorkspace(0, 1, 0)
                 .putWidget(pendingAppPkg, "PlaceholderWidget", 2, 2));
diff --git a/tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java b/robolectric_tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
similarity index 89%
rename from tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
rename to robolectric_tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
index 005389e5bb..87b08870c8 100644
--- a/tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
+++ b/robolectric_tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
@@ -30,31 +30,26 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
 import android.content.Context;
-import android.content.Intent;
 import android.database.Cursor;
 import android.database.sqlite.SQLiteDatabase;
 import android.graphics.Point;
 import android.os.Process;
 
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.util.LauncherModelHelper;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
 
-import java.util.HashMap;
 import java.util.HashSet;
 
 /** Unit tests for {@link GridSizeMigrationTaskV2} */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
 public class GridSizeMigrationTaskV2Test {
 
     private LauncherModelHelper mModelHelper;
@@ -78,7 +73,7 @@ public class GridSizeMigrationTaskV2Test {
     @Before
     public void setUp() {
         mModelHelper = new LauncherModelHelper();
-        mContext = mModelHelper.sandboxContext;
+        mContext = RuntimeEnvironment.application;
         mDb = mModelHelper.provider.getDb();
 
         mValidPackages = new HashSet<>();
@@ -103,13 +98,8 @@ public class GridSizeMigrationTaskV2Test {
                 LauncherSettings.Favorites.TMP_TABLE);
     }
 
-    @After
-    public void tearDown() {
-        mModelHelper.destroy();
-    }
-
     @Test
-    public void testMigration() throws Exception {
+    public void testMigration() {
         int[] srcHotseatItems = {
                 mModelHelper.addItem(APP_ICON, 0, HOTSEAT, 0, 0, testPackage1, 1, TMP_CONTENT_URI),
                 mModelHelper.addItem(SHORTCUT, 1, HOTSEAT, 0, 0, testPackage2, 2, TMP_CONTENT_URI),
@@ -144,17 +134,17 @@ public class GridSizeMigrationTaskV2Test {
         // Check hotseat items
         Cursor c = mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
                 new String[]{LauncherSettings.Favorites.SCREEN, LauncherSettings.Favorites.INTENT},
-                "container=" + CONTAINER_HOTSEAT, null, LauncherSettings.Favorites.SCREEN, null);
+                "container=" + CONTAINER_HOTSEAT, null, null, null);
         assertEquals(c.getCount(), mIdp.numDatabaseHotseatIcons);
         int screenIndex = c.getColumnIndex(LauncherSettings.Favorites.SCREEN);
         int intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
         c.moveToNext();
-        assertEquals(c.getInt(screenIndex), 0);
-        assertTrue(c.getString(intentIndex).contains(testPackage1));
-        c.moveToNext();
         assertEquals(c.getInt(screenIndex), 1);
         assertTrue(c.getString(intentIndex).contains(testPackage2));
         c.moveToNext();
+        assertEquals(c.getInt(screenIndex), 0);
+        assertTrue(c.getString(intentIndex).contains(testPackage1));
+        c.moveToNext();
         assertEquals(c.getInt(screenIndex), 2);
         assertTrue(c.getString(intentIndex).contains(testPackage3));
         c.moveToNext();
@@ -167,24 +157,35 @@ public class GridSizeMigrationTaskV2Test {
                 new String[]{LauncherSettings.Favorites.CELLX, LauncherSettings.Favorites.CELLY,
                         LauncherSettings.Favorites.INTENT},
                 "container=" + CONTAINER_DESKTOP, null, null, null);
+        assertEquals(c.getCount(), 6);
         intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
         int cellXIndex = c.getColumnIndex(LauncherSettings.Favorites.CELLX);
         int cellYIndex = c.getColumnIndex(LauncherSettings.Favorites.CELLY);
 
-        HashMap<String, Point> locMap = new HashMap<>();
-        while (c.moveToNext()) {
-            locMap.put(
-                    Intent.parseUri(c.getString(intentIndex), 0).getPackage(),
-                    new Point(c.getInt(cellXIndex), c.getInt(cellYIndex)));
-        }
-        c.close();
+        c.moveToNext();
+        assertTrue(c.getString(intentIndex).contains(testPackage7));
+        c.moveToNext();
+        assertTrue(c.getString(intentIndex).contains(testPackage6));
+        assertEquals(c.getInt(cellXIndex), 0);
+        assertEquals(c.getInt(cellYIndex), 3);
+        c.moveToNext();
+        assertTrue(c.getString(intentIndex).contains(testPackage10));
+        assertEquals(c.getInt(cellXIndex), 1);
+        assertEquals(c.getInt(cellYIndex), 3);
+        c.moveToNext();
+        assertTrue(c.getString(intentIndex).contains(testPackage5));
+        assertEquals(c.getInt(cellXIndex), 2);
+        assertEquals(c.getInt(cellYIndex), 3);
+        c.moveToNext();
+        assertTrue(c.getString(intentIndex).contains(testPackage9));
+        assertEquals(c.getInt(cellXIndex), 3);
+        assertEquals(c.getInt(cellYIndex), 3);
+        c.moveToNext();
+        assertTrue(c.getString(intentIndex).contains(testPackage8));
+        assertEquals(c.getInt(cellXIndex), 0);
+        assertEquals(c.getInt(cellYIndex), 2);
 
-        assertEquals(locMap.size(), 6);
-        assertEquals(new Point(0, 2), locMap.get(testPackage8));
-        assertEquals(new Point(0, 3), locMap.get(testPackage6));
-        assertEquals(new Point(1, 3), locMap.get(testPackage10));
-        assertEquals(new Point(2, 3), locMap.get(testPackage5));
-        assertEquals(new Point(3, 3), locMap.get(testPackage9));
+        c.close();
     }
 
     @Test
@@ -211,7 +212,7 @@ public class GridSizeMigrationTaskV2Test {
         // Check hotseat items
         Cursor c = mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
                 new String[]{LauncherSettings.Favorites.SCREEN, LauncherSettings.Favorites.INTENT},
-                "container=" + CONTAINER_HOTSEAT, null, LauncherSettings.Favorites.SCREEN, null);
+                "container=" + CONTAINER_HOTSEAT, null, null, null);
         assertEquals(c.getCount(), numSrcDatabaseHotseatIcons);
         int screenIndex = c.getColumnIndex(LauncherSettings.Favorites.SCREEN);
         int intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
@@ -256,7 +257,7 @@ public class GridSizeMigrationTaskV2Test {
         // Check hotseat items
         Cursor c = mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
                 new String[]{LauncherSettings.Favorites.SCREEN, LauncherSettings.Favorites.INTENT},
-                "container=" + CONTAINER_HOTSEAT, null, LauncherSettings.Favorites.SCREEN, null);
+                "container=" + CONTAINER_HOTSEAT, null, null, null);
         assertEquals(c.getCount(), mIdp.numDatabaseHotseatIcons);
         int screenIndex = c.getColumnIndex(LauncherSettings.Favorites.SCREEN);
         int intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
diff --git a/tests/src/com/android/launcher3/model/LoaderCursorTest.java b/robolectric_tests/src/com/android/launcher3/model/LoaderCursorTest.java
similarity index 93%
rename from tests/src/com/android/launcher3/model/LoaderCursorTest.java
rename to robolectric_tests/src/com/android/launcher3/model/LoaderCursorTest.java
index 6444ef6927..800311afb6 100644
--- a/tests/src/com/android/launcher3/model/LoaderCursorTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/LoaderCursorTest.java
@@ -16,8 +16,6 @@
 
 package com.android.launcher3.model;
 
-import static androidx.test.InstrumentationRegistry.getContext;
-
 import static com.android.launcher3.LauncherSettings.Favorites.CELLX;
 import static com.android.launcher3.LauncherSettings.Favorites.CELLY;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER;
@@ -35,7 +33,7 @@ import static com.android.launcher3.LauncherSettings.Favorites.RESTORED;
 import static com.android.launcher3.LauncherSettings.Favorites.SCREEN;
 import static com.android.launcher3.LauncherSettings.Favorites.TITLE;
 import static com.android.launcher3.LauncherSettings.Favorites._ID;
-import static com.android.launcher3.util.LauncherModelHelper.TEST_ACTIVITY;
+import static com.android.launcher3.util.LauncherModelHelper.TEST_PACKAGE;
 
 import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertFalse;
@@ -43,37 +41,37 @@ import static junit.framework.Assert.assertNotNull;
 import static junit.framework.Assert.assertNull;
 import static junit.framework.Assert.assertTrue;
 
+import static org.robolectric.Shadows.shadowOf;
+
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
 import android.database.MatrixCursor;
 import android.os.Process;
 
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.util.Executors;
-import com.android.launcher3.util.LauncherModelHelper;
 import com.android.launcher3.util.PackageManagerHelper;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.annotation.LooperMode.Mode;
 
 /**
  * Tests for {@link LoaderCursor}
  */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
+@LooperMode(Mode.PAUSED)
 public class LoaderCursorTest {
 
-    private LauncherModelHelper mModelHelper;
     private LauncherAppState mApp;
 
     private MatrixCursor mCursor;
@@ -84,8 +82,7 @@ public class LoaderCursorTest {
 
     @Before
     public void setup() {
-        mModelHelper = new LauncherModelHelper();
-        mContext = mModelHelper.sandboxContext;
+        mContext = RuntimeEnvironment.application;
         mIDP = InvariantDeviceProfile.INSTANCE.get(mContext);
         mApp = LauncherAppState.getInstance(mContext);
 
@@ -100,11 +97,6 @@ public class LoaderCursorTest {
         ums.allUsers.put(0, Process.myUserHandle());
     }
 
-    @After
-    public void tearDown() {
-        mModelHelper.destroy();
-    }
-
     private void initCursor(int itemType, String title) {
         mCursor.newRow()
                 .add(_ID, 1)
@@ -125,7 +117,9 @@ public class LoaderCursorTest {
 
     @Test
     public void getAppShortcutInfo_dontAllowMissing_validComponent() throws Exception {
-        ComponentName cn = new ComponentName(getContext(), TEST_ACTIVITY);
+        ComponentName cn = new ComponentName(TEST_PACKAGE, TEST_PACKAGE);
+        shadowOf(mContext.getPackageManager()).addActivityIfNotPresent(cn);
+
         initCursor(ITEM_TYPE_APPLICATION, "");
         assertTrue(mLoaderCursor.moveToNext());
 
diff --git a/tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java b/robolectric_tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
similarity index 65%
rename from tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
rename to robolectric_tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
index 42c9f11a52..43193553fe 100644
--- a/tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
@@ -15,60 +15,69 @@
  */
 package com.android.launcher3.model;
 
+import static com.android.launcher3.util.Executors.createAndStartNewLooper;
 import static com.android.launcher3.util.LauncherModelHelper.TEST_PACKAGE;
 
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.spy;
+import static org.robolectric.Shadows.shadowOf;
 
 import android.os.Process;
 
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-
 import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.shadows.ShadowLooperExecutor;
 import com.android.launcher3.util.Executors;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.LauncherLayoutBuilder;
 import com.android.launcher3.util.LauncherModelHelper;
+import com.android.launcher3.util.LooperExecutor;
 import com.android.launcher3.util.RunnableList;
-import com.android.launcher3.util.TestUtil;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.annotation.LooperMode.Mode;
+import org.robolectric.shadow.api.Shadow;
+import org.robolectric.shadows.ShadowPackageManager;
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashSet;
 import java.util.List;
-import java.util.Set;
 import java.util.stream.Collectors;
 
 /**
  * Tests to verify multiple callbacks in Loader
  */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
+@LooperMode(Mode.PAUSED)
 public class ModelMultiCallbacksTest {
 
     private LauncherModelHelper mModelHelper;
 
-    @Before
-    public void setUp() {
-        mModelHelper = new LauncherModelHelper();
-    }
+    private ShadowPackageManager mSpm;
+    private LooperExecutor mTempMainExecutor;
 
-    @After
-    public void tearDown() throws Exception {
-        mModelHelper.destroy();
-        TestUtil.uninstallDummyApp();
+    @Before
+    public void setUp() throws Exception {
+        mModelHelper = new LauncherModelHelper();
+        mModelHelper.installApp(TEST_PACKAGE);
+
+        mSpm = shadowOf(RuntimeEnvironment.application.getPackageManager());
+
+        // Since robolectric tests run on main thread, we run the loader-UI calls on a temp thread,
+        // so that we can wait appropriately for the loader to complete.
+        mTempMainExecutor = new LooperExecutor(createAndStartNewLooper("tempMain"));
+        ShadowLooperExecutor sle = Shadow.extract(Executors.MAIN_EXECUTOR);
+        sle.setHandler(mTempMainExecutor.getHandler());
     }
 
     @Test
@@ -76,7 +85,7 @@ public class ModelMultiCallbacksTest {
         setupWorkspacePages(3);
 
         MyCallbacks cb1 = spy(MyCallbacks.class);
-        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().addCallbacksAndLoad(cb1));
+        mModelHelper.getModel().addCallbacksAndLoad(cb1);
 
         waitForLoaderAndTempMainThread();
         cb1.verifySynchronouslyBound(3);
@@ -85,10 +94,10 @@ public class ModelMultiCallbacksTest {
         cb1.reset();
         MyCallbacks cb2 = spy(MyCallbacks.class);
         cb2.mPageToBindSync = IntSet.wrap(2);
-        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().addCallbacksAndLoad(cb2));
+        mModelHelper.getModel().addCallbacksAndLoad(cb2);
 
         waitForLoaderAndTempMainThread();
-        assertFalse(cb1.bindStarted);
+        cb1.verifySynchronouslyBound(3);
         cb2.verifySynchronouslyBound(3);
 
         // Remove callbacks
@@ -96,7 +105,7 @@ public class ModelMultiCallbacksTest {
         cb2.reset();
 
         // No effect on callbacks when removing an callback
-        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().removeCallbacks(cb2));
+        mModelHelper.getModel().removeCallbacks(cb2);
         waitForLoaderAndTempMainThread();
         assertNull(cb1.mPendingTasks);
         assertNull(cb2.mPendingTasks);
@@ -110,48 +119,52 @@ public class ModelMultiCallbacksTest {
 
     @Test
     public void testTwoCallbacks_receiveUpdates() throws Exception {
-        TestUtil.uninstallDummyApp();
-
         setupWorkspacePages(1);
 
         MyCallbacks cb1 = spy(MyCallbacks.class);
         MyCallbacks cb2 = spy(MyCallbacks.class);
-        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().addCallbacksAndLoad(cb1));
-        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().addCallbacksAndLoad(cb2));
+        mModelHelper.getModel().addCallbacksAndLoad(cb1);
+        mModelHelper.getModel().addCallbacksAndLoad(cb2);
         waitForLoaderAndTempMainThread();
 
-        assertTrue(cb1.allApps().contains(TEST_PACKAGE));
-        assertTrue(cb2.allApps().contains(TEST_PACKAGE));
+        cb1.verifyApps(TEST_PACKAGE);
+        cb2.verifyApps(TEST_PACKAGE);
 
         // Install package 1
-        TestUtil.installDummyApp();
-        mModelHelper.getModel().onPackageAdded(TestUtil.DUMMY_PACKAGE, Process.myUserHandle());
+        String pkg1 = "com.test.pkg1";
+        mModelHelper.installApp(pkg1);
+        mModelHelper.getModel().onPackageAdded(pkg1, Process.myUserHandle());
         waitForLoaderAndTempMainThread();
-        assertTrue(cb1.allApps().contains(TestUtil.DUMMY_PACKAGE));
-        assertTrue(cb2.allApps().contains(TestUtil.DUMMY_PACKAGE));
+        cb1.verifyApps(TEST_PACKAGE, pkg1);
+        cb2.verifyApps(TEST_PACKAGE, pkg1);
+
+        // Install package 2
+        String pkg2 = "com.test.pkg2";
+        mModelHelper.installApp(pkg2);
+        mModelHelper.getModel().onPackageAdded(pkg2, Process.myUserHandle());
+        waitForLoaderAndTempMainThread();
+        cb1.verifyApps(TEST_PACKAGE, pkg1, pkg2);
+        cb2.verifyApps(TEST_PACKAGE, pkg1, pkg2);
 
         // Uninstall package 2
-        TestUtil.uninstallDummyApp();
-        mModelHelper.getModel().onPackageRemoved(TestUtil.DUMMY_PACKAGE, Process.myUserHandle());
+        mSpm.removePackage(pkg1);
+        mModelHelper.getModel().onPackageRemoved(pkg1, Process.myUserHandle());
         waitForLoaderAndTempMainThread();
-        assertFalse(cb1.allApps().contains(TestUtil.DUMMY_PACKAGE));
-        assertFalse(cb2.allApps().contains(TestUtil.DUMMY_PACKAGE));
+        cb1.verifyApps(TEST_PACKAGE, pkg2);
+        cb2.verifyApps(TEST_PACKAGE, pkg2);
 
         // Unregister a callback and verify updates no longer received
-        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().removeCallbacks(cb2));
-        TestUtil.installDummyApp();
-        mModelHelper.getModel().onPackageAdded(TestUtil.DUMMY_PACKAGE, Process.myUserHandle());
+        mModelHelper.getModel().removeCallbacks(cb2);
+        mSpm.removePackage(pkg2);
+        mModelHelper.getModel().onPackageRemoved(pkg2, Process.myUserHandle());
         waitForLoaderAndTempMainThread();
-
-        // cb2 didn't get the update
-        assertTrue(cb1.allApps().contains(TestUtil.DUMMY_PACKAGE));
-        assertFalse(cb2.allApps().contains(TestUtil.DUMMY_PACKAGE));
+        cb1.verifyApps(TEST_PACKAGE);
+        cb2.verifyApps(TEST_PACKAGE, pkg2);
     }
 
     private void waitForLoaderAndTempMainThread() throws Exception {
-        Executors.MAIN_EXECUTOR.submit(() -> { }).get();
         Executors.MODEL_EXECUTOR.submit(() -> { }).get();
-        Executors.MAIN_EXECUTOR.submit(() -> { }).get();
+        mTempMainExecutor.submit(() -> { }).get();
     }
 
     private void setupWorkspacePages(int pageCount) throws Exception {
@@ -170,15 +183,9 @@ public class ModelMultiCallbacksTest {
         IntSet mPageBoundSync = new IntSet();
         RunnableList mPendingTasks;
         AppInfo[] mAppInfos;
-        boolean bindStarted;
 
         MyCallbacks() { }
 
-        @Override
-        public void startBinding() {
-            bindStarted = true;
-        }
-
         @Override
         public void onInitialBindComplete(IntSet boundPages, RunnableList pendingTasks) {
             mPageBoundSync = boundPages;
@@ -205,12 +212,10 @@ public class ModelMultiCallbacksTest {
             mPageBoundSync = new IntSet();
             mPendingTasks = null;
             mAppInfos = null;
-            bindStarted = false;
         }
 
         public void verifySynchronouslyBound(int totalItems) {
             // Verify that the requested page is bound synchronously
-            assertTrue(bindStarted);
             assertEquals(mPageToBindSync, mPageBoundSync);
             assertEquals(mItems.size(), 1);
             assertEquals(IntSet.wrap(mItems.get(0).screenId), mPageBoundSync);
@@ -221,14 +226,12 @@ public class ModelMultiCallbacksTest {
             assertEquals(mItems.size(), totalItems);
         }
 
-        public Set<String> allApps() {
-            return Arrays.stream(mAppInfos)
-                    .map(ai -> ai.getTargetComponent().getPackageName())
-                    .collect(Collectors.toSet());
-        }
-
         public void verifyApps(String... apps) {
-            assertTrue(allApps().containsAll(Arrays.asList(apps)));
+            assertEquals(apps.length, mAppInfos.length);
+            assertEquals(Arrays.stream(mAppInfos)
+                    .map(ai -> ai.getTargetComponent().getPackageName())
+                    .collect(Collectors.toSet()),
+                    new HashSet<>(Arrays.asList(apps)));
         }
     }
 }
diff --git a/tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java b/robolectric_tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
similarity index 87%
rename from tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
rename to robolectric_tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
index 519191e251..412ace03ff 100644
--- a/tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
@@ -2,19 +2,18 @@ package com.android.launcher3.model;
 
 import static org.junit.Assert.assertEquals;
 
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.pm.PackageInstallInfo;
 import com.android.launcher3.util.LauncherModelHelper;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.annotation.LooperMode.Mode;
 
 import java.util.Arrays;
 import java.util.HashSet;
@@ -22,8 +21,8 @@ import java.util.HashSet;
 /**
  * Tests for {@link PackageInstallStateChangedTask}
  */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
+@LooperMode(Mode.PAUSED)
 public class PackageInstallStateChangedTaskTest {
 
     private LauncherModelHelper mModelHelper;
@@ -31,12 +30,7 @@ public class PackageInstallStateChangedTaskTest {
     @Before
     public void setup() throws Exception {
         mModelHelper = new LauncherModelHelper();
-        mModelHelper.initializeData("package_install_state_change_task_data");
-    }
-
-    @After
-    public void tearDown() {
-        mModelHelper.destroy();
+        mModelHelper.initializeData("/package_install_state_change_task_data.txt");
     }
 
     private PackageInstallStateChangedTask newTask(String pkg, int progress) {
@@ -72,7 +66,7 @@ public class PackageInstallStateChangedTaskTest {
         HashSet<Integer> updates = new HashSet<>(Arrays.asList(idsUpdated));
         for (ItemInfo info : mModelHelper.getBgDataModel().itemsIdMap) {
             if (info instanceof WorkspaceItemInfo) {
-                assertEquals(updates.contains(info.id) ? progress: 100,
+                assertEquals(updates.contains(info.id) ? progress: 0,
                         ((WorkspaceItemInfo) info).getProgressLevel());
             } else {
                 assertEquals(updates.contains(info.id) ? progress: -1,
diff --git a/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java b/robolectric_tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
similarity index 93%
rename from tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
rename to robolectric_tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
index 48305eebdf..4184d33a68 100644
--- a/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
+++ b/robolectric_tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
@@ -21,21 +21,18 @@ import android.content.ContentValues;
 import android.database.Cursor;
 import android.database.sqlite.SQLiteDatabase;
 
-import androidx.test.InstrumentationRegistry;
-import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
-
 import com.android.launcher3.LauncherProvider.DatabaseHelper;
 import com.android.launcher3.LauncherSettings.Favorites;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
 
 /**
  * Tests for {@link RestoreDbTask}
  */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
 public class RestoreDbTaskTest {
 
     @Test
@@ -98,7 +95,7 @@ public class RestoreDbTaskTest {
         private final long mProfileId;
 
         MyDatabaseHelper(long profileId) {
-            super(InstrumentationRegistry.getInstrumentation().getTargetContext(), null, false);
+            super(RuntimeEnvironment.application, null, false);
             mProfileId = profileId;
         }
 
diff --git a/robolectric_tests/src/com/android/launcher3/secondarydisplay/SDWorkModeTest.java b/robolectric_tests/src/com/android/launcher3/secondarydisplay/SDWorkModeTest.java
new file mode 100644
index 0000000000..e3694aeddc
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/secondarydisplay/SDWorkModeTest.java
@@ -0,0 +1,125 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.secondarydisplay;
+
+import static com.android.launcher3.util.LauncherModelHelper.TEST_PACKAGE;
+import static com.android.launcher3.util.LauncherUIHelper.doLayout;
+import static com.android.launcher3.util.Preconditions.assertNotNull;
+
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertTrue;
+
+import android.content.Context;
+import android.os.UserManager;
+import android.provider.Settings;
+
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.allapps.AllAppsPagedView;
+import com.android.launcher3.allapps.AllAppsRecyclerView;
+import com.android.launcher3.util.LauncherLayoutBuilder;
+import com.android.launcher3.util.LauncherModelHelper;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.robolectric.Robolectric;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.annotation.LooperMode.Mode;
+import org.robolectric.shadow.api.Shadow;
+import org.robolectric.shadows.ShadowUserManager;
+
+/**
+ * Tests for {@link SecondaryDisplayLauncher} with work profile
+ */
+@RunWith(RobolectricTestRunner.class)
+@LooperMode(Mode.PAUSED)
+public class SDWorkModeTest {
+
+    private static final int SYSTEM_USER = 0;
+    private static final int FLAG_SYSTEM = 0x00000800;
+    private static final int WORK_PROFILE_ID = 10;
+    private static final int FLAG_PROFILE = 0x00001000;
+
+    private Context mTargetContext;
+    private InvariantDeviceProfile mIdp;
+    private LauncherModelHelper mModelHelper;
+
+    @Before
+    public void setup() throws Exception {
+        mModelHelper = new LauncherModelHelper();
+        mTargetContext = RuntimeEnvironment.application;
+        mIdp = InvariantDeviceProfile.INSTANCE.get(mTargetContext);
+        Settings.Global.putFloat(mTargetContext.getContentResolver(),
+                Settings.Global.WINDOW_ANIMATION_SCALE, 0);
+
+        mModelHelper.installApp(TEST_PACKAGE);
+    }
+
+    @Test
+    public void testAllAppsList_noWorkProfile() throws Exception {
+        SecondaryDisplayLauncher launcher = loadLauncher();
+        launcher.showAppDrawer(true);
+        doLayout(launcher);
+
+        verifyRecyclerViewCount(launcher.getAppsView().getActiveRecyclerView());
+    }
+
+    @Test
+    public void testAllAppsList_workProfile() throws Exception {
+        ShadowUserManager sum = Shadow.extract(mTargetContext.getSystemService(UserManager.class));
+        sum.addUser(SYSTEM_USER, "me", FLAG_SYSTEM);
+        sum.addProfile(SYSTEM_USER, WORK_PROFILE_ID, "work", FLAG_PROFILE);
+
+        SecondaryDisplayLauncher launcher = loadLauncher();
+        launcher.showAppDrawer(true);
+        doLayout(launcher);
+
+        AllAppsRecyclerView rv1 = launcher.getAppsView().getActiveRecyclerView();
+        verifyRecyclerViewCount(rv1);
+
+        assertNotNull(launcher.getAppsView().getWorkModeSwitch());
+        assertTrue(launcher.getAppsView().getRecyclerViewContainer() instanceof AllAppsPagedView);
+
+        AllAppsPagedView pagedView =
+                (AllAppsPagedView) launcher.getAppsView().getRecyclerViewContainer();
+        pagedView.snapToPageImmediately(1);
+        doLayout(launcher);
+
+        AllAppsRecyclerView rv2 = launcher.getAppsView().getActiveRecyclerView();
+        verifyRecyclerViewCount(rv2);
+        assertNotSame(rv1, rv2);
+    }
+
+    private SecondaryDisplayLauncher loadLauncher() throws Exception {
+        // Install 100 apps
+        for (int i = 0; i < 100; i++) {
+            mModelHelper.installApp(TEST_PACKAGE + i);
+        }
+        mModelHelper.setupDefaultLayoutProvider(new LauncherLayoutBuilder()).loadModelSync();
+        SecondaryDisplayLauncher launcher =
+                Robolectric.buildActivity(SecondaryDisplayLauncher.class).setup().get();
+        doLayout(launcher);
+        return launcher;
+    }
+
+    private void verifyRecyclerViewCount(AllAppsRecyclerView rv) {
+        int childCount = rv.getChildCount();
+        assertTrue(childCount > 0);
+        assertTrue(childCount < 100);
+    }
+}
diff --git a/robolectric_tests/src/com/android/launcher3/testing/TestActivity.java b/robolectric_tests/src/com/android/launcher3/testing/TestActivity.java
new file mode 100644
index 0000000000..17d0ac1a9f
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/testing/TestActivity.java
@@ -0,0 +1,46 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.testing;
+
+import com.android.launcher3.BaseActivity;
+import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.views.ActivityContext;
+import com.android.launcher3.views.BaseDragLayer;
+
+/** An empty activity for {@link android.app.Fragment}s, {@link android.view.View}s testing. */
+public class TestActivity extends BaseActivity implements ActivityContext {
+
+    private DeviceProfile mDeviceProfile;
+
+    @Override
+    public BaseDragLayer getDragLayer() {
+        return new BaseDragLayer(this, /* attrs= */ null, /* alphaChannelCount= */ 1) {
+            @Override
+            public void recreateControllers() {
+                // Do nothing.
+            }
+        };
+    }
+
+    @Override
+    public DeviceProfile getDeviceProfile() {
+        return mDeviceProfile;
+    }
+
+    public void setDeviceProfile(DeviceProfile deviceProfile) {
+        mDeviceProfile = deviceProfile;
+    }
+}
diff --git a/tests/src/com/android/launcher3/util/LauncherLayoutBuilder.java b/robolectric_tests/src/com/android/launcher3/util/LauncherLayoutBuilder.java
similarity index 100%
rename from tests/src/com/android/launcher3/util/LauncherLayoutBuilder.java
rename to robolectric_tests/src/com/android/launcher3/util/LauncherLayoutBuilder.java
diff --git a/tests/src/com/android/launcher3/util/LauncherModelHelper.java b/robolectric_tests/src/com/android/launcher3/util/LauncherModelHelper.java
similarity index 60%
rename from tests/src/com/android/launcher3/util/LauncherModelHelper.java
rename to robolectric_tests/src/com/android/launcher3/util/LauncherModelHelper.java
index c9b63aeb11..846e2019a0 100644
--- a/tests/src/com/android/launcher3/util/LauncherModelHelper.java
+++ b/robolectric_tests/src/com/android/launcher3/util/LauncherModelHelper.java
@@ -15,40 +15,26 @@
  */
 package com.android.launcher3.util;
 
-import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
-import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
+import static android.content.Intent.ACTION_CREATE_SHORTCUT;
 
 import static com.android.launcher3.LauncherSettings.Favorites.CONTENT_URI;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 
-import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.atLeast;
-import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
+import static org.robolectric.Shadows.shadowOf;
 
 import android.content.ComponentName;
-import android.content.ContentProvider;
-import android.content.ContentResolver;
 import android.content.ContentValues;
 import android.content.Context;
 import android.content.Intent;
-import android.content.pm.PackageManager;
-import android.content.pm.ProviderInfo;
-import android.content.res.Resources;
+import android.content.IntentFilter;
+import android.content.pm.PackageManager.NameNotFoundException;
 import android.database.sqlite.SQLiteDatabase;
 import android.net.Uri;
-import android.os.ParcelFileDescriptor;
-import android.os.ParcelFileDescriptor.AutoCloseOutputStream;
 import android.os.Process;
 import android.provider.Settings;
-import android.test.mock.MockContentResolver;
-import android.util.ArrayMap;
-
-import androidx.test.core.app.ApplicationProvider;
-import androidx.test.uiautomator.UiDevice;
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
@@ -59,30 +45,27 @@ import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.model.AllAppsList;
 import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.BgDataModel.Callbacks;
-import com.android.launcher3.model.ItemInstallQueue;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfo;
-import com.android.launcher3.pm.InstallSessionHelper;
 import com.android.launcher3.pm.UserCache;
-import com.android.launcher3.testing.TestInformationProvider;
-import com.android.launcher3.uioverrides.plugins.PluginManagerWrapper;
-import com.android.launcher3.util.MainThreadInitializedObject.ObjectProvider;
-import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext;
-import com.android.launcher3.widget.custom.CustomWidgetManager;
+import com.android.launcher3.shadows.ShadowLooperExecutor;
 
 import org.mockito.ArgumentCaptor;
+import org.robolectric.Robolectric;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.shadow.api.Shadow;
+import org.robolectric.shadows.ShadowContentResolver;
+import org.robolectric.shadows.ShadowPackageManager;
+import org.robolectric.util.ReflectionHelpers;
 
 import java.io.BufferedReader;
+import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
-import java.io.File;
-import java.io.FileNotFoundException;
 import java.io.InputStreamReader;
 import java.io.OutputStreamWriter;
 import java.lang.reflect.Field;
 import java.util.HashMap;
 import java.util.List;
-import java.util.UUID;
-import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
 import java.util.function.Function;
@@ -98,9 +81,7 @@ public class LauncherModelHelper {
     public static final int APP_ICON = LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
     public static final int SHORTCUT = LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT;
     public static final int NO__ICON = -1;
-
-    public static final String TEST_PACKAGE = testContext().getPackageName();
-    public static final String TEST_ACTIVITY = "com.android.launcher3.tests.Activity2";
+    public static final String TEST_PACKAGE = "com.android.launcher3.validpackage";
 
     // Authority for providing a test default-workspace-layout data.
     private static final String TEST_PROVIDER_AUTHORITY =
@@ -109,42 +90,21 @@ public class LauncherModelHelper {
     private static final int DEFAULT_GRID_SIZE = 4;
 
     private final HashMap<Class, HashMap<String, Field>> mFieldCache = new HashMap<>();
-    private final MockContentResolver mMockResolver = new MockContentResolver();
     public final TestLauncherProvider provider;
-    public final SanboxModelContext sandboxContext;
-
-    public final long defaultProfileId;
+    private final long mDefaultProfileId;
 
     private BgDataModel mDataModel;
     private AllAppsList mAllAppsList;
 
     public LauncherModelHelper() {
-        Context context = getApplicationContext();
-        // System settings cache content provider. Ensure that they are statically initialized
-        Settings.Secure.getString(context.getContentResolver(), "test");
-        Settings.System.getString(context.getContentResolver(), "test");
-        Settings.Global.getString(context.getContentResolver(), "test");
-
-        provider = new TestLauncherProvider();
-        sandboxContext = new SanboxModelContext();
-        defaultProfileId = UserCache.INSTANCE.get(sandboxContext)
+        provider = Robolectric.setupContentProvider(TestLauncherProvider.class);
+        mDefaultProfileId = UserCache.INSTANCE.get(RuntimeEnvironment.application)
                 .getSerialNumberForUser(Process.myUserHandle());
-        setupProvider(LauncherProvider.AUTHORITY, provider);
-    }
-
-    protected void setupProvider(String authority, ContentProvider provider) {
-        ProviderInfo providerInfo = new ProviderInfo();
-        providerInfo.authority = authority;
-        providerInfo.applicationInfo = sandboxContext.getApplicationInfo();
-        provider.attachInfo(sandboxContext, providerInfo);
-        mMockResolver.addProvider(providerInfo.authority, provider);
-        doReturn(providerInfo)
-                .when(sandboxContext.mPm)
-                .resolveContentProvider(eq(authority), anyInt());
+        ShadowContentResolver.registerProviderInternal(LauncherProvider.AUTHORITY, provider);
     }
 
     public LauncherModel getModel() {
-        return LauncherAppState.getInstance(sandboxContext).getModel();
+        return LauncherAppState.getInstance(RuntimeEnvironment.application).getModel();
     }
 
     public synchronized BgDataModel getBgDataModel() {
@@ -161,28 +121,6 @@ public class LauncherModelHelper {
         return mAllAppsList;
     }
 
-    public void destroy() {
-        // When destroying the context, make sure that the model thread is blocked, so that no
-        // new jobs get posted while we are cleaning up
-        CountDownLatch l1 = new CountDownLatch(1);
-        CountDownLatch l2 = new CountDownLatch(1);
-        MODEL_EXECUTOR.execute(() -> {
-            l1.countDown();
-            waitOrThrow(l2);
-        });
-        waitOrThrow(l1);
-        sandboxContext.onDestroy();
-        l2.countDown();
-    }
-
-    private void waitOrThrow(CountDownLatch latch) {
-        try {
-            latch.await();
-        } catch (Exception e) {
-            throw new RuntimeException(e);
-        }
-    }
-
     /**
      * Synchronously executes the task and returns all the UI callbacks posted.
      */
@@ -223,16 +161,13 @@ public class LauncherModelHelper {
      * Initializes mock data for the test.
      */
     public void initializeData(String resourceName) throws Exception {
+        Context targetContext = RuntimeEnvironment.application;
         BgDataModel bgDataModel = getBgDataModel();
         AllAppsList allAppsList = getAllAppsList();
 
         MODEL_EXECUTOR.submit(() -> {
-            // Copy apk from resources to a local file and install from there.
-            Resources resources = testContext().getResources();
-            int resId = resources.getIdentifier(
-                    resourceName, "raw", testContext().getPackageName());
             try (BufferedReader reader = new BufferedReader(new InputStreamReader(
-                    resources.openRawResource(resId)))) {
+                    this.getClass().getResourceAsStream(resourceName)))) {
                 String line;
                 HashMap<String, Class> classMap = new HashMap<>();
                 while ((line = reader.readLine()) != null) {
@@ -246,7 +181,7 @@ public class LauncherModelHelper {
                             classMap.put(commands[1], Class.forName(commands[2]));
                             break;
                         case "bgItem":
-                            bgDataModel.addItem(sandboxContext,
+                            bgDataModel.addItem(targetContext,
                                     (ItemInfo) initItem(classMap.get(commands[1]), commands, 2),
                                     false);
                             break;
@@ -301,7 +236,7 @@ public class LauncherModelHelper {
     }
 
     public int addItem(int type, int screen, int container, int x, int y) {
-        return addItem(type, screen, container, x, y, defaultProfileId, TEST_PACKAGE);
+        return addItem(type, screen, container, x, y, mDefaultProfileId, TEST_PACKAGE);
     }
 
     public int addItem(int type, int screen, int container, int x, int y, long profileId) {
@@ -309,12 +244,12 @@ public class LauncherModelHelper {
     }
 
     public int addItem(int type, int screen, int container, int x, int y, String packageName) {
-        return addItem(type, screen, container, x, y, defaultProfileId, packageName);
+        return addItem(type, screen, container, x, y, mDefaultProfileId, packageName);
     }
 
     public int addItem(int type, int screen, int container, int x, int y, String packageName,
             int id, Uri contentUri) {
-        addItem(type, screen, container, x, y, defaultProfileId, packageName, id, contentUri);
+        addItem(type, screen, container, x, y, mDefaultProfileId, packageName, id, contentUri);
         return id;
     }
 
@@ -325,7 +260,8 @@ public class LauncherModelHelper {
      */
     public int addItem(int type, int screen, int container, int x, int y, long profileId,
             String packageName) {
-        int id = LauncherSettings.Settings.call(sandboxContext.getContentResolver(),
+        Context context = RuntimeEnvironment.application;
+        int id = LauncherSettings.Settings.call(context.getContentResolver(),
                 LauncherSettings.Settings.METHOD_NEW_ITEM_ID)
                 .getInt(LauncherSettings.Settings.EXTRA_VALUE);
         addItem(type, screen, container, x, y, profileId, packageName, id, CONTENT_URI);
@@ -334,6 +270,8 @@ public class LauncherModelHelper {
 
     public void addItem(int type, int screen, int container, int x, int y, long profileId,
             String packageName, int id, Uri contentUri) {
+        Context context = RuntimeEnvironment.application;
+
         ContentValues values = new ContentValues();
         values.put(LauncherSettings.Favorites._ID, id);
         values.put(LauncherSettings.Favorites.CONTAINER, container);
@@ -357,7 +295,7 @@ public class LauncherModelHelper {
             }
         }
 
-        sandboxContext.getContentResolver().insert(contentUri, values);
+        context.getContentResolver().insert(contentUri, values);
     }
 
     public int[][][] createGrid(int[][][] typeArray) {
@@ -365,11 +303,12 @@ public class LauncherModelHelper {
     }
 
     public int[][][] createGrid(int[][][] typeArray, int startScreen) {
-        LauncherSettings.Settings.call(sandboxContext.getContentResolver(),
+        final Context context = RuntimeEnvironment.application;
+        LauncherSettings.Settings.call(context.getContentResolver(),
                 LauncherSettings.Settings.METHOD_CREATE_EMPTY_DB);
-        LauncherSettings.Settings.call(sandboxContext.getContentResolver(),
+        LauncherSettings.Settings.call(context.getContentResolver(),
                 LauncherSettings.Settings.METHOD_CLEAR_EMPTY_DB_FLAG);
-        return createGrid(typeArray, startScreen, defaultProfileId);
+        return createGrid(typeArray, startScreen, mDefaultProfileId);
     }
 
     /**
@@ -381,13 +320,14 @@ public class LauncherModelHelper {
      * @return the same grid representation where each entry is the corresponding item id.
      */
     public int[][][] createGrid(int[][][] typeArray, int startScreen, long profileId) {
+        Context context = RuntimeEnvironment.application;
         int[][][] ids = new int[typeArray.length][][];
         for (int i = 0; i < typeArray.length; i++) {
             // Add screen to DB
             int screenId = startScreen + i;
 
             // Keep the screen id counter up to date
-            LauncherSettings.Settings.call(sandboxContext.getContentResolver(),
+            LauncherSettings.Settings.call(context.getContentResolver(),
                     LauncherSettings.Settings.METHOD_NEW_SCREEN_ID);
 
             ids[i] = new int[typeArray[i].length][];
@@ -413,45 +353,69 @@ public class LauncherModelHelper {
      */
     public LauncherModelHelper setupDefaultLayoutProvider(LauncherLayoutBuilder builder)
             throws Exception {
-        InvariantDeviceProfile idp = InvariantDeviceProfile.INSTANCE.get(sandboxContext);
+        Context context = RuntimeEnvironment.application;
+        InvariantDeviceProfile idp = InvariantDeviceProfile.INSTANCE.get(context);
         idp.numRows = idp.numColumns = idp.numDatabaseHotseatIcons = DEFAULT_GRID_SIZE;
         idp.iconBitmapSize = DEFAULT_BITMAP_SIZE;
 
-        UiDevice.getInstance(getInstrumentation()).executeShellCommand(
-                "settings put secure launcher3.layout.provider " + TEST_PROVIDER_AUTHORITY);
-        ContentProvider cp = new TestInformationProvider() {
+        Settings.Secure.putString(context.getContentResolver(),
+                "launcher3.layout.provider", TEST_PROVIDER_AUTHORITY);
 
-            @Override
-            public ParcelFileDescriptor openFile(Uri uri, String mode)
-                    throws FileNotFoundException {
-                try {
-                    ParcelFileDescriptor[] pipe = ParcelFileDescriptor.createPipe();
-                    AutoCloseOutputStream outputStream = new AutoCloseOutputStream(pipe[1]);
-                    ByteArrayOutputStream bos = new ByteArrayOutputStream();
-                    builder.build(new OutputStreamWriter(bos));
-                    outputStream.write(bos.toByteArray());
-                    outputStream.flush();
-                    outputStream.close();
-                    return pipe[0];
-                } catch (Exception e) {
-                    throw new FileNotFoundException(e.getMessage());
-                }
-            }
-        };
-        setupProvider(TEST_PROVIDER_AUTHORITY, cp);
+        shadowOf(context.getPackageManager())
+                .addProviderIfNotPresent(new ComponentName("com.test", "Mock")).authority =
+                TEST_PROVIDER_AUTHORITY;
+
+        ByteArrayOutputStream bos = new ByteArrayOutputStream();
+        builder.build(new OutputStreamWriter(bos));
+        Uri layoutUri = LauncherProvider.getLayoutUri(TEST_PROVIDER_AUTHORITY, context);
+        shadowOf(context.getContentResolver()).registerInputStream(layoutUri,
+                new ByteArrayInputStream(bos.toByteArray()));
         return this;
     }
 
+    /**
+     * Simulates an apk install with a default main activity with same class and package name
+     */
+    public void installApp(String component) throws NameNotFoundException {
+        IntentFilter filter = new IntentFilter(Intent.ACTION_MAIN);
+        filter.addCategory(Intent.CATEGORY_LAUNCHER);
+        installApp(component, component, filter);
+    }
+
+    /**
+     * Simulates a custom shortcut install
+     */
+    public void installCustomShortcut(String pkg, String clazz) throws NameNotFoundException {
+        installApp(pkg, clazz, new IntentFilter(ACTION_CREATE_SHORTCUT));
+    }
+
+    private void installApp(String pkg, String clazz, IntentFilter filter)
+            throws NameNotFoundException {
+        ShadowPackageManager spm = shadowOf(RuntimeEnvironment.application.getPackageManager());
+        ComponentName cn = new ComponentName(pkg, clazz);
+        spm.addActivityIfNotPresent(cn);
+
+        filter.addCategory(Intent.CATEGORY_DEFAULT);
+        spm.addIntentFilterForActivity(cn, filter);
+    }
+
     /**
      * Loads the model in memory synchronously
      */
     public void loadModelSync() throws ExecutionException, InterruptedException {
-        Callbacks mockCb = new Callbacks() { };
-        Executors.MAIN_EXECUTOR.submit(() -> getModel().addCallbacksAndLoad(mockCb)).get();
+        // Since robolectric tests run on main thread, we run the loader-UI calls on a temp thread,
+        // so that we can wait appropriately for the loader to complete.
+        ShadowLooperExecutor sle = Shadow.extract(Executors.MAIN_EXECUTOR);
+        sle.setHandler(Executors.UI_HELPER_EXECUTOR.getHandler());
+
+        Callbacks mockCb = mock(Callbacks.class);
+        getModel().addCallbacksAndLoad(mockCb);
 
         Executors.MODEL_EXECUTOR.submit(() -> { }).get();
-        Executors.MAIN_EXECUTOR.submit(() -> { }).get();
-        Executors.MAIN_EXECUTOR.submit(() -> getModel().removeCallbacks(mockCb)).get();
+        Executors.UI_HELPER_EXECUTOR.submit(() -> { }).get();
+
+        sle.setHandler(null);
+        getModel().removeCallbacks(mockCb);
     }
 
     /**
@@ -473,97 +437,4 @@ public class LauncherModelHelper {
             return mOpenHelper;
         }
     }
-
-    public static boolean deleteContents(File dir) {
-        File[] files = dir.listFiles();
-        boolean success = true;
-        if (files != null) {
-            for (File file : files) {
-                if (file.isDirectory()) {
-                    success &= deleteContents(file);
-                }
-                if (!file.delete()) {
-                    success = false;
-                }
-            }
-        }
-        return success;
-    }
-
-    public class SanboxModelContext extends SandboxContext {
-
-        private final ArrayMap<String, Object> mSpiedServices = new ArrayMap<>();
-        private final PackageManager mPm;
-        private final File mDbDir;
-
-        SanboxModelContext() {
-            super(ApplicationProvider.getApplicationContext(),
-                    UserCache.INSTANCE, InstallSessionHelper.INSTANCE,
-                    LauncherAppState.INSTANCE, InvariantDeviceProfile.INSTANCE,
-                    DisplayController.INSTANCE, CustomWidgetManager.INSTANCE,
-                    SettingsCache.INSTANCE, PluginManagerWrapper.INSTANCE,
-                    ItemInstallQueue.INSTANCE);
-            mPm = spy(getBaseContext().getPackageManager());
-            mDbDir = new File(getCacheDir(), UUID.randomUUID().toString());
-        }
-
-        public SanboxModelContext allow(MainThreadInitializedObject object) {
-            mAllowedObjects.add(object);
-            return this;
-        }
-
-        @Override
-        public File getDatabasePath(String name) {
-            if (!mDbDir.exists()) {
-                mDbDir.mkdirs();
-            }
-            return new File(mDbDir, name);
-        }
-
-        @Override
-        public ContentResolver getContentResolver() {
-            return mMockResolver;
-        }
-
-        @Override
-        public void onDestroy() {
-            if (deleteContents(mDbDir)) {
-                mDbDir.delete();
-            }
-            super.onDestroy();
-        }
-
-
-        @Override
-        protected <T> T createObject(ObjectProvider<T> provider) {
-            return spy(provider.get(this));
-        }
-
-        @Override
-        public PackageManager getPackageManager() {
-            return mPm;
-        }
-
-        @Override
-        public Object getSystemService(String name) {
-            Object service = mSpiedServices.get(name);
-            return service != null ? service : super.getSystemService(name);
-        }
-
-        public <T> T spyService(Class<T> tClass) {
-            String name = getSystemServiceName(tClass);
-            Object service = mSpiedServices.get(name);
-            if (service != null) {
-                return (T) service;
-            }
-
-            T result = spy(getSystemService(tClass));
-            mSpiedServices.put(name, result);
-            return result;
-        }
-    }
-
-    private static Context testContext() {
-        return getInstrumentation().getContext();
-    }
 }
diff --git a/robolectric_tests/src/com/android/launcher3/util/LauncherUIHelper.java b/robolectric_tests/src/com/android/launcher3/util/LauncherUIHelper.java
new file mode 100644
index 0000000000..caad40e260
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/util/LauncherUIHelper.java
@@ -0,0 +1,101 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.util;
+
+import static android.view.View.MeasureSpec.EXACTLY;
+import static android.view.View.MeasureSpec.makeMeasureSpec;
+
+import static com.android.launcher3.Utilities.createHomeIntent;
+
+import android.app.Activity;
+import android.content.Context;
+import android.content.Intent;
+import android.content.pm.ResolveInfo;
+import android.graphics.Point;
+import android.view.View;
+import android.view.WindowManager;
+
+import com.android.launcher3.Launcher;
+
+import org.robolectric.Robolectric;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.android.controller.ActivityController;
+import org.robolectric.shadows.ShadowLooper;
+import org.robolectric.util.ReflectionHelpers;
+
+import java.util.List;
+
+/**
+ * Utility class to help manage Launcher UI and related objects for test.
+ */
+public class LauncherUIHelper {
+
+    /**
+     * Returns the class name for the Launcher activity as defined in the manifest
+     */
+    public static String getLauncherClassName() {
+        Context context = RuntimeEnvironment.application;
+        Intent homeIntent = createHomeIntent().setPackage(context.getPackageName());
+
+        List<ResolveInfo> launchers = context.getPackageManager()
+                .queryIntentActivities(homeIntent, 0);
+        if (launchers.size() != 1) {
+            return null;
+        }
+        return launchers.get(0).activityInfo.name;
+    }
+
+    /**
+     * Returns an activity controller for Launcher activity defined in the manifest
+     */
+    public static <T extends Launcher> ActivityController<T> buildLauncher() {
+        try {
+            Class<T> tClass = (Class<T>) Class.forName(getLauncherClassName());
+            return Robolectric.buildActivity(tClass);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    /**
+     * Creates and binds a Launcher activity defined in the manifest.
+     * Note that the model must be bound before calling this
+     */
+    public static <T extends Launcher> T buildAndBindLauncher() {
+        ActivityController<T> controller = buildLauncher();
+
+        T launcher = controller.setup().get();
+        doLayout(launcher);
+        ViewOnDrawExecutor executor = ReflectionHelpers.getField(launcher, "mPendingExecutor");
+        if (executor != null) {
+            executor.markCompleted();
+        }
+        return launcher;
+    }
+
+    /**
+     * Performs a measure and layout pass for the given activity
+     */
+    public static void doLayout(Activity activity) {
+        Point size = new Point();
+        RuntimeEnvironment.application.getSystemService(WindowManager.class)
+                .getDefaultDisplay().getSize(size);
+        View view = activity.getWindow().getDecorView();
+        view.measure(makeMeasureSpec(size.x, EXACTLY), makeMeasureSpec(size.y, EXACTLY));
+        view.layout(0, 0, size.x, size.y);
+        ShadowLooper.idleMainLooper();
+    }
+}
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index 50808244ad..702b73afba 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -50,7 +50,7 @@ import com.android.launcher3.util.SimpleBroadcastReceiver;
 import com.android.launcher3.util.Themes;
 import com.android.launcher3.widget.custom.CustomWidgetManager;
 
-public class LauncherAppState implements SafeCloseable {
+public class LauncherAppState {
 
     public static final String ACTION_FORCE_ROLOAD = "force-reload-launcher";
     private static final String KEY_ICON_STATE = "pref_icon_shape_path";
@@ -158,8 +158,7 @@ public class LauncherAppState implements SafeCloseable {
     /**
      * Call from Application.onTerminate(), which is not guaranteed to ever be called.
      */
-    @Override
-    public void close() {
+    public void onTerminate() {
         mModel.destroy();
         mContext.getSystemService(LauncherApps.class).unregisterCallback(mModel);
         CustomWidgetManager.INSTANCE.get(mContext).setWidgetRefreshCallback(null);
diff --git a/src/com/android/launcher3/LauncherModel.java b/src/com/android/launcher3/LauncherModel.java
index 9ebec0a2ea..7b6a5bfeb8 100644
--- a/src/com/android/launcher3/LauncherModel.java
+++ b/src/com/android/launcher3/LauncherModel.java
@@ -96,10 +96,9 @@ public class LauncherModel extends LauncherApps.Callback implements InstallSessi
     // our monitoring of the package manager provides all updates and we never
     // need to do a requery. This is only ever touched from the loader thread.
     private boolean mModelLoaded;
-    private boolean mModelDestroyed = false;
     public boolean isModelLoaded() {
         synchronized (mLock) {
-            return mModelLoaded && mLoaderTask == null && !mModelDestroyed;
+            return mModelLoaded && mLoaderTask == null;
         }
     }
 
@@ -246,7 +245,6 @@ public class LauncherModel extends LauncherApps.Callback implements InstallSessi
      * Called when the model is destroyed
      */
     public void destroy() {
-        mModelDestroyed = true;
         MODEL_EXECUTOR.execute(mModelDelegate::destroy);
     }
 
@@ -559,9 +557,6 @@ public class LauncherModel extends LauncherApps.Callback implements InstallSessi
     }
 
     public void enqueueModelUpdateTask(ModelUpdateTask task) {
-        if (mModelDestroyed) {
-            return;
-        }
         task.init(mApp, this, mBgDataModel, mBgAllAppsList, MAIN_EXECUTOR);
         MODEL_EXECUTOR.execute(task);
     }
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index 1c8954defa..a96de31241 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -86,7 +86,7 @@ import com.android.launcher3.uioverrides.plugins.PluginManagerWrapper;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSet;
-import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext;
+import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.BaseDragLayer;
 import com.android.launcher3.widget.BaseLauncherAppWidgetHostView;
@@ -97,10 +97,13 @@ import com.android.launcher3.widget.NavigableAppWidgetHostView;
 import com.android.launcher3.widget.custom.CustomWidgetManager;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.ConcurrentLinkedQueue;
 
 /**
@@ -119,16 +122,22 @@ public class LauncherPreviewRenderer extends ContextWrapper
      * Context used just for preview. It also provides a few objects (e.g. UserCache) just for
      * preview purposes.
      */
-    public static class PreviewContext extends SandboxContext {
+    public static class PreviewContext extends ContextWrapper {
+
+        private final Set<MainThreadInitializedObject> mAllowedObjects = new HashSet<>(
+                Arrays.asList(UserCache.INSTANCE, InstallSessionHelper.INSTANCE,
+                        LauncherAppState.INSTANCE, InvariantDeviceProfile.INSTANCE,
+                        CustomWidgetManager.INSTANCE, PluginManagerWrapper.INSTANCE));
 
         private final InvariantDeviceProfile mIdp;
+        private final Map<MainThreadInitializedObject, Object> mObjectMap = new HashMap<>();
         private final ConcurrentLinkedQueue<LauncherIconsForPreview> mIconPool =
                 new ConcurrentLinkedQueue<>();
 
+        private boolean mDestroyed = false;
+
         public PreviewContext(Context base, InvariantDeviceProfile idp) {
-            super(base, UserCache.INSTANCE, InstallSessionHelper.INSTANCE,
-                    LauncherAppState.INSTANCE, InvariantDeviceProfile.INSTANCE,
-                    CustomWidgetManager.INSTANCE, PluginManagerWrapper.INSTANCE);
+            super(base);
             mIdp = idp;
             mObjectMap.put(InvariantDeviceProfile.INSTANCE, idp);
             mObjectMap.put(LauncherAppState.INSTANCE,
@@ -136,6 +145,37 @@ public class LauncherPreviewRenderer extends ContextWrapper
 
         }
 
+        @Override
+        public Context getApplicationContext() {
+            return this;
+        }
+
+        public void onDestroy() {
+            CustomWidgetManager.INSTANCE.get(this).onDestroy();
+            LauncherAppState.INSTANCE.get(this).onTerminate();
+            mDestroyed = true;
+        }
+
+        /**
+         * Find a cached object from mObjectMap if we have already created one. If not, generate
+         * an object using the provider.
+         */
+        public <T> T getObject(MainThreadInitializedObject<T> mainThreadInitializedObject,
+                MainThreadInitializedObject.ObjectProvider<T> provider) {
+            if (FeatureFlags.IS_STUDIO_BUILD && mDestroyed) {
+                throw new RuntimeException("Context already destroyed");
+            }
+            if (!mAllowedObjects.contains(mainThreadInitializedObject)) {
+                throw new IllegalStateException("Leaking unknown objects");
+            }
+            if (mObjectMap.containsKey(mainThreadInitializedObject)) {
+                return (T) mObjectMap.get(mainThreadInitializedObject);
+            }
+            T t = provider.get(this);
+            mObjectMap.put(mainThreadInitializedObject, t);
+            return t;
+        }
+
         public LauncherIcons newLauncherIcons(Context context, boolean shapeDetection) {
             LauncherIconsForPreview launcherIconsForPreview = mIconPool.poll();
             if (launcherIconsForPreview != null) {
diff --git a/src/com/android/launcher3/icons/IconCache.java b/src/com/android/launcher3/icons/IconCache.java
index 1a468aeb88..cd13cd0ec5 100644
--- a/src/com/android/launcher3/icons/IconCache.java
+++ b/src/com/android/launcher3/icons/IconCache.java
@@ -134,9 +134,6 @@ public class IconCache extends BaseIconCache {
      * Closes the cache DB. This will clear any in-memory cache.
      */
     public void close() {
-        // This will clear all pending updates
-        getUpdateHandler();
-
         mIconDb.close();
     }
 
diff --git a/src/com/android/launcher3/util/DisplayController.java b/src/com/android/launcher3/util/DisplayController.java
index 94f29dbbf7..e2c0a32bb0 100644
--- a/src/com/android/launcher3/util/DisplayController.java
+++ b/src/com/android/launcher3/util/DisplayController.java
@@ -55,7 +55,7 @@ import java.util.Set;
  * Utility class to cache properties of default display to avoid a system RPC on every call.
  */
 @SuppressLint("NewApi")
-public class DisplayController implements DisplayListener, ComponentCallbacks, SafeCloseable {
+public class DisplayController implements DisplayListener, ComponentCallbacks {
 
     private static final String TAG = "DisplayController";
 
@@ -79,7 +79,6 @@ public class DisplayController implements DisplayListener, ComponentCallbacks, S
 
     private final ArrayList<DisplayInfoChangeListener> mListeners = new ArrayList<>();
     private Info mInfo;
-    private boolean mDestroyed = false;
 
     private DisplayController(Context context) {
         mContext = context;
@@ -111,17 +110,6 @@ public class DisplayController implements DisplayListener, ComponentCallbacks, S
         mDM.registerDisplayListener(this, UI_HELPER_EXECUTOR.getHandler());
     }
 
-    @Override
-    public void close() {
-        mDestroyed = true;
-        if (mWindowContext != null) {
-            mWindowContext.unregisterComponentCallbacks(this);
-        } else {
-            // TODO: unregister broadcast receiver
-        }
-        mDM.unregisterDisplayListener(this);
-    }
-
     @Override
     public final void onDisplayAdded(int displayId) { }
 
@@ -169,9 +157,6 @@ public class DisplayController implements DisplayListener, ComponentCallbacks, S
      * Only used for pre-S
      */
     private void onConfigChanged(Intent intent) {
-        if (mDestroyed) {
-            return;
-        }
         Configuration config = mContext.getResources().getConfiguration();
         if (mInfo.fontScale != config.fontScale || mInfo.densityDpi != config.densityDpi) {
             Log.d(TAG, "Configuration changed, notifying listeners");
diff --git a/src/com/android/launcher3/util/MainThreadInitializedObject.java b/src/com/android/launcher3/util/MainThreadInitializedObject.java
index ef160b199a..f6003dd7bf 100644
--- a/src/com/android/launcher3/util/MainThreadInitializedObject.java
+++ b/src/com/android/launcher3/util/MainThreadInitializedObject.java
@@ -18,20 +18,13 @@ package com.android.launcher3.util;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
 import android.content.Context;
-import android.content.ContextWrapper;
 import android.os.Looper;
 
-import androidx.annotation.UiThread;
 import androidx.annotation.VisibleForTesting;
 
+import com.android.launcher3.graphics.LauncherPreviewRenderer.PreviewContext;
 import com.android.launcher3.util.ResourceBasedOverride.Overrides;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.ExecutionException;
 
 /**
@@ -47,8 +40,8 @@ public class MainThreadInitializedObject<T> {
     }
 
     public T get(Context context) {
-        if (context instanceof SandboxContext) {
-            return ((SandboxContext) context).getObject(this, mProvider);
+        if (context instanceof PreviewContext) {
+            return ((PreviewContext) context).getObject(this, mProvider);
         }
 
         if (mValue == null) {
@@ -87,78 +80,4 @@ public class MainThreadInitializedObject<T> {
 
         T get(Context context);
     }
-
-    /**
-     * Abstract Context which allows custom implementations for
-     * {@link MainThreadInitializedObject} providers
-     */
-    public static abstract class SandboxContext extends ContextWrapper {
-
-        protected final Set<MainThreadInitializedObject> mAllowedObjects;
-        protected final Map<MainThreadInitializedObject, Object> mObjectMap = new HashMap<>();
-        protected final ArrayList<Object> mOrderedObjects = new ArrayList<>();
-
-        private final Object mDestroyLock = new Object();
-        private boolean mDestroyed = false;
-
-        public SandboxContext(Context base, MainThreadInitializedObject... allowedObjects) {
-            super(base);
-            mAllowedObjects = new HashSet<>(Arrays.asList(allowedObjects));
-        }
-
-        @Override
-        public Context getApplicationContext() {
-            return this;
-        }
-
-        public void onDestroy() {
-            synchronized (mDestroyLock) {
-                // Destroy in reverse order
-                for (int i = mOrderedObjects.size() - 1; i >= 0; i--) {
-                    Object o = mOrderedObjects.get(i);
-                    if (o instanceof SafeCloseable) {
-                        ((SafeCloseable) o).close();
-                    }
-                }
-                mDestroyed = true;
-            }
-        }
-
-        /**
-         * Find a cached object from mObjectMap if we have already created one. If not, generate
-         * an object using the provider.
-         */
-        private <T> T getObject(MainThreadInitializedObject<T> object, ObjectProvider<T> provider) {
-            synchronized (mDestroyLock) {
-                if (mDestroyed) {
-                    throw new RuntimeException("Context already destroyed");
-                }
-                if (!mAllowedObjects.contains(object)) {
-                    throw new IllegalStateException(
-                            "Leaking unknown objects " + object + "  " + provider);
-                }
-                T t = (T) mObjectMap.get(object);
-                if (t != null) {
-                    return t;
-                }
-                if (Looper.myLooper() == Looper.getMainLooper()) {
-                    t = createObject(provider);
-                    mObjectMap.put(object, t);
-                    mOrderedObjects.add(t);
-                    return t;
-                }
-            }
-
-            try {
-                return MAIN_EXECUTOR.submit(() -> getObject(object, provider)).get();
-            } catch (InterruptedException | ExecutionException e) {
-                throw new RuntimeException(e);
-            }
-        }
-
-        @UiThread
-        protected <T> T createObject(ObjectProvider<T> provider) {
-            return provider.get(this);
-        }
-    }
 }
diff --git a/src/com/android/launcher3/util/SettingsCache.java b/src/com/android/launcher3/util/SettingsCache.java
index 0c5b7225d3..10611c7211 100644
--- a/src/com/android/launcher3/util/SettingsCache.java
+++ b/src/com/android/launcher3/util/SettingsCache.java
@@ -47,7 +47,7 @@ import java.util.concurrent.CopyOnWriteArrayList;
  *
  * Cache will also be updated if a key queried is missing (even if it has no listeners registered).
  */
-public class SettingsCache extends ContentObserver implements SafeCloseable {
+public class SettingsCache extends ContentObserver {
 
     /** Hidden field Settings.Secure.NOTIFICATION_BADGING */
     public static final Uri NOTIFICATION_BADGING_URI =
@@ -69,6 +69,7 @@ public class SettingsCache extends ContentObserver implements SafeCloseable {
     private final Map<Uri, CopyOnWriteArrayList<OnChangeListener>> mListenerMap = new HashMap<>();
     protected final ContentResolver mResolver;
 
+
     /**
      * Singleton instance
      */
@@ -80,11 +81,6 @@ public class SettingsCache extends ContentObserver implements SafeCloseable {
         mResolver = context.getContentResolver();
     }
 
-    @Override
-    public void close() {
-        mResolver.unregisterContentObserver(this);
-    }
-
     @Override
     public void onChange(boolean selfChange, Uri uri) {
         // We use default of 1, but if we're getting an onChange call, can assume a non-default
diff --git a/src/com/android/launcher3/util/UiThreadHelper.java b/src/com/android/launcher3/util/UiThreadHelper.java
index ac5368c8b4..0f40179f79 100644
--- a/src/com/android/launcher3/util/UiThreadHelper.java
+++ b/src/com/android/launcher3/util/UiThreadHelper.java
@@ -28,7 +28,7 @@ import android.os.Message;
 import android.view.View;
 import android.view.inputmethod.InputMethodManager;
 
-import com.android.launcher3.BaseActivity;
+import com.android.launcher3.Launcher;
 import com.android.launcher3.views.ActivityContext;
 
 /**
@@ -56,7 +56,7 @@ public class UiThreadHelper {
                 STATS_LOGGER_KEY,
                 Message.obtain(
                         HANDLER.get(root.getContext()),
-                        () -> BaseActivity.fromContext(root.getContext())
+                        () -> Launcher.cast(activityContext)
                                 .getStatsLogManager()
                                 .logger()
                                 .log(LAUNCHER_ALLAPPS_KEYBOARD_CLOSED)
diff --git a/src/com/android/launcher3/widget/custom/CustomWidgetManager.java b/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
index 2e2a968a44..329a44452b 100644
--- a/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
+++ b/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
@@ -33,7 +33,6 @@ import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.uioverrides.plugins.PluginManagerWrapper;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.PackageUserKey;
-import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.widget.LauncherAppWidgetHostView;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.systemui.plugins.CustomWidgetPlugin;
@@ -47,7 +46,7 @@ import java.util.stream.Stream;
 /**
  * CustomWidgetManager handles custom widgets implemented as a plugin.
  */
-public class CustomWidgetManager implements PluginListener<CustomWidgetPlugin>, SafeCloseable {
+public class CustomWidgetManager implements PluginListener<CustomWidgetPlugin> {
 
     public static final MainThreadInitializedObject<CustomWidgetManager> INSTANCE =
             new MainThreadInitializedObject<>(CustomWidgetManager::new);
@@ -72,8 +71,7 @@ public class CustomWidgetManager implements PluginListener<CustomWidgetPlugin>,
                 .addPluginListener(this, CustomWidgetPlugin.class, true);
     }
 
-    @Override
-    public void close() {
+    public void onDestroy() {
         PluginManagerWrapper.INSTANCE.get(mContext).removePluginListener(this);
     }
 
diff --git a/tests/src/com/android/launcher3/secondarydisplay/SDLauncherTest.java b/tests/src/com/android/launcher3/secondarydisplay/SDLauncherTest.java
deleted file mode 100644
index fd86cf1144..0000000000
--- a/tests/src/com/android/launcher3/secondarydisplay/SDLauncherTest.java
+++ /dev/null
@@ -1,53 +0,0 @@
-/*
- * Copyright (C) 2021 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.android.launcher3.secondarydisplay;
-
-import static androidx.test.core.app.ActivityScenario.launch;
-
-import androidx.test.core.app.ActivityScenario;
-import androidx.test.espresso.intent.Intents;
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.MediumTest;
-
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-
-/**
- * Tests for {@link SecondaryDisplayLauncher}
- */
-@MediumTest
-@RunWith(AndroidJUnit4.class)
-public class SDLauncherTest {
-
-    @Before
-    public void setUp() {
-        Intents.init();
-    }
-
-    @After
-    public void tearDown() {
-        Intents.release();
-    }
-
-    @Test
-    public void testAllAppsListOpens() {
-        ActivityScenario<SecondaryDisplayLauncher> launcher =
-                launch(SecondaryDisplayLauncher.class);
-        launcher.onActivity(l -> l.showAppDrawer(true));
-    }
-}
diff --git a/tests/src/com/android/launcher3/util/ReflectionHelpers.java b/tests/src/com/android/launcher3/util/ReflectionHelpers.java
deleted file mode 100644
index d89975de8a..0000000000
--- a/tests/src/com/android/launcher3/util/ReflectionHelpers.java
+++ /dev/null
@@ -1,58 +0,0 @@
-/*
- * Copyright (C) 2021 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.launcher3.util;
-
-import java.lang.reflect.Field;
-
-public class ReflectionHelpers {
-
-    /**
-     * Reflectively get the value of a field.
-     *
-     * @param object Target object.
-     * @param fieldName The field name.
-     * @param <R> The return type.
-     * @return Value of the field on the object.
-     */
-    public static <R> R getField(Object object, String fieldName) {
-        try {
-            Field field = object.getClass().getDeclaredField(fieldName);
-            field.setAccessible(true);
-            return (R) field.get(object);
-        } catch (Exception e) {
-            throw new RuntimeException(e);
-        }
-    }
-
-    /**
-     * Reflectively set the value of a field.
-     *
-     * @param object Target object.
-     * @param fieldName The field name.
-     * @param fieldNewValue New value.
-     */
-    public static void setField(Object object, String fieldName, Object fieldNewValue) {
-        try {
-            Field field = object.getClass().getDeclaredField(fieldName);
-            field.setAccessible(true);
-            field.set(object, fieldNewValue);
-        } catch (Exception e) {
-            throw new RuntimeException(e);
-        }
-    }
-
-}