diff --git a/quickstep/src/com/android/launcher3/statehandlers/DepthController.java b/quickstep/src/com/android/launcher3/statehandlers/DepthController.java
index ab88b2b..3e2fb63 100644
--- a/quickstep/src/com/android/launcher3/statehandlers/DepthController.java
+++ b/quickstep/src/com/android/launcher3/statehandlers/DepthController.java
@@ -96,7 +96,11 @@
                 public void onDraw() {
                     View view = mLauncher.getDragLayer();
                     ViewRootImpl viewRootImpl = view.getViewRootImpl();
-                    setSurface(viewRootImpl != null ? viewRootImpl.getSurfaceControl() : null);
+                    boolean applied = setSurface(
+                            viewRootImpl != null ? viewRootImpl.getSurfaceControl() : null);
+                    if (!applied) {
+                        dispatchTransactionSurface(mDepth);
+                    }
                     view.post(() -> view.getViewTreeObserver().removeOnDrawListener(this));
                 }
             };
@@ -202,20 +206,22 @@
 
     /**
      * Sets the specified app target surface to apply the blur to.
+     * @return true when surface was valid and transaction was dispatched.
      */
-    public void setSurface(SurfaceControl surface) {
+    public boolean setSurface(SurfaceControl surface) {
         // Set launcher as the SurfaceControl when we don't need an external target anymore.
         if (surface == null) {
             ViewRootImpl viewRootImpl = mLauncher.getDragLayer().getViewRootImpl();
             surface = viewRootImpl != null ? viewRootImpl.getSurfaceControl() : null;
         }
-
         if (mSurface != surface) {
             mSurface = surface;
             if (surface != null) {
                 dispatchTransactionSurface(mDepth);
+                return true;
             }
         }
+        return false;
     }
 
     @Override
@@ -229,6 +235,8 @@
             setDepth(toDepth);
         } else if (toState == LauncherState.OVERVIEW) {
             dispatchTransactionSurface(mDepth);
+        } else if (toState == LauncherState.BACKGROUND_APP) {
+            mLauncher.getDragLayer().getViewTreeObserver().addOnDrawListener(mOnDrawListener);
         }
     }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
index 5144d9a..59393d7 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
@@ -243,7 +243,19 @@
         if (!mTouchEnabled) {
             return true;
         }
-        mControllerCallbacks.onTouchEvent(event);
+        if (mIconLayoutBounds.contains((int) event.getX(), (int) event.getY())) {
+            // Don't allow long pressing between icons.
+            return true;
+        }
+        if (mControllerCallbacks.onTouchEvent(event)) {
+            int oldAction = event.getAction();
+            try {
+                event.setAction(MotionEvent.ACTION_CANCEL);
+                return super.onTouchEvent(event);
+            } finally {
+                event.setAction(oldAction);
+            }
+        }
         return super.onTouchEvent(event);
     }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
index 40b0e18..f1748af 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
@@ -223,7 +223,11 @@
             return view -> mControllers.taskbarStashController.updateAndAnimateIsStashedInApp(true);
         }
 
-        public void onTouchEvent(MotionEvent motionEvent) {
+        /**
+         * Get the first chance to handle TaskbarView#onTouchEvent, and return whether we want to
+         * consume the touch so TaskbarView treats it as an ACTION_CANCEL.
+         */
+        public boolean onTouchEvent(MotionEvent motionEvent) {
             final float x = motionEvent.getRawX();
             final float y = motionEvent.getRawY();
             switch (motionEvent.getAction()) {
@@ -239,6 +243,7 @@
                         mControllers.taskbarStashController.startStashHint(
                                 /* animateForward= */ false);
                         mCanceledStashHint = true;
+                        return true;
                     }
                     break;
                 case MotionEvent.ACTION_UP:
@@ -249,6 +254,7 @@
                     }
                     break;
             }
+            return false;
         }
     }
 }
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 6e90a3a..4c46683 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -805,7 +805,9 @@
         mActivityInitListener.unregister();
         mStateCallback.setStateOnUiThread(STATE_GESTURE_CANCELLED | STATE_HANDLER_INVALIDATED);
 
-        TaskViewUtils.setDividerBarShown(mRecentsAnimationTargets.nonApps, true);
+        if (mRecentsAnimationTargets != null) {
+            TaskViewUtils.setDividerBarShown(mRecentsAnimationTargets.nonApps, true);
+        }
 
         // Defer clearing the controller and the targets until after we've updated the state
         mRecentsAnimationController = null;
diff --git a/quickstep/src/com/android/quickstep/TaskViewUtils.java b/quickstep/src/com/android/quickstep/TaskViewUtils.java
index 5b9e214..d4c94db 100644
--- a/quickstep/src/com/android/quickstep/TaskViewUtils.java
+++ b/quickstep/src/com/android/quickstep/TaskViewUtils.java
@@ -366,8 +366,8 @@
      * device is considered in multiWindowMode and things like insets and stuff change
      * and calculations have to be adjusted in the animations for that
      */
-    public static void composeRecentsSplitLaunchAnimator(@NonNull TaskView initialView,
-            @NonNull TaskView v, @NonNull TransitionInfo transitionInfo,
+    public static void composeRecentsSplitLaunchAnimator(@NonNull Task initalTask,
+            @NonNull Task secondTask, @NonNull TransitionInfo transitionInfo,
             SurfaceControl.Transaction t, @NonNull Runnable finishCallback) {
 
         final TransitionInfo.Change[] splitRoots = new TransitionInfo.Change[2];
@@ -377,7 +377,7 @@
             final int mode = change.getMode();
             // Find the target tasks' root tasks since those are the split stages that need to
             // be animated (the tasks themselves are children and thus inherit animation).
-            if (taskId == initialView.getTask().key.id || taskId == v.getTask().key.id) {
+            if (taskId == initalTask.key.id || taskId == secondTask.key.id) {
                 if (!(mode == TRANSIT_OPEN || mode == TRANSIT_TO_FRONT)) {
                     throw new IllegalStateException(
                             "Expected task to be showing, but it is " + mode);
@@ -386,7 +386,7 @@
                     throw new IllegalStateException("Initiating multi-split launch but the split"
                             + "root of " + taskId + " is already visible or has broken hierarchy.");
                 }
-                splitRoots[taskId == initialView.getTask().key.id ? 0 : 1] =
+                splitRoots[taskId == initalTask.key.id ? 0 : 1] =
                         transitionInfo.getChange(change.getParent());
             }
         }
@@ -406,8 +406,8 @@
     }
 
     /** Legacy version (until shell transitions are enabled) */
-    public static void composeRecentsSplitLaunchAnimatorLegacy(@NonNull TaskView initialView,
-            @NonNull TaskView v, @NonNull RemoteAnimationTargetCompat[] appTargets,
+    public static void composeRecentsSplitLaunchAnimatorLegacy(@NonNull Task initialTask,
+            @NonNull Task secondTask, @NonNull RemoteAnimationTargetCompat[] appTargets,
             @NonNull RemoteAnimationTargetCompat[] wallpaperTargets,
             @NonNull RemoteAnimationTargetCompat[] nonAppTargets,
             @NonNull Runnable finishCallback) {
@@ -416,12 +416,12 @@
         for (int i = 0; i < appTargets.length; ++i) {
             final int taskId = appTargets[i].taskInfo != null ? appTargets[i].taskInfo.taskId : -1;
             final int mode = appTargets[i].mode;
-            if (taskId == initialView.getTask().key.id || taskId == v.getTask().key.id) {
+            if (taskId == initialTask.key.id || taskId == secondTask.key.id) {
                 if (mode != MODE_OPENING) {
                     throw new IllegalStateException(
                             "Expected task to be opening, but it is " + mode);
                 }
-                splitRoots[taskId == initialView.getTask().key.id ? 0 : 1] = i;
+                splitRoots[taskId == initialTask.key.id ? 0 : 1] = i;
             }
         }
 
diff --git a/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java b/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
index cbc5584..ac5f5d8 100644
--- a/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
+++ b/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
@@ -33,7 +33,7 @@
 import com.android.quickstep.SystemUiProxy;
 import com.android.quickstep.TaskAnimationManager;
 import com.android.quickstep.TaskViewUtils;
-import com.android.quickstep.views.TaskView;
+import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.system.RemoteAnimationAdapterCompat;
 import com.android.systemui.shared.system.RemoteAnimationRunnerCompat;
 import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
@@ -47,23 +47,21 @@
 public class SplitSelectStateController {
 
     private final SystemUiProxy mSystemUiProxy;
-    private TaskView mInitialTaskView;
-    private TaskView mSecondTaskView;
     private @StagePosition int mStagePosition;
+    private Task mInitialTask;
+    private Task mSecondTask;
     private Rect mInitialBounds;
-    private final Handler mHandler;
 
     public SplitSelectStateController(Handler handler, SystemUiProxy systemUiProxy) {
         mSystemUiProxy = systemUiProxy;
-        mHandler = handler;
     }
 
     /**
      * To be called after first task selected
      */
-    public void setInitialTaskSelect(TaskView taskView, @StagePosition int stagePosition,
+    public void setInitialTaskSelect(Task taskView, @StagePosition int stagePosition,
             Rect initialBounds) {
-        mInitialTaskView = taskView;
+        mInitialTask = taskView;
         mStagePosition = stagePosition;
         mInitialBounds = initialBounds;
     }
@@ -71,21 +69,28 @@
     /**
      * To be called after second task selected
      */
-    public void setSecondTaskId(TaskView taskView) {
-        mSecondTaskView = taskView;
+    public void setSecondTaskId(Task taskView) {
+        mSecondTask = taskView;
+        launchTasks(mInitialTask, mSecondTask, mStagePosition);
+    }
+
+    /**
+     * @param stagePosition representing location of task1
+     */
+    public void launchTasks(Task task1, Task task2, @StagePosition int stagePosition) {
         // Assume initial task is for top/left part of screen
-        final int[] taskIds = mStagePosition == STAGE_POSITION_TOP_OR_LEFT
-                ? new int[]{mInitialTaskView.getTask().key.id, taskView.getTask().key.id}
-                : new int[]{taskView.getTask().key.id, mInitialTaskView.getTask().key.id};
+        final int[] taskIds = stagePosition == STAGE_POSITION_TOP_OR_LEFT
+                ? new int[]{task1.key.id, task2.key.id}
+                : new int[]{task2.key.id, task1.key.id};
         if (TaskAnimationManager.ENABLE_SHELL_TRANSITIONS) {
             RemoteSplitLaunchTransitionRunner animationRunner =
-                    new RemoteSplitLaunchTransitionRunner(mInitialTaskView, taskView);
+                    new RemoteSplitLaunchTransitionRunner(task1, task2);
             mSystemUiProxy.startTasks(taskIds[0], null /* mainOptions */, taskIds[1],
                     null /* sideOptions */, STAGE_POSITION_BOTTOM_OR_RIGHT,
                     new RemoteTransitionCompat(animationRunner, MAIN_EXECUTOR));
         } else {
             RemoteSplitLaunchAnimationRunner animationRunner =
-                    new RemoteSplitLaunchAnimationRunner(mInitialTaskView, taskView);
+                    new RemoteSplitLaunchAnimationRunner(task1, task2);
             final RemoteAnimationAdapter adapter = new RemoteAnimationAdapter(
                     RemoteAnimationAdapterCompat.wrapRemoteAnimationRunner(animationRunner),
                     300, 150,
@@ -105,19 +110,19 @@
      */
     private class RemoteSplitLaunchTransitionRunner implements RemoteTransitionRunner {
 
-        private final TaskView mInitialTaskView;
-        private final TaskView mTaskView;
+        private final Task mInitialTask;
+        private final Task mSecondTask;
 
-        RemoteSplitLaunchTransitionRunner(TaskView initialTaskView, TaskView taskView) {
-            mInitialTaskView = initialTaskView;
-            mTaskView = taskView;
+        RemoteSplitLaunchTransitionRunner(Task initialTask, Task secondTask) {
+            mInitialTask = initialTask;
+            mSecondTask = secondTask;
         }
 
         @Override
         public void startAnimation(IBinder transition, TransitionInfo info,
                 SurfaceControl.Transaction t, Runnable finishCallback) {
-            TaskViewUtils.composeRecentsSplitLaunchAnimator(mInitialTaskView, mTaskView,
-                    info, t, finishCallback);
+            TaskViewUtils.composeRecentsSplitLaunchAnimator(mInitialTask,
+                    mSecondTask, info, t, finishCallback);
             // After successful launch, call resetState
             resetState();
         }
@@ -129,20 +134,20 @@
      */
     private class RemoteSplitLaunchAnimationRunner implements RemoteAnimationRunnerCompat {
 
-        private final TaskView mInitialTaskView;
-        private final TaskView mTaskView;
+        private final Task mInitialTask;
+        private final Task mSecondTask;
 
-        RemoteSplitLaunchAnimationRunner(TaskView initialTaskView, TaskView taskView) {
-            mInitialTaskView = initialTaskView;
-            mTaskView = taskView;
+        RemoteSplitLaunchAnimationRunner(Task initialTask, Task secondTask) {
+            mInitialTask = initialTask;
+            mSecondTask = secondTask;
         }
 
         @Override
         public void onAnimationStart(int transit, RemoteAnimationTargetCompat[] apps,
                 RemoteAnimationTargetCompat[] wallpapers, RemoteAnimationTargetCompat[] nonApps,
                 Runnable finishedCallback) {
-            TaskViewUtils.composeRecentsSplitLaunchAnimatorLegacy(mInitialTaskView, mTaskView, apps,
-                    wallpapers, nonApps, finishedCallback);
+            TaskViewUtils.composeRecentsSplitLaunchAnimatorLegacy(mInitialTask,
+                    mSecondTask, apps, wallpapers, nonApps, finishedCallback);
             // After successful launch, call resetState
             resetState();
         }
@@ -157,8 +162,8 @@
      * To be called if split select was cancelled
      */
     public void resetState() {
-        mInitialTaskView = null;
-        mSecondTaskView = null;
+        mInitialTask = null;
+        mSecondTask = null;
         mStagePosition = SplitConfigurationOptions.STAGE_POSITION_UNDEFINED;
         mInitialBounds = null;
     }
@@ -168,7 +173,7 @@
      *         chosen
      */
     public boolean isSplitSelectActive() {
-        return mInitialTaskView != null && mSecondTaskView == null;
+        return mInitialTask != null && mSecondTask == null;
     }
 
     public Rect getInitialBounds() {
diff --git a/quickstep/src/com/android/quickstep/views/GroupedTaskView.java b/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
index 1bc7c75..e1e1c65 100644
--- a/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
+++ b/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
@@ -4,6 +4,7 @@
 import android.util.AttributeSet;
 
 import com.android.launcher3.R;
+import com.android.launcher3.util.RunnableList;
 import com.android.launcher3.util.SplitConfigurationOptions;
 import com.android.quickstep.RecentsModel;
 import com.android.quickstep.TaskThumbnailCache;
@@ -95,6 +96,13 @@
     }
 
     @Override
+    public RunnableList launchTaskAnimated() {
+        getRecentsView().getSplitPlaceholder().launchTasks(mTask, mSecondaryTask,
+                SplitConfigurationOptions.STAGE_POSITION_TOP_OR_LEFT);
+        return null;
+    }
+
+    @Override
     public void onRecycle() {
         super.onRecycle();
         mSnapshotView2.setThumbnail(mSecondaryTask, null);
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 18387dc..6aa7e06 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -1351,8 +1351,10 @@
                 }
                 Log.d(TASK_VIEW_ID_CRASH, "taskViewCount: " + getTaskViewCount()
                         + " " + sb.toString());
+                mRunningTaskViewId = -1;
+            } else {
+                mRunningTaskViewId = newRunningTaskView.getTaskViewId();
             }
-            mRunningTaskViewId = newRunningTaskView.getTaskViewId();
         }
 
         if (mNextPage == INVALID_PAGE) {
@@ -3568,7 +3570,8 @@
         mSplitHiddenTaskView = taskView;
         Rect initialBounds = new Rect(taskView.getLeft(), taskView.getTop(), taskView.getRight(),
                 taskView.getBottom());
-        mSplitSelectStateController.setInitialTaskSelect(taskView, stagePosition, initialBounds);
+        mSplitSelectStateController.setInitialTaskSelect(taskView.getTask(),
+                stagePosition, initialBounds);
         mSplitHiddenTaskViewIndex = indexOfChild(taskView);
         if (ENABLE_QUICKSTEP_LIVE_TILE.get()) {
             finishRecentsAnimation(true, null);
@@ -3609,7 +3612,7 @@
                 secondTaskEndingBounds, taskView.getThumbnail(),
                 true /*fadeWithThumbnail*/);
         pendingAnimation.addEndListener(aBoolean -> {
-            mSplitSelectStateController.setSecondTaskId(taskView);
+            mSplitSelectStateController.setSecondTaskId(taskView.getTask());
             resetFromSplitSelectionState();
         });
         mSecondSplitHiddenTaskView = taskView;
diff --git a/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
similarity index 65%
rename from quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
rename to quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
index f82fbcc..c1b3beb 100644
--- a/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
+++ b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
@@ -15,24 +15,31 @@
  */
 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.robolectric.Shadows.shadowOf;
+import static org.mockito.Mockito.doReturn;
 
 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;
@@ -40,38 +47,34 @@
 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;
 
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public final class WidgetsPredicationUpdateTaskTest {
 
-    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 AppWidgetProviderInfo mApp1Provider1;
+    private AppWidgetProviderInfo mApp1Provider2;
+    private AppWidgetProviderInfo mApp2Provider1;
+    private AppWidgetProviderInfo mApp4Provider1;
+    private AppWidgetProviderInfo mApp4Provider2;
+    private AppWidgetProviderInfo mApp5Provider1;
+    private List<AppWidgetProviderInfo> allWidgets;
 
     private FakeBgDataModelCallback mCallback = new FakeBgDataModelCallback();
-    private Context mContext;
     private LauncherModelHelper mModelHelper;
     private UserHandle mUserHandle;
 
@@ -80,54 +83,56 @@
 
     @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());
 
-        mContext = RuntimeEnvironment.application;
-        mModelHelper = new LauncherModelHelper();
-        mUserHandle = Process.myUserHandle();
+        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()));
+
         // 2 widgets, app4/provider1 & app5/provider1, have already been added to the workspace.
-        mModelHelper.initializeData("/widgets_predication_update_task_data.txt");
+        mModelHelper.initializeData("widgets_predication_update_task_data");
 
-        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(mContext), /* packageUser= */ null));
-        waitUntilIdle();
+                LauncherAppState.getInstance(mModelHelper.sandboxContext),
+                /* packageUser= */ null));
+
+        MODEL_EXECUTOR.submit(() -> { }).get();
+        MAIN_EXECUTOR.submit(() -> { }).get();
     }
 
+    @After
+    public void tearDown() {
+        mModelHelper.destroy();
+    }
 
     @Test
     public void widgetsRecommendationRan_shouldOnlyReturnNotAddedWidgetsInAppPredictionOrder()
@@ -165,9 +170,9 @@
     @Test
     public void widgetsRecommendationRan_localFilterDisabled_shouldReturnWidgetsInPredicationOrder()
             throws Exception {
-        ShadowDeviceFlag shadowDeviceFlag = Shadow.extract(
-                FeatureFlags.ENABLE_LOCAL_RECOMMENDED_WIDGETS_FILTER);
-        shadowDeviceFlag.setValue(false);
+        if (FeatureFlags.ENABLE_LOCAL_RECOMMENDED_WIDGETS_FILTER.get()) {
+            return;
+        }
 
         // WHEN newPredicationTask is executed with 5 predicated widgets.
         AppTarget widget1 = new AppTarget(new AppTargetId("app1"), "app1", "provider1",
@@ -203,13 +208,8 @@
         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/res/values/config.xml b/res/values/config.xml
index 04c359e..72959b2 100644
--- a/res/values/config.xml
+++ b/res/values/config.xml
@@ -34,6 +34,9 @@
     <!-- View tag key used to determine if we should fade in the child views.. -->
     <string name="popup_container_iterate_children" translatable="false">popup_container_iterate_children</string>
 
+    <!-- config used to determine if header protection is supported in AllApps -->
+    <bool name="config_header_protection_supported">false</bool>
+
     <!-- Workspace -->
     <!-- The duration (in ms) of the fade animation on the object outlines, used when
          we are dragging objects around on the home screen. -->
diff --git a/robolectric_tests/src/com/android/launcher3/secondarydisplay/SDWorkModeTest.java b/robolectric_tests/src/com/android/launcher3/secondarydisplay/SDWorkModeTest.java
deleted file mode 100644
index e3694ae..0000000
--- a/robolectric_tests/src/com/android/launcher3/secondarydisplay/SDWorkModeTest.java
+++ /dev/null
@@ -1,125 +0,0 @@
-/*
- * 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
deleted file mode 100644
index 17d0ac1..0000000
--- a/robolectric_tests/src/com/android/launcher3/testing/TestActivity.java
+++ /dev/null
@@ -1,46 +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.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/robolectric_tests/src/com/android/launcher3/util/LauncherUIHelper.java b/robolectric_tests/src/com/android/launcher3/util/LauncherUIHelper.java
deleted file mode 100644
index caad40e..0000000
--- a/robolectric_tests/src/com/android/launcher3/util/LauncherUIHelper.java
+++ /dev/null
@@ -1,101 +0,0 @@
-/*
- * 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/BubbleTextView.java b/src/com/android/launcher3/BubbleTextView.java
index 02ec5e8..e52d1be 100644
--- a/src/com/android/launcher3/BubbleTextView.java
+++ b/src/com/android/launcher3/BubbleTextView.java
@@ -409,6 +409,10 @@
      * Returns true if the touch down at the provided position be ignored
      */
     protected boolean shouldIgnoreTouchDown(float x, float y) {
+        if (mDisplay == DISPLAY_TASKBAR) {
+            // Allow touching within padding on taskbar, given icon sizes are smaller.
+            return false;
+        }
         return y < getPaddingTop()
                 || x < getPaddingLeft()
                 || y > getHeight() - getPaddingBottom()
diff --git a/src/com/android/launcher3/ExtendedEditText.java b/src/com/android/launcher3/ExtendedEditText.java
index a4e1af6..21bc479 100644
--- a/src/com/android/launcher3/ExtendedEditText.java
+++ b/src/com/android/launcher3/ExtendedEditText.java
@@ -99,8 +99,18 @@
         }
     }
 
-    // inherited class can override to change the appearance of the edit text.
-    public void show() {}
+    /**
+     * Sets whether EditText background should be visible
+     * @param maxAlpha defines the maximum alpha the background should animates to
+     */
+    public void setBackgroundVisibility(boolean visible, float maxAlpha) {}
+
+    /**
+     * Returns whether a visible background is set on EditText
+     */
+    public boolean getBackgroundVisibility() {
+        return getBackground() != null;
+    }
 
     public void showKeyboard() {
         mShowImeAfterFirstLayout = !showSoftInput();
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 3754dc1..8d14318 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -1195,7 +1195,7 @@
         // Until the workspace is bound, ensure that we keep the wallpaper offset locked to the
         // default state, otherwise we will update to the wrong offsets in RTL
         mWorkspace.lockWallpaperToDefaultPage();
-        mWorkspace.bindAndInitFirstWorkspaceScreen(null /* recycled qsb */);
+        mWorkspace.bindAndInitFirstWorkspaceScreen();
         mDragController.addDragListener(mWorkspace);
 
         // Get the search/delete/uninstall bar
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index 8adbcd9..86217d7 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -50,7 +50,7 @@
 import com.android.launcher3.util.Themes;
 import com.android.launcher3.widget.custom.CustomWidgetManager;
 
-public class LauncherAppState {
+public class LauncherAppState implements SafeCloseable {
 
     public static final String ACTION_FORCE_ROLOAD = "force-reload-launcher";
     private static final String KEY_ICON_STATE = "pref_icon_shape_path";
@@ -162,7 +162,8 @@
     /**
      * Call from Application.onTerminate(), which is not guaranteed to ever be called.
      */
-    public void onTerminate() {
+    @Override
+    public void close() {
         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 7b6a5bf..9ebec0a 100644
--- a/src/com/android/launcher3/LauncherModel.java
+++ b/src/com/android/launcher3/LauncherModel.java
@@ -96,9 +96,10 @@
     // 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;
+            return mModelLoaded && mLoaderTask == null && !mModelDestroyed;
         }
     }
 
@@ -245,6 +246,7 @@
      * Called when the model is destroyed
      */
     public void destroy() {
+        mModelDestroyed = true;
         MODEL_EXECUTOR.execute(mModelDelegate::destroy);
     }
 
@@ -557,6 +559,9 @@
     }
 
     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/PagedView.java b/src/com/android/launcher3/PagedView.java
index 696e897..2f9b5af 100644
--- a/src/com/android/launcher3/PagedView.java
+++ b/src/com/android/launcher3/PagedView.java
@@ -1046,7 +1046,7 @@
     /**
      * If being flinged and user touches the screen, initiate drag; otherwise don't.
      */
-    private void updateIsBeingDraggedOnTouchDown(MotionEvent ev) {
+    protected void updateIsBeingDraggedOnTouchDown(MotionEvent ev) {
         // mScroller.isFinished should be false when being flinged.
         final int xDist = Math.abs(mScroller.getFinalX() - mScroller.getCurrX());
         final boolean finishedScrolling = (mScroller.isFinished() || xDist < mPageSlop / 3);
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index 94ec903..b197e1c 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -223,6 +223,9 @@
     // Variables relating to touch disambiguation (scrolling workspace vs. scrolling a widget)
     private float mXDown;
     private float mYDown;
+    private View mQsb;
+    private boolean mIsEventOverQsb;
+
     final static float START_DAMPING_TOUCH_SLOP_ANGLE = (float) Math.PI / 6;
     final static float MAX_SWIPE_ANGLE = (float) Math.PI / 3;
     final static float TOUCH_SLOP_DAMPING_FACTOR = 4;
@@ -548,9 +551,8 @@
 
     /**
      * Initializes and binds the first page
-     * @param qsb an existing qsb to recycle or null.
      */
-    public void bindAndInitFirstWorkspaceScreen(View qsb) {
+    public void bindAndInitFirstWorkspaceScreen() {
         if (!FeatureFlags.QSB_ON_FIRST_SCREEN) {
             return;
         }
@@ -558,10 +560,10 @@
         // Add the first page
         CellLayout firstPage = insertNewWorkspaceScreen(Workspace.FIRST_SCREEN_ID, getChildCount());
         // Always add a QSB on the first screen.
-        if (qsb == null) {
+        if (mQsb == null) {
             // In transposed layout, we add the QSB in the Grid. As workspace does not touch the
             // edges, we do not need a full width QSB.
-            qsb = LayoutInflater.from(getContext())
+            mQsb = LayoutInflater.from(getContext())
                     .inflate(R.layout.search_container_workspace, firstPage, false);
         }
 
@@ -570,8 +572,9 @@
         CellLayout.LayoutParams lp = new CellLayout.LayoutParams(0, 0, firstPage.getCountX(),
                 cellVSpan);
         lp.canReorder = false;
-        if (!firstPage.addViewToCellLayout(qsb, 0, R.id.search_container_workspace, lp, true)) {
+        if (!firstPage.addViewToCellLayout(mQsb, 0, R.id.search_container_workspace, lp, true)) {
             Log.e(TAG, "Failed to add to item at (0, 0) to CellLayout");
+            mQsb = null;
         }
     }
 
@@ -581,9 +584,8 @@
         disableLayoutTransitions();
 
         // Recycle the QSB widget
-        View qsb = findViewById(R.id.search_container_workspace);
-        if (qsb != null) {
-            ((ViewGroup) qsb.getParent()).removeView(qsb);
+        if (mQsb != null) {
+            ((ViewGroup) mQsb.getParent()).removeView(mQsb);
         }
 
         // Remove the pages and clear the screen models
@@ -596,7 +598,7 @@
         mLauncher.mHandler.removeCallbacksAndMessages(DeferredWidgetRefresh.class);
 
         // Ensure that the first page is always present
-        bindAndInitFirstWorkspaceScreen(qsb);
+        bindAndInitFirstWorkspaceScreen();
 
         // Re-enable the layout transitions
         enableLayoutTransitions();
@@ -922,17 +924,25 @@
     }
 
     @Override
-    public boolean onInterceptTouchEvent(MotionEvent ev) {
-        if (ev.getActionMasked() == MotionEvent.ACTION_DOWN) {
-            mXDown = ev.getX();
-            mYDown = ev.getY();
+    protected void updateIsBeingDraggedOnTouchDown(MotionEvent ev) {
+        super.updateIsBeingDraggedOnTouchDown(ev);
+
+        mXDown = ev.getX();
+        mYDown = ev.getY();
+        if (mQsb != null) {
+            mTempFXY[0] = mXDown + getScrollX();
+            mTempFXY[1] = mYDown + getScrollY();
+            Utilities.mapCoordInSelfToDescendant(mQsb, this, mTempFXY);
+            mIsEventOverQsb = mQsb.getLeft() <= mTempFXY[0] && mQsb.getRight() >= mTempFXY[0]
+                    && mQsb.getTop() <= mTempFXY[1] && mQsb.getBottom() >= mTempFXY[1];
+        } else {
+            mIsEventOverQsb = false;
         }
-        return super.onInterceptTouchEvent(ev);
     }
 
     @Override
     protected void determineScrollingStart(MotionEvent ev) {
-        if (!isFinishedSwitchingState()) return;
+        if (!isFinishedSwitchingState() || mIsEventOverQsb) return;
 
         float deltaX = ev.getX() - mXDown;
         float absDeltaX = Math.abs(deltaX);
diff --git a/src/com/android/launcher3/allapps/AllAppsContainerView.java b/src/com/android/launcher3/allapps/AllAppsContainerView.java
index 9c61c1b..57a3e1c 100644
--- a/src/com/android/launcher3/allapps/AllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/AllAppsContainerView.java
@@ -59,6 +59,7 @@
 import com.android.launcher3.DeviceProfile.OnDeviceProfileChangeListener;
 import com.android.launcher3.DragSource;
 import com.android.launcher3.DropTarget.DragObject;
+import com.android.launcher3.ExtendedEditText;
 import com.android.launcher3.Insettable;
 import com.android.launcher3.InsettableFrameLayout;
 import com.android.launcher3.R;
@@ -118,7 +119,7 @@
     private SpannableStringBuilder mSearchQueryBuilder = null;
 
     protected boolean mUsingTabs;
-    private boolean mSearchModeWhileUsingTabs = false;
+    private boolean mIsSearching;
 
     protected RecyclerViewFastScroller mTouchHandler;
     protected final Point mFastScrollerOffset = new Point();
@@ -132,6 +133,7 @@
     private final float mHeaderThreshold;
     private ScrimView mScrimView;
     private int mHeaderColor;
+    private int mTabsProtectionAlpha;
 
     public AllAppsContainerView(Context context) {
         this(context, null);
@@ -625,18 +627,19 @@
         for (int i = 0; i < mAH.length; i++) {
             mAH[i].adapter.setLastSearchQuery(query);
         }
+        mIsSearching = true;
         if (mUsingTabs) {
-            mSearchModeWhileUsingTabs = true;
             rebindAdapters(false); // hide tabs
         }
         mHeader.setCollapsed(true);
     }
 
     public void onClearSearchResult() {
-        if (mSearchModeWhileUsingTabs) {
+        if (mUsingTabs) {
             rebindAdapters(true); // show tabs
-            mSearchModeWhileUsingTabs = false;
         }
+        mIsSearching = false;
+        getActiveRecyclerView().scrollToTop();
     }
 
     public void onSearchResultsChanged() {
@@ -710,13 +713,12 @@
         mHeaderPaint.setColor(mHeaderColor);
         mHeaderPaint.setAlpha((int) (getAlpha() * Color.alpha(mHeaderColor)));
         if (mHeaderPaint.getColor() != mScrimColor && mHeaderPaint.getColor() != 0) {
-            int bottom = mUsingTabs && mHeader.mHeaderCollapsed ? mHeader.getVisibleBottomBound()
-                    : mSearchContainer.getBottom();
-            canvas.drawRect(0, 0, canvas.getWidth(), bottom + getTranslationY(),
-                    mHeaderPaint);
-
-            if (FeatureFlags.ENABLE_DEVICE_SEARCH.get() && getTranslationY() == 0) {
-                mSearchUiManager.getEditText().setBackground(null);
+            int bottom = (int) (mSearchContainer.getBottom() + getTranslationY());
+            canvas.drawRect(0, 0, canvas.getWidth(), bottom, mHeaderPaint);
+            int tabsHeight = getFloatingHeaderView().getPeripheralProtectionHeight();
+            if (mTabsProtectionAlpha > 0 && tabsHeight != 0) {
+                mHeaderPaint.setAlpha((int) (getAlpha() * mTabsProtectionAlpha));
+                canvas.drawRect(0, bottom, canvas.getWidth(), bottom + tabsHeight, mHeaderPaint);
             }
         }
     }
@@ -796,18 +798,29 @@
 
 
     protected void updateHeaderScroll(int scrolledOffset) {
-        float prog = Math.max(0, Math.min(1, (float) scrolledOffset / mHeaderThreshold));
+
+        float prog = Utilities.boundToRange((float) scrolledOffset / mHeaderThreshold, 0f, 1f);
         int viewBG = ColorUtils.blendARGB(mScrimColor, mHeaderProtectionColor, prog);
         int headerColor = ColorUtils.setAlphaComponent(viewBG,
                 (int) (getSearchView().getAlpha() * 255));
-        if (headerColor != mHeaderColor) {
+        int tabsAlpha = mHeader.getPeripheralProtectionHeight() == 0 ? 0
+                : (int) (Utilities.boundToRange(
+                        (scrolledOffset + mHeader.mSnappedScrolledY) / mHeaderThreshold, 0f, 1f)
+                        * 255);
+        if (headerColor != mHeaderColor || mTabsProtectionAlpha != tabsAlpha) {
             mHeaderColor = headerColor;
-            getSearchView().setBackgroundColor(viewBG);
-            getFloatingHeaderView().setHeaderColor(viewBG);
+            mTabsProtectionAlpha = tabsAlpha;
             invalidateHeader();
-            if (scrolledOffset == 0 && mSearchUiManager.getEditText() != null) {
-                mSearchUiManager.getEditText().show();
+        }
+        if (mSearchUiManager.getEditText() != null) {
+            ExtendedEditText editText = mSearchUiManager.getEditText();
+            boolean bgVisible = editText.getBackgroundVisibility();
+            if (scrolledOffset == 0 && !mIsSearching) {
+                bgVisible = true;
+            } else if (scrolledOffset > mHeaderThreshold) {
+                bgVisible = false;
             }
+            editText.setBackgroundVisibility(bgVisible, 1 - prog);
         }
     }
 
@@ -815,7 +828,7 @@
      * redraws header protection
      */
     public void invalidateHeader() {
-        if (mScrimView != null && FeatureFlags.ENABLE_DEVICE_SEARCH.get()) {
+        if (mScrimView != null && mHeader.isHeaderProtectionSupported()) {
             mScrimView.invalidate();
         }
     }
diff --git a/src/com/android/launcher3/allapps/FloatingHeaderView.java b/src/com/android/launcher3/allapps/FloatingHeaderView.java
index 8ea83d5..debb5b2 100644
--- a/src/com/android/launcher3/allapps/FloatingHeaderView.java
+++ b/src/com/android/launcher3/allapps/FloatingHeaderView.java
@@ -17,9 +17,6 @@
 
 import android.animation.ValueAnimator;
 import android.content.Context;
-import android.graphics.Canvas;
-import android.graphics.Color;
-import android.graphics.Paint;
 import android.graphics.Point;
 import android.graphics.Rect;
 import android.util.ArrayMap;
@@ -50,11 +47,10 @@
         ValueAnimator.AnimatorUpdateListener, PluginListener<AllAppsRow>, Insettable,
         OnHeightUpdatedListener {
 
-    private final Rect mClip = new Rect(0, 0, Integer.MAX_VALUE, Integer.MAX_VALUE);
+    private final Rect mRVClip = new Rect(0, 0, Integer.MAX_VALUE, Integer.MAX_VALUE);
+    private final Rect mHeaderClip = new Rect(0, 0, Integer.MAX_VALUE, Integer.MAX_VALUE);
     private final ValueAnimator mAnimator = ValueAnimator.ofInt(0, 0);
-    private final ValueAnimator mHeaderAnimator = ValueAnimator.ofInt(0, 1).setDuration(100);
     private final Point mTempOffset = new Point();
-    private final Paint mBGPaint = new Paint(Paint.ANTI_ALIAS_FLAG);
     private final RecyclerView.OnScrollListener mOnScrollListener =
             new RecyclerView.OnScrollListener() {
                 @Override
@@ -82,19 +78,19 @@
                 }
             };
 
-    private final int mHeaderTopPadding;
-
     protected final Map<AllAppsRow, PluginHeaderRow> mPluginRows = new ArrayMap<>();
 
+    private final int mHeaderTopPadding;
+    private final boolean mHeaderProtectionSupported;
+
     protected ViewGroup mTabLayout;
     private AllAppsRecyclerView mMainRV;
     private AllAppsRecyclerView mWorkRV;
     private AllAppsRecyclerView mCurrentRV;
     private ViewGroup mParent;
     public boolean mHeaderCollapsed;
-    private int mSnappedScrolledY;
+    protected int mSnappedScrolledY;
     private int mTranslationY;
-    private int mHeaderColor;
 
     private boolean mForwardToRecyclerView;
 
@@ -120,6 +116,8 @@
         super(context, attrs);
         mHeaderTopPadding = context.getResources()
                 .getDimensionPixelSize(R.dimen.all_apps_header_top_padding);
+        mHeaderProtectionSupported = context.getResources().getBoolean(
+                R.bool.config_header_protection_supported);
     }
 
     @Override
@@ -138,7 +136,6 @@
         }
         mFixedRows = rows.toArray(new FloatingHeaderRow[rows.size()]);
         mAllRows = mFixedRows;
-        mHeaderAnimator.addUpdateListener(valueAnimator -> invalidate());
     }
 
     @Override
@@ -285,7 +282,7 @@
                 mHeaderCollapsed = false;
             }
             mTranslationY = currentScrollY;
-        } else if (!mHeaderCollapsed) {
+        } else {
             mTranslationY = currentScrollY - mSnappedScrolledY - mMaxTranslation;
 
             // update state vars
@@ -295,31 +292,10 @@
             } else if (mTranslationY <= -mMaxTranslation) { // hide or stay hidden
                 mHeaderCollapsed = true;
                 mSnappedScrolledY = -mMaxTranslation;
-                mHeaderAnimator.setCurrentFraction(0);
-                mHeaderAnimator.start();
             }
         }
     }
 
-    /**
-     * Set current header protection background color
-     */
-    public void setHeaderColor(int color) {
-        mHeaderColor = color;
-        invalidate();
-    }
-
-    @Override
-    protected void dispatchDraw(Canvas canvas) {
-        if (mHeaderCollapsed && !mCollapsed && mTabLayout.getVisibility() == VISIBLE
-                && mHeaderColor != Color.TRANSPARENT && FeatureFlags.ENABLE_DEVICE_SEARCH.get()) {
-            mBGPaint.setColor(mHeaderColor);
-            mBGPaint.setAlpha((int) (255 * mHeaderAnimator.getAnimatedFraction()));
-            canvas.drawRect(0, 0, getWidth(), getHeight() + mTranslationY, mBGPaint);
-        }
-        super.dispatchDraw(canvas);
-    }
-
     protected void applyVerticalMove() {
         int uncappedTranslationY = mTranslationY;
         mTranslationY = Math.max(mTranslationY, -mMaxTranslation);
@@ -336,11 +312,15 @@
         }
 
         mTabLayout.setTranslationY(mTranslationY);
-        mClip.top = mMaxTranslation + mTranslationY;
+
+        int clipHeight = mHeaderTopPadding - getPaddingBottom();
+        mRVClip.top = mTabsHidden ? clipHeight : 0;
+        mHeaderClip.top = clipHeight;
         // clipping on a draw might cause additional redraw
-        mMainRV.setClipBounds(mClip);
+        setClipBounds(mHeaderClip);
+        mMainRV.setClipBounds(mRVClip);
         if (mWorkRV != null) {
-            mWorkRV.setClipBounds(mClip);
+            mWorkRV.setClipBounds(mRVClip);
         }
     }
 
@@ -421,6 +401,10 @@
         return false;
     }
 
+    public boolean isHeaderProtectionSupported() {
+        return mHeaderProtectionSupported;
+    }
+
     @Override
     public boolean hasOverlappingRendering() {
         return false;
@@ -444,10 +428,19 @@
     }
 
     /**
-     * Returns visible height of FloatingHeaderView contents
+     * Returns visible height of FloatingHeaderView contents requiring header protection
      */
-    public int getVisibleBottomBound() {
-        return getBottom() + mTranslationY;
+    public int getPeripheralProtectionHeight() {
+        if (!mHeaderProtectionSupported) {
+            return 0;
+        }
+
+        // we only want to show protection when work tab is available and header is either
+        // collapsed or animating to/from collapsed state
+        if (mTabsHidden || !mHeaderCollapsed) {
+            return 0;
+        }
+        return Math.max(getHeight() - getPaddingTop() + mTranslationY, 0);
     }
 }
 
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index 74c7310..d5972a7 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -219,7 +219,7 @@
     public static final BooleanFlag ENABLE_TASKBAR = getDebugFlag(
             "ENABLE_TASKBAR", true, "Allows a system Taskbar to be shown on larger devices.");
 
-    public static final BooleanFlag ENABLE_TASKBAR_EDU = getDebugFlag("ENABLE_TASKBAR_EDU", false,
+    public static final BooleanFlag ENABLE_TASKBAR_EDU = getDebugFlag("ENABLE_TASKBAR_EDU", true,
             "Enables showing taskbar education the first time an app is opened.");
 
     public static final BooleanFlag ENABLE_OVERVIEW_GRID = getDebugFlag(
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index 94c578e..43ac8f9 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -87,7 +87,7 @@
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSet;
-import com.android.launcher3.util.MainThreadInitializedObject;
+import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.BaseDragLayer;
 import com.android.launcher3.widget.BaseLauncherAppWidgetHostView;
@@ -98,13 +98,10 @@
 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;
 
 /**
@@ -123,22 +120,16 @@
      * Context used just for preview. It also provides a few objects (e.g. UserCache) just for
      * preview purposes.
      */
-    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));
+    public static class PreviewContext extends SandboxContext {
 
         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);
+            super(base, UserCache.INSTANCE, InstallSessionHelper.INSTANCE,
+                    LauncherAppState.INSTANCE, InvariantDeviceProfile.INSTANCE,
+                    CustomWidgetManager.INSTANCE, PluginManagerWrapper.INSTANCE);
             mIdp = idp;
             mObjectMap.put(InvariantDeviceProfile.INSTANCE, idp);
             mObjectMap.put(LauncherAppState.INSTANCE,
@@ -146,37 +137,6 @@
 
         }
 
-        @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 cd13cd0..1a468ae 100644
--- a/src/com/android/launcher3/icons/IconCache.java
+++ b/src/com/android/launcher3/icons/IconCache.java
@@ -134,6 +134,9 @@
      * 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/logging/StatsLogManager.java b/src/com/android/launcher3/logging/StatsLogManager.java
index d959ee2..5ed651f 100644
--- a/src/com/android/launcher3/logging/StatsLogManager.java
+++ b/src/com/android/launcher3/logging/StatsLogManager.java
@@ -76,6 +76,22 @@
     }
 
     public interface EventEnum {
+
+        /**
+         * Tag used to request new UI Event IDs via presubmit analysis.
+         *
+         * <p>Use RESERVE_NEW_UI_EVENT_ID as the constructor parameter for a new {@link EventEnum}
+         * to signal the presubmit analyzer to reserve a new ID for the event. The new ID will be
+         * returned as a Gerrit presubmit finding.  Do not submit {@code RESERVE_NEW_UI_EVENT_ID} as
+         * the constructor parameter for any event.
+         *
+         * <pre>
+         * &#064;UiEvent(doc = "Briefly describe the interaction when this event will be logged")
+         * UNIQUE_EVENT_NAME(RESERVE_NEW_UI_EVENT_ID);
+         * </pre>
+         */
+        int RESERVE_NEW_UI_EVENT_ID = Integer.MIN_VALUE; // Negative IDs are ignored by the logger.
+
         int getId();
     }
 
@@ -489,8 +505,7 @@
         LAUNCHER_TURN_OFF_WORK_APPS_TAP(839),
 
         @UiEvent(doc = "Launcher item drop failed since there was not enough room on the screen.")
-        LAUNCHER_ITEM_DROP_FAILED_INSUFFICIENT_SPACE(872)
-        ;
+        LAUNCHER_ITEM_DROP_FAILED_INSUFFICIENT_SPACE(872);
 
         // ADD MORE
 
diff --git a/src/com/android/launcher3/util/DisplayController.java b/src/com/android/launcher3/util/DisplayController.java
index 7446181..3a2b961 100644
--- a/src/com/android/launcher3/util/DisplayController.java
+++ b/src/com/android/launcher3/util/DisplayController.java
@@ -58,7 +58,7 @@
  * Utility class to cache properties of default display to avoid a system RPC on every call.
  */
 @SuppressLint("NewApi")
-public class DisplayController implements DisplayListener, ComponentCallbacks {
+public class DisplayController implements DisplayListener, ComponentCallbacks, SafeCloseable {
 
     private static final String TAG = "DisplayController";
 
@@ -82,6 +82,7 @@
     private final ArrayList<DisplayInfoChangeListener> mListeners = new ArrayList<>();
 
     private Info mInfo;
+    private boolean mDestroyed = false;
 
     private DisplayController(Context context) {
         mContext = context;
@@ -119,6 +120,17 @@
     }
 
     @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) { }
 
     @Override
@@ -165,6 +177,9 @@
      * 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 f6003dd..badcd35 100644
--- a/src/com/android/launcher3/util/MainThreadInitializedObject.java
+++ b/src/com/android/launcher3/util/MainThreadInitializedObject.java
@@ -18,13 +18,21 @@
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
 import android.content.Context;
+import android.content.ContextWrapper;
 import android.os.Looper;
+import android.util.Log;
 
+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;
 
 /**
@@ -40,8 +48,8 @@
     }
 
     public T get(Context context) {
-        if (context instanceof PreviewContext) {
-            return ((PreviewContext) context).getObject(this, mProvider);
+        if (context instanceof SandboxContext) {
+            return ((SandboxContext) context).getObject(this, mProvider);
         }
 
         if (mValue == null) {
@@ -80,4 +88,80 @@
 
         T get(Context context);
     }
+
+    /**
+     * Abstract Context which allows custom implementations for
+     * {@link MainThreadInitializedObject} providers
+     */
+    public static abstract class SandboxContext extends ContextWrapper {
+
+        private static final String TAG = "SandboxContext";
+
+        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) {
+                    Log.e(TAG, "Static object access with a destroyed context");
+                }
+                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 10611c7..0c5b722 100644
--- a/src/com/android/launcher3/util/SettingsCache.java
+++ b/src/com/android/launcher3/util/SettingsCache.java
@@ -47,7 +47,7 @@
  *
  * Cache will also be updated if a key queried is missing (even if it has no listeners registered).
  */
-public class SettingsCache extends ContentObserver {
+public class SettingsCache extends ContentObserver implements SafeCloseable {
 
     /** Hidden field Settings.Secure.NOTIFICATION_BADGING */
     public static final Uri NOTIFICATION_BADGING_URI =
@@ -69,7 +69,6 @@
     private final Map<Uri, CopyOnWriteArrayList<OnChangeListener>> mListenerMap = new HashMap<>();
     protected final ContentResolver mResolver;
 
-
     /**
      * Singleton instance
      */
@@ -82,6 +81,11 @@
     }
 
     @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
         // value will exist
diff --git a/src/com/android/launcher3/util/UiThreadHelper.java b/src/com/android/launcher3/util/UiThreadHelper.java
index 0f40179..ac5368c 100644
--- a/src/com/android/launcher3/util/UiThreadHelper.java
+++ b/src/com/android/launcher3/util/UiThreadHelper.java
@@ -28,7 +28,7 @@
 import android.view.View;
 import android.view.inputmethod.InputMethodManager;
 
-import com.android.launcher3.Launcher;
+import com.android.launcher3.BaseActivity;
 import com.android.launcher3.views.ActivityContext;
 
 /**
@@ -56,7 +56,7 @@
                 STATS_LOGGER_KEY,
                 Message.obtain(
                         HANDLER.get(root.getContext()),
-                        () -> Launcher.cast(activityContext)
+                        () -> BaseActivity.fromContext(root.getContext())
                                 .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 329a444..2e2a968 100644
--- a/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
+++ b/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
@@ -33,6 +33,7 @@
 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;
@@ -46,7 +47,7 @@
 /**
  * CustomWidgetManager handles custom widgets implemented as a plugin.
  */
-public class CustomWidgetManager implements PluginListener<CustomWidgetPlugin> {
+public class CustomWidgetManager implements PluginListener<CustomWidgetPlugin>, SafeCloseable {
 
     public static final MainThreadInitializedObject<CustomWidgetManager> INSTANCE =
             new MainThreadInitializedObject<>(CustomWidgetManager::new);
@@ -71,7 +72,8 @@
                 .addPluginListener(this, CustomWidgetPlugin.class, true);
     }
 
-    public void onDestroy() {
+    @Override
+    public void close() {
         PluginManagerWrapper.INSTANCE.get(mContext).removePluginListener(this);
     }
 
diff --git a/robolectric_tests/resources/cache_data_updated_task_data.txt b/tests/res/raw/cache_data_updated_task_data.txt
similarity index 100%
rename from robolectric_tests/resources/cache_data_updated_task_data.txt
rename to tests/res/raw/cache_data_updated_task_data.txt
diff --git a/robolectric_tests/resources/db_schema_v10.json b/tests/res/raw/db_schema_v10.json
similarity index 100%
rename from robolectric_tests/resources/db_schema_v10.json
rename to tests/res/raw/db_schema_v10.json
diff --git a/robolectric_tests/resources/package_install_state_change_task_data.txt b/tests/res/raw/package_install_state_change_task_data.txt
similarity index 100%
rename from robolectric_tests/resources/package_install_state_change_task_data.txt
rename to tests/res/raw/package_install_state_change_task_data.txt
diff --git a/robolectric_tests/resources/widgets_predication_update_task_data.txt b/tests/res/raw/widgets_predication_update_task_data.txt
similarity index 100%
rename from robolectric_tests/resources/widgets_predication_update_task_data.txt
rename to tests/res/raw/widgets_predication_update_task_data.txt
diff --git a/robolectric_tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java b/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
similarity index 84%
rename from robolectric_tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
rename to tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
index 8aa6f37..16f024e 100644
--- a/robolectric_tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
+++ b/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
@@ -13,6 +13,9 @@
 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;
@@ -21,19 +24,17 @@
 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;
@@ -41,8 +42,8 @@
 /**
  * Tests for {@link AddWorkspaceItemsTask}
  */
-@RunWith(RobolectricTestRunner.class)
-@LooperMode(Mode.PAUSED)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class AddWorkspaceItemsTaskTest {
 
     private final ComponentName mComponent1 = new ComponentName("a", "b");
@@ -60,7 +61,7 @@
     @Before
     public void setup() {
         mModelHelper = new LauncherModelHelper();
-        mTargetContext = RuntimeEnvironment.application;
+        mTargetContext = mModelHelper.sandboxContext;
         mIdp = InvariantDeviceProfile.INSTANCE.get(mTargetContext);
         mIdp.numColumns = mIdp.numRows = 5;
         mAppState = LauncherAppState.getInstance(mTargetContext);
@@ -70,6 +71,11 @@
         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) {
@@ -80,6 +86,8 @@
 
     @Test
     public void testFindSpaceForItem_prefers_second() throws Exception {
+        mIdp.isSplitDisplay = false;
+
         // First screen has only one hole of size 1
         int nextId = setupWorkspaceWithHoles(1, 1, new Rect(2, 2, 3, 3));
 
@@ -88,7 +96,7 @@
 
         int[] spaceFound = newTask().findSpaceForItem(
                 mAppState, mModelHelper.getBgDataModel(), mExistingScreens, mNewScreens, 1, 1);
-        assertEquals(2, spaceFound[0]);
+        assertEquals(1, spaceFound[0]);
         assertTrue(mScreenOccupancy.get(spaceFound[0])
                 .isRegionVacant(spaceFound[1], spaceFound[2], 1, 1));
 
@@ -101,6 +109,24 @@
     }
 
     @Test
+    public void testFindSpaceForItem_prefers_third_on_split_display() throws Exception {
+        mIdp.isSplitDisplay = true;
+        // First screen has only one hole of size 1
+        int nextId = setupWorkspaceWithHoles(1, 1, new Rect(2, 2, 3, 3));
+
+        // Second screen has 2 holes of sizes 3x2 and 2x3
+        setupWorkspaceWithHoles(nextId, 2, new Rect(2, 0, 5, 2), new Rect(0, 2, 2, 5));
+
+        int[] spaceFound = newTask().findSpaceForItem(
+                mAppState, mModelHelper.getBgDataModel(), mExistingScreens, mNewScreens, 1, 1);
+        // For split display, it picks the next screen, even if there is enough space
+        // on previous screen
+        assertEquals(2, spaceFound[0]);
+        assertTrue(mScreenOccupancy.get(spaceFound[0])
+                .isRegionVacant(spaceFound[1], spaceFound[2], 1, 1));
+    }
+
+    @Test
     public void testFindSpaceForItem_adds_new_screen() throws Exception {
         // First screen has 2 holes of sizes 3x2 and 2x3
         setupWorkspaceWithHoles(1, 1, new Rect(2, 0, 5, 2), new Rect(0, 2, 2, 5));
@@ -127,7 +153,7 @@
     @Test
     public void testAddItem_some_items_added() throws Exception {
         Callbacks callbacks = mock(Callbacks.class);
-        mModelHelper.getModel().addCallbacks(callbacks);
+        Executors.MAIN_EXECUTOR.submit(() -> mModelHelper.getModel().addCallbacks(callbacks)).get();
 
         WorkspaceItemInfo info = new WorkspaceItemInfo();
         info.intent = new Intent().setComponent(mComponent1);
diff --git a/robolectric_tests/src/com/android/launcher3/model/BackupRestoreTest.java b/tests/src/com/android/launcher3/model/BackupRestoreTest.java
similarity index 62%
rename from robolectric_tests/src/com/android/launcher3/model/BackupRestoreTest.java
rename to tests/src/com/android/launcher3/model/BackupRestoreTest.java
index a397db5..41914de 100644
--- a/robolectric_tests/src/com/android/launcher3/model/BackupRestoreTest.java
+++ b/tests/src/com/android/launcher3/model/BackupRestoreTest.java
@@ -17,6 +17,7 @@
 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;
@@ -26,74 +27,110 @@
 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.robolectric.util.ReflectionHelpers.setField;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
 
 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.os.UserManager;
+import android.util.ArrayMap;
+import android.util.LongSparseArray;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
 
 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.
  */
-@RunWith(RobolectricTestRunner.class)
-@LooperMode(LooperMode.Mode.PAUSED)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class BackupRestoreTest {
 
-    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 static final int PER_USER_RANGE = 200000;
 
-    private ShadowUserManager mUserManager;
+
+    private long mCurrentMyProfileId;
+    private long mOldMyProfileId;
+
+    private long mCurrentWorkProfileId;
+    private long mOldWorkProfileId;
+
     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();
-        mModelHelper = new LauncherModelHelper();
-        RestoreDbTask.setPending(RuntimeEnvironment.application);
+        RestoreDbTask.setPending(mModelHelper.sandboxContext);
         mDb = mModelHelper.provider.getDb();
-        mIdp = InvariantDeviceProfile.INSTANCE.get(RuntimeEnvironment.application);
+        mIdp = InvariantDeviceProfile.INSTANCE.get(mModelHelper.sandboxContext);
+
+    }
+
+    @After
+    public void tearDown() {
+        mUserChangeListener.close();
+        mModelHelper.destroy();
     }
 
     private void setupUserManager() {
-        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);
+        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);
+        }
     }
 
     private void setupBackupManager() {
-        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));
+        mBackupManager = spy(new BackupManager(mModelHelper.sandboxContext));
+        doReturn(myUserHandle()).when(mBackupManager)
+                .getUserForAncestralSerialNumber(eq(mOldMyProfileId));
+        doReturn(mWorkUserHandle).when(mBackupManager)
+                .getUserForAncestralSerialNumber(eq(mOldWorkProfileId));
     }
 
     @Test
@@ -118,18 +155,18 @@
                 { SHORTCUT, SHORTCUT, NO__ICON, NO__ICON},
                 { NO__ICON, NO__ICON, SHORTCUT, SHORTCUT},
                 { APP_ICON, SHORTCUT, SHORTCUT, APP_ICON},
-            }}, 1, MY_OLD_PROFILE_ID);
+            }}, 1, mOldMyProfileId);
         // 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, OLD_WORK_PROFILE_ID);
+            }}, 2, mOldWorkProfileId);
         // simulates the creation of backup upon restore
-        new GridBackupTable(RuntimeEnvironment.application, mDb, mIdp.numDatabaseHotseatIcons,
+        new GridBackupTable(mModelHelper.sandboxContext, mDb, mIdp.numDatabaseHotseatIcons,
                 mIdp.numColumns, mIdp.numRows).doBackup(
-                        MY_OLD_PROFILE_ID, GridBackupTable.OPTION_REQUIRES_SANITIZATION);
+                mOldMyProfileId, GridBackupTable.OPTION_REQUIRES_SANITIZATION);
         // reset favorites table
         createTableUsingOldProfileId();
     }
@@ -141,28 +178,28 @@
     private void verifyTableIsFilled(String tableName, boolean sanitized) {
         assertEquals(sanitized ? 12 : 13, getCount(mDb,
                 "SELECT * FROM " + tableName + " WHERE profileId = "
-                        + (sanitized ? MY_PROFILE_ID : MY_OLD_PROFILE_ID)));
+                        + (sanitized ? mCurrentMyProfileId : mOldMyProfileId)));
         assertEquals(10, getCount(mDb, "SELECT * FROM " + tableName + " WHERE profileId = "
-                + (sanitized ? WORK_PROFILE_ID : OLD_WORK_PROFILE_ID)));
+                + (sanitized ? mCurrentWorkProfileId : mOldWorkProfileId)));
     }
 
     private void createTableUsingOldProfileId() {
         // simulates the creation of favorites table on old device
         dropTable(mDb, TABLE_NAME);
-        addTableToDb(mDb, MY_OLD_PROFILE_ID, false);
+        addTableToDb(mDb, mOldMyProfileId, false);
     }
 
     private void createRestoreSession() throws Exception {
         final PackageInstaller.SessionParams params = new PackageInstaller.SessionParams(
                 PackageInstaller.SessionParams.MODE_FULL_INSTALL);
-        final PackageInstaller installer = RuntimeEnvironment.application.getPackageManager()
+        final PackageInstaller installer = mModelHelper.sandboxContext.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(RuntimeEnvironment.application,
+        RestoreDbTask.restoreIfPossible(mModelHelper.sandboxContext,
                 mModelHelper.provider.getHelper(), mBackupManager);
     }
 
diff --git a/robolectric_tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java b/tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
similarity index 92%
rename from robolectric_tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
rename to tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
index 9ac3fe7..dba0a40 100644
--- a/robolectric_tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
+++ b/tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
@@ -16,6 +16,8 @@
 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;
@@ -26,13 +28,10 @@
 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;
@@ -40,8 +39,8 @@
 /**
  * Tests for {@link CacheDataUpdatedTask}
  */
-@RunWith(RobolectricTestRunner.class)
-@LooperMode(Mode.PAUSED)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class CacheDataUpdatedTaskTest {
 
     private static final String NEW_LABEL_PREFIX = "new-label-";
@@ -51,10 +50,10 @@
     @Before
     public void setup() throws Exception {
         mModelHelper = new LauncherModelHelper();
-        mModelHelper.initializeData("/cache_data_updated_task_data.txt");
+        mModelHelper.initializeData("cache_data_updated_task_data");
 
         // Add placeholder entries in the cache to simulate update
-        Context context = RuntimeEnvironment.application;
+        Context context = mModelHelper.sandboxContext;
         IconCache iconCache = LauncherAppState.getInstance(context).getIconCache();
         CachingLogic<ItemInfo> placeholderLogic = new CachingLogic<ItemInfo>() {
             @Override
@@ -86,6 +85,11 @@
         }
     }
 
+    @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/robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java b/tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
similarity index 91%
rename from robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
rename to tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
index be03c7d..d849c8f 100644
--- a/robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
+++ b/tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
@@ -15,12 +15,13 @@
  */
 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;
@@ -32,6 +33,10 @@
 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;
@@ -40,15 +45,14 @@
 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}
  */
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class DbDowngradeHelperTest {
 
     private static final String SCHEMA_FILE = "test_schema.json";
@@ -60,7 +64,7 @@
 
     @Before
     public void setup() {
-        mContext = RuntimeEnvironment.application;
+        mContext = InstrumentationRegistry.getInstrumentation().getTargetContext();
         mSchemaFile = mContext.getFileStreamPath(SCHEMA_FILE);
         mDbFile = mContext.getDatabasePath(DB_FILE);
     }
@@ -77,8 +81,10 @@
     public void testUpdateSchemaFile() throws Exception {
         // Setup mock resources
         Resources res = spy(mContext.getResources());
-        doAnswer(i ->this.getClass().getResourceAsStream("/db_schema_v10.json"))
-                .when(res).openRawResource(eq(R.raw.downgrade_schema));
+        Resources myRes = getContext().getResources();
+        doAnswer(i -> myRes.openRawResource(
+                myRes.getIdentifier("db_schema_v10", "raw", getContext().getPackageName())))
+                .when(res).openRawResource(R.raw.downgrade_schema);
         Context context = spy(mContext);
         when(context.getResources()).thenReturn(res);
 
diff --git a/robolectric_tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java b/tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
similarity index 77%
rename from robolectric_tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
rename to tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
index 655237d..004ed06 100644
--- a/robolectric_tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
+++ b/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,19 +35,16 @@
 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
  */
-@RunWith(RobolectricTestRunner.class)
-@LooperMode(Mode.PAUSED)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class DefaultLayoutProviderTest {
 
     private LauncherModelHelper mModelHelper;
@@ -56,16 +53,18 @@
     @Before
     public void setUp() {
         mModelHelper = new LauncherModelHelper();
-        mTargetContext = RuntimeEnvironment.application;
+        mTargetContext = mModelHelper.sandboxContext;
+    }
 
-        shadowOf(mTargetContext.getPackageManager())
-                .addActivityIfNotPresent(new ComponentName(TEST_PACKAGE, TEST_PACKAGE));
+    @After
+    public void tearDown() {
+        mModelHelper.destroy();
     }
 
     @Test
     public void testCustomProfileLoaded_with_icon_on_hotseat() throws Exception {
         writeLayoutAndLoad(new LauncherLayoutBuilder().atHotseat(0)
-                .putApp(TEST_PACKAGE, TEST_PACKAGE));
+                .putApp(TEST_PACKAGE, TEST_ACTIVITY));
 
         // Verify one item in hotseat
         assertEquals(1, mModelHelper.getBgDataModel().workspaceItems.size());
@@ -77,9 +76,9 @@
     @Test
     public void testCustomProfileLoaded_with_folder() throws Exception {
         writeLayoutAndLoad(new LauncherLayoutBuilder().atHotseat(0).putFolder(android.R.string.copy)
-                .addApp(TEST_PACKAGE, TEST_PACKAGE)
-                .addApp(TEST_PACKAGE, TEST_PACKAGE)
-                .addApp(TEST_PACKAGE, TEST_PACKAGE)
+                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
+                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
+                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
                 .build());
 
         // Verify folder
@@ -92,9 +91,9 @@
     @Test
     public void testCustomProfileLoaded_with_folder_custom_title() throws Exception {
         writeLayoutAndLoad(new LauncherLayoutBuilder().atHotseat(0).putFolder("CustomFolder")
-                .addApp(TEST_PACKAGE, TEST_PACKAGE)
-                .addApp(TEST_PACKAGE, TEST_PACKAGE)
-                .addApp(TEST_PACKAGE, TEST_PACKAGE)
+                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
+                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
+                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
                 .build());
 
         // Verify folder
@@ -112,12 +111,10 @@
         // 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();
-        int sessionId = installer.createSession(params);
-        SessionInfo sessionInfo = installer.getSessionInfo(sessionId);
-        setField(sessionInfo, "installerPackageName", "com.test");
-        setField(sessionInfo, "appIcon", BitmapInfo.LOW_RES_ICON);
+        installer.createSession(params);
 
         writeLayoutAndLoad(new LauncherLayoutBuilder().atWorkspace(0, 1, 0)
                 .putWidget(pendingAppPkg, "PlaceholderWidget", 2, 2));
diff --git a/robolectric_tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java b/tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
similarity index 88%
rename from robolectric_tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
rename to tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
index 87b0887..005389e 100644
--- a/robolectric_tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
+++ b/tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
@@ -30,26 +30,31 @@
 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} */
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class GridSizeMigrationTaskV2Test {
 
     private LauncherModelHelper mModelHelper;
@@ -73,7 +78,7 @@
     @Before
     public void setUp() {
         mModelHelper = new LauncherModelHelper();
-        mContext = RuntimeEnvironment.application;
+        mContext = mModelHelper.sandboxContext;
         mDb = mModelHelper.provider.getDb();
 
         mValidPackages = new HashSet<>();
@@ -98,8 +103,13 @@
                 LauncherSettings.Favorites.TMP_TABLE);
     }
 
+    @After
+    public void tearDown() {
+        mModelHelper.destroy();
+    }
+
     @Test
-    public void testMigration() {
+    public void testMigration() throws Exception {
         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),
@@ -134,17 +144,17 @@
         // Check hotseat items
         Cursor c = mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
                 new String[]{LauncherSettings.Favorites.SCREEN, LauncherSettings.Favorites.INTENT},
-                "container=" + CONTAINER_HOTSEAT, null, null, null);
+                "container=" + CONTAINER_HOTSEAT, null, LauncherSettings.Favorites.SCREEN, 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), 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), 1);
+        assertTrue(c.getString(intentIndex).contains(testPackage2));
+        c.moveToNext();
         assertEquals(c.getInt(screenIndex), 2);
         assertTrue(c.getString(intentIndex).contains(testPackage3));
         c.moveToNext();
@@ -157,35 +167,24 @@
                 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);
 
-        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);
-
+        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();
+
+        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));
     }
 
     @Test
@@ -212,7 +211,7 @@
         // Check hotseat items
         Cursor c = mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
                 new String[]{LauncherSettings.Favorites.SCREEN, LauncherSettings.Favorites.INTENT},
-                "container=" + CONTAINER_HOTSEAT, null, null, null);
+                "container=" + CONTAINER_HOTSEAT, null, LauncherSettings.Favorites.SCREEN, null);
         assertEquals(c.getCount(), numSrcDatabaseHotseatIcons);
         int screenIndex = c.getColumnIndex(LauncherSettings.Favorites.SCREEN);
         int intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
@@ -257,7 +256,7 @@
         // Check hotseat items
         Cursor c = mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
                 new String[]{LauncherSettings.Favorites.SCREEN, LauncherSettings.Favorites.INTENT},
-                "container=" + CONTAINER_HOTSEAT, null, null, null);
+                "container=" + CONTAINER_HOTSEAT, null, LauncherSettings.Favorites.SCREEN, null);
         assertEquals(c.getCount(), mIdp.numDatabaseHotseatIcons);
         int screenIndex = c.getColumnIndex(LauncherSettings.Favorites.SCREEN);
         int intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
diff --git a/robolectric_tests/src/com/android/launcher3/model/LoaderCursorTest.java b/tests/src/com/android/launcher3/model/LoaderCursorTest.java
similarity index 93%
rename from robolectric_tests/src/com/android/launcher3/model/LoaderCursorTest.java
rename to tests/src/com/android/launcher3/model/LoaderCursorTest.java
index 800311a..6444ef6 100644
--- a/robolectric_tests/src/com/android/launcher3/model/LoaderCursorTest.java
+++ b/tests/src/com/android/launcher3/model/LoaderCursorTest.java
@@ -16,6 +16,8 @@
 
 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;
@@ -33,7 +35,7 @@
 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_PACKAGE;
+import static com.android.launcher3.util.LauncherModelHelper.TEST_ACTIVITY;
 
 import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertFalse;
@@ -41,37 +43,37 @@
 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}
  */
-@RunWith(RobolectricTestRunner.class)
-@LooperMode(Mode.PAUSED)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class LoaderCursorTest {
 
+    private LauncherModelHelper mModelHelper;
     private LauncherAppState mApp;
 
     private MatrixCursor mCursor;
@@ -82,7 +84,8 @@
 
     @Before
     public void setup() {
-        mContext = RuntimeEnvironment.application;
+        mModelHelper = new LauncherModelHelper();
+        mContext = mModelHelper.sandboxContext;
         mIDP = InvariantDeviceProfile.INSTANCE.get(mContext);
         mApp = LauncherAppState.getInstance(mContext);
 
@@ -97,6 +100,11 @@
         ums.allUsers.put(0, Process.myUserHandle());
     }
 
+    @After
+    public void tearDown() {
+        mModelHelper.destroy();
+    }
+
     private void initCursor(int itemType, String title) {
         mCursor.newRow()
                 .add(_ID, 1)
@@ -117,9 +125,7 @@
 
     @Test
     public void getAppShortcutInfo_dontAllowMissing_validComponent() throws Exception {
-        ComponentName cn = new ComponentName(TEST_PACKAGE, TEST_PACKAGE);
-        shadowOf(mContext.getPackageManager()).addActivityIfNotPresent(cn);
-
+        ComponentName cn = new ComponentName(getContext(), TEST_ACTIVITY);
         initCursor(ITEM_TYPE_APPLICATION, "");
         assertTrue(mLoaderCursor.moveToNext());
 
diff --git a/robolectric_tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java b/tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
similarity index 65%
rename from robolectric_tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
rename to tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
index 4319355..42c9f11 100644
--- a/robolectric_tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
+++ b/tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
@@ -15,69 +15,60 @@
  */
 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
  */
-@RunWith(RobolectricTestRunner.class)
-@LooperMode(Mode.PAUSED)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class ModelMultiCallbacksTest {
 
     private LauncherModelHelper mModelHelper;
 
-    private ShadowPackageManager mSpm;
-    private LooperExecutor mTempMainExecutor;
-
     @Before
-    public void setUp() throws Exception {
+    public void setUp() {
         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());
+    @After
+    public void tearDown() throws Exception {
+        mModelHelper.destroy();
+        TestUtil.uninstallDummyApp();
     }
 
     @Test
@@ -85,7 +76,7 @@
         setupWorkspacePages(3);
 
         MyCallbacks cb1 = spy(MyCallbacks.class);
-        mModelHelper.getModel().addCallbacksAndLoad(cb1);
+        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().addCallbacksAndLoad(cb1));
 
         waitForLoaderAndTempMainThread();
         cb1.verifySynchronouslyBound(3);
@@ -94,10 +85,10 @@
         cb1.reset();
         MyCallbacks cb2 = spy(MyCallbacks.class);
         cb2.mPageToBindSync = IntSet.wrap(2);
-        mModelHelper.getModel().addCallbacksAndLoad(cb2);
+        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().addCallbacksAndLoad(cb2));
 
         waitForLoaderAndTempMainThread();
-        cb1.verifySynchronouslyBound(3);
+        assertFalse(cb1.bindStarted);
         cb2.verifySynchronouslyBound(3);
 
         // Remove callbacks
@@ -105,7 +96,7 @@
         cb2.reset();
 
         // No effect on callbacks when removing an callback
-        mModelHelper.getModel().removeCallbacks(cb2);
+        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().removeCallbacks(cb2));
         waitForLoaderAndTempMainThread();
         assertNull(cb1.mPendingTasks);
         assertNull(cb2.mPendingTasks);
@@ -119,52 +110,48 @@
 
     @Test
     public void testTwoCallbacks_receiveUpdates() throws Exception {
+        TestUtil.uninstallDummyApp();
+
         setupWorkspacePages(1);
 
         MyCallbacks cb1 = spy(MyCallbacks.class);
         MyCallbacks cb2 = spy(MyCallbacks.class);
-        mModelHelper.getModel().addCallbacksAndLoad(cb1);
-        mModelHelper.getModel().addCallbacksAndLoad(cb2);
+        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().addCallbacksAndLoad(cb1));
+        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().addCallbacksAndLoad(cb2));
         waitForLoaderAndTempMainThread();
 
-        cb1.verifyApps(TEST_PACKAGE);
-        cb2.verifyApps(TEST_PACKAGE);
+        assertTrue(cb1.allApps().contains(TEST_PACKAGE));
+        assertTrue(cb2.allApps().contains(TEST_PACKAGE));
 
         // Install package 1
-        String pkg1 = "com.test.pkg1";
-        mModelHelper.installApp(pkg1);
-        mModelHelper.getModel().onPackageAdded(pkg1, Process.myUserHandle());
+        TestUtil.installDummyApp();
+        mModelHelper.getModel().onPackageAdded(TestUtil.DUMMY_PACKAGE, Process.myUserHandle());
         waitForLoaderAndTempMainThread();
-        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);
+        assertTrue(cb1.allApps().contains(TestUtil.DUMMY_PACKAGE));
+        assertTrue(cb2.allApps().contains(TestUtil.DUMMY_PACKAGE));
 
         // Uninstall package 2
-        mSpm.removePackage(pkg1);
-        mModelHelper.getModel().onPackageRemoved(pkg1, Process.myUserHandle());
+        TestUtil.uninstallDummyApp();
+        mModelHelper.getModel().onPackageRemoved(TestUtil.DUMMY_PACKAGE, Process.myUserHandle());
         waitForLoaderAndTempMainThread();
-        cb1.verifyApps(TEST_PACKAGE, pkg2);
-        cb2.verifyApps(TEST_PACKAGE, pkg2);
+        assertFalse(cb1.allApps().contains(TestUtil.DUMMY_PACKAGE));
+        assertFalse(cb2.allApps().contains(TestUtil.DUMMY_PACKAGE));
 
         // Unregister a callback and verify updates no longer received
-        mModelHelper.getModel().removeCallbacks(cb2);
-        mSpm.removePackage(pkg2);
-        mModelHelper.getModel().onPackageRemoved(pkg2, Process.myUserHandle());
+        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().removeCallbacks(cb2));
+        TestUtil.installDummyApp();
+        mModelHelper.getModel().onPackageAdded(TestUtil.DUMMY_PACKAGE, Process.myUserHandle());
         waitForLoaderAndTempMainThread();
-        cb1.verifyApps(TEST_PACKAGE);
-        cb2.verifyApps(TEST_PACKAGE, pkg2);
+
+        // cb2 didn't get the update
+        assertTrue(cb1.allApps().contains(TestUtil.DUMMY_PACKAGE));
+        assertFalse(cb2.allApps().contains(TestUtil.DUMMY_PACKAGE));
     }
 
     private void waitForLoaderAndTempMainThread() throws Exception {
+        Executors.MAIN_EXECUTOR.submit(() -> { }).get();
         Executors.MODEL_EXECUTOR.submit(() -> { }).get();
-        mTempMainExecutor.submit(() -> { }).get();
+        Executors.MAIN_EXECUTOR.submit(() -> { }).get();
     }
 
     private void setupWorkspacePages(int pageCount) throws Exception {
@@ -183,10 +170,16 @@
         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;
             mPendingTasks = pendingTasks;
@@ -212,10 +205,12 @@
             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);
@@ -226,12 +221,14 @@
             assertEquals(mItems.size(), totalItems);
         }
 
-        public void verifyApps(String... apps) {
-            assertEquals(apps.length, mAppInfos.length);
-            assertEquals(Arrays.stream(mAppInfos)
+        public Set<String> allApps() {
+            return Arrays.stream(mAppInfos)
                     .map(ai -> ai.getTargetComponent().getPackageName())
-                    .collect(Collectors.toSet()),
-                    new HashSet<>(Arrays.asList(apps)));
+                    .collect(Collectors.toSet());
+        }
+
+        public void verifyApps(String... apps) {
+            assertTrue(allApps().containsAll(Arrays.asList(apps)));
         }
     }
 }
diff --git a/robolectric_tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java b/tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
similarity index 87%
rename from robolectric_tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
rename to tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
index 412ace0..519191e 100644
--- a/robolectric_tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
+++ b/tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
@@ -2,18 +2,19 @@
 
 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;
@@ -21,8 +22,8 @@
 /**
  * Tests for {@link PackageInstallStateChangedTask}
  */
-@RunWith(RobolectricTestRunner.class)
-@LooperMode(Mode.PAUSED)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class PackageInstallStateChangedTaskTest {
 
     private LauncherModelHelper mModelHelper;
@@ -30,7 +31,12 @@
     @Before
     public void setup() throws Exception {
         mModelHelper = new LauncherModelHelper();
-        mModelHelper.initializeData("/package_install_state_change_task_data.txt");
+        mModelHelper.initializeData("package_install_state_change_task_data");
+    }
+
+    @After
+    public void tearDown() {
+        mModelHelper.destroy();
     }
 
     private PackageInstallStateChangedTask newTask(String pkg, int progress) {
@@ -66,7 +72,7 @@
         HashSet<Integer> updates = new HashSet<>(Arrays.asList(idsUpdated));
         for (ItemInfo info : mModelHelper.getBgDataModel().itemsIdMap) {
             if (info instanceof WorkspaceItemInfo) {
-                assertEquals(updates.contains(info.id) ? progress: 0,
+                assertEquals(updates.contains(info.id) ? progress: 100,
                         ((WorkspaceItemInfo) info).getProgressLevel());
             } else {
                 assertEquals(updates.contains(info.id) ? progress: -1,
diff --git a/robolectric_tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java b/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
similarity index 92%
rename from robolectric_tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
rename to tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
index 4184d33..48305ee 100644
--- a/robolectric_tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
+++ b/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
@@ -21,18 +21,21 @@
 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}
  */
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class RestoreDbTaskTest {
 
     @Test
@@ -95,7 +98,7 @@
         private final long mProfileId;
 
         MyDatabaseHelper(long profileId) {
-            super(RuntimeEnvironment.application, null, false);
+            super(InstrumentationRegistry.getInstrumentation().getTargetContext(), null, false);
             mProfileId = profileId;
         }
 
diff --git a/tests/src/com/android/launcher3/secondarydisplay/SDLauncherTest.java b/tests/src/com/android/launcher3/secondarydisplay/SDLauncherTest.java
new file mode 100644
index 0000000..fd86cf1
--- /dev/null
+++ b/tests/src/com/android/launcher3/secondarydisplay/SDLauncherTest.java
@@ -0,0 +1,53 @@
+/*
+ * 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/robolectric_tests/src/com/android/launcher3/util/LauncherLayoutBuilder.java b/tests/src/com/android/launcher3/util/LauncherLayoutBuilder.java
similarity index 100%
rename from robolectric_tests/src/com/android/launcher3/util/LauncherLayoutBuilder.java
rename to tests/src/com/android/launcher3/util/LauncherLayoutBuilder.java
diff --git a/robolectric_tests/src/com/android/launcher3/util/LauncherModelHelper.java b/tests/src/com/android/launcher3/util/LauncherModelHelper.java
similarity index 60%
rename from robolectric_tests/src/com/android/launcher3/util/LauncherModelHelper.java
rename to tests/src/com/android/launcher3/util/LauncherModelHelper.java
index 846e201..c9b63ae 100644
--- a/robolectric_tests/src/com/android/launcher3/util/LauncherModelHelper.java
+++ b/tests/src/com/android/launcher3/util/LauncherModelHelper.java
@@ -15,26 +15,40 @@
  */
 package com.android.launcher3.util;
 
-import static android.content.Intent.ACTION_CREATE_SHORTCUT;
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 
 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.IntentFilter;
-import android.content.pm.PackageManager.NameNotFoundException;
+import android.content.pm.PackageManager;
+import android.content.pm.ProviderInfo;
+import android.content.res.Resources;
 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;
@@ -45,27 +59,30 @@
 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.shadows.ShadowLooperExecutor;
+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 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;
@@ -81,7 +98,9 @@
     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 = "com.android.launcher3.validpackage";
+
+    public static final String TEST_PACKAGE = testContext().getPackageName();
+    public static final String TEST_ACTIVITY = "com.android.launcher3.tests.Activity2";
 
     // Authority for providing a test default-workspace-layout data.
     private static final String TEST_PROVIDER_AUTHORITY =
@@ -90,21 +109,42 @@
     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;
-    private final long mDefaultProfileId;
+    public final SanboxModelContext sandboxContext;
+
+    public final long defaultProfileId;
 
     private BgDataModel mDataModel;
     private AllAppsList mAllAppsList;
 
     public LauncherModelHelper() {
-        provider = Robolectric.setupContentProvider(TestLauncherProvider.class);
-        mDefaultProfileId = UserCache.INSTANCE.get(RuntimeEnvironment.application)
+        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)
                 .getSerialNumberForUser(Process.myUserHandle());
-        ShadowContentResolver.registerProviderInternal(LauncherProvider.AUTHORITY, provider);
+        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());
     }
 
     public LauncherModel getModel() {
-        return LauncherAppState.getInstance(RuntimeEnvironment.application).getModel();
+        return LauncherAppState.getInstance(sandboxContext).getModel();
     }
 
     public synchronized BgDataModel getBgDataModel() {
@@ -121,6 +161,28 @@
         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.
      */
@@ -161,13 +223,16 @@
      * 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(
-                    this.getClass().getResourceAsStream(resourceName)))) {
+                    resources.openRawResource(resId)))) {
                 String line;
                 HashMap<String, Class> classMap = new HashMap<>();
                 while ((line = reader.readLine()) != null) {
@@ -181,7 +246,7 @@
                             classMap.put(commands[1], Class.forName(commands[2]));
                             break;
                         case "bgItem":
-                            bgDataModel.addItem(targetContext,
+                            bgDataModel.addItem(sandboxContext,
                                     (ItemInfo) initItem(classMap.get(commands[1]), commands, 2),
                                     false);
                             break;
@@ -236,7 +301,7 @@
     }
 
     public int addItem(int type, int screen, int container, int x, int y) {
-        return addItem(type, screen, container, x, y, mDefaultProfileId, TEST_PACKAGE);
+        return addItem(type, screen, container, x, y, defaultProfileId, TEST_PACKAGE);
     }
 
     public int addItem(int type, int screen, int container, int x, int y, long profileId) {
@@ -244,12 +309,12 @@
     }
 
     public int addItem(int type, int screen, int container, int x, int y, String packageName) {
-        return addItem(type, screen, container, x, y, mDefaultProfileId, packageName);
+        return addItem(type, screen, container, x, y, defaultProfileId, 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, mDefaultProfileId, packageName, id, contentUri);
+        addItem(type, screen, container, x, y, defaultProfileId, packageName, id, contentUri);
         return id;
     }
 
@@ -260,8 +325,7 @@
      */
     public int addItem(int type, int screen, int container, int x, int y, long profileId,
             String packageName) {
-        Context context = RuntimeEnvironment.application;
-        int id = LauncherSettings.Settings.call(context.getContentResolver(),
+        int id = LauncherSettings.Settings.call(sandboxContext.getContentResolver(),
                 LauncherSettings.Settings.METHOD_NEW_ITEM_ID)
                 .getInt(LauncherSettings.Settings.EXTRA_VALUE);
         addItem(type, screen, container, x, y, profileId, packageName, id, CONTENT_URI);
@@ -270,8 +334,6 @@
 
     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);
@@ -295,7 +357,7 @@
             }
         }
 
-        context.getContentResolver().insert(contentUri, values);
+        sandboxContext.getContentResolver().insert(contentUri, values);
     }
 
     public int[][][] createGrid(int[][][] typeArray) {
@@ -303,12 +365,11 @@
     }
 
     public int[][][] createGrid(int[][][] typeArray, int startScreen) {
-        final Context context = RuntimeEnvironment.application;
-        LauncherSettings.Settings.call(context.getContentResolver(),
+        LauncherSettings.Settings.call(sandboxContext.getContentResolver(),
                 LauncherSettings.Settings.METHOD_CREATE_EMPTY_DB);
-        LauncherSettings.Settings.call(context.getContentResolver(),
+        LauncherSettings.Settings.call(sandboxContext.getContentResolver(),
                 LauncherSettings.Settings.METHOD_CLEAR_EMPTY_DB_FLAG);
-        return createGrid(typeArray, startScreen, mDefaultProfileId);
+        return createGrid(typeArray, startScreen, defaultProfileId);
     }
 
     /**
@@ -320,14 +381,13 @@
      * @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(context.getContentResolver(),
+            LauncherSettings.Settings.call(sandboxContext.getContentResolver(),
                     LauncherSettings.Settings.METHOD_NEW_SCREEN_ID);
 
             ids[i] = new int[typeArray[i].length][];
@@ -353,69 +413,45 @@
      */
     public LauncherModelHelper setupDefaultLayoutProvider(LauncherLayoutBuilder builder)
             throws Exception {
-        Context context = RuntimeEnvironment.application;
-        InvariantDeviceProfile idp = InvariantDeviceProfile.INSTANCE.get(context);
+        InvariantDeviceProfile idp = InvariantDeviceProfile.INSTANCE.get(sandboxContext);
         idp.numRows = idp.numColumns = idp.numDatabaseHotseatIcons = DEFAULT_GRID_SIZE;
         idp.iconBitmapSize = DEFAULT_BITMAP_SIZE;
 
-        Settings.Secure.putString(context.getContentResolver(),
-                "launcher3.layout.provider", TEST_PROVIDER_AUTHORITY);
+        UiDevice.getInstance(getInstrumentation()).executeShellCommand(
+                "settings put secure launcher3.layout.provider " + TEST_PROVIDER_AUTHORITY);
+        ContentProvider cp = new TestInformationProvider() {
 
-        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()));
+            @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);
         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 {
-        // 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);
+        Callbacks mockCb = new Callbacks() { };
+        Executors.MAIN_EXECUTOR.submit(() -> getModel().addCallbacksAndLoad(mockCb)).get();
 
         Executors.MODEL_EXECUTOR.submit(() -> { }).get();
-        Executors.UI_HELPER_EXECUTOR.submit(() -> { }).get();
-
-        sle.setHandler(null);
-        getModel().removeCallbacks(mockCb);
+        Executors.MAIN_EXECUTOR.submit(() -> { }).get();
+        Executors.MAIN_EXECUTOR.submit(() -> getModel().removeCallbacks(mockCb)).get();
     }
 
     /**
@@ -437,4 +473,97 @@
             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/tests/src/com/android/launcher3/util/ReflectionHelpers.java b/tests/src/com/android/launcher3/util/ReflectionHelpers.java
new file mode 100644
index 0000000..d89975d
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/ReflectionHelpers.java
@@ -0,0 +1,58 @@
+/*
+ * 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);
+        }
+    }
+
+}
diff --git a/tests/src/com/android/launcher3/util/rule/FailureWatcher.java b/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
index fae55cc..f9a9997 100644
--- a/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
+++ b/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
@@ -15,6 +15,7 @@
 import org.junit.runner.Description;
 import org.junit.runners.model.Statement;
 
+import java.io.BufferedOutputStream;
 import java.io.File;
 import java.io.FileOutputStream;
 import java.io.IOException;
@@ -66,15 +67,18 @@
         onError(mDevice, description, e);
     }
 
+    static File diagFile(Description description, String prefix, String ext) {
+        return new File(getInstrumentation().getTargetContext().getFilesDir(),
+                prefix + "-" + description.getTestClass().getSimpleName() + "."
+                        + description.getMethodName() + "." + ext);
+    }
+
     public static void onError(UiDevice device, Description description, Throwable e) {
         Log.d("b/196820244", "onError 1");
         if (device == null) return;
         Log.d("b/196820244", "onError 2");
-        final File parentFile = getInstrumentation().getTargetContext().getFilesDir();
-        final File sceenshot = new File(parentFile,
-                "TestScreenshot-" + description.getMethodName() + ".png");
-        final File hierarchy = new File(parentFile,
-                "Hierarchy-" + description.getMethodName() + ".zip");
+        final File sceenshot = diagFile(description, "TestScreenshot", "png");
+        final File hierarchy = diagFile(description, "Hierarchy", "zip");
 
         // Dump window hierarchy
         try (ZipOutputStream out = new ZipOutputStream(new FileOutputStream(hierarchy))) {
@@ -97,13 +101,13 @@
         device.takeScreenshot(sceenshot);
 
         // Dump accessibility hierarchy
-        final File accessibilityHierarchyFile = new File(parentFile,
-                "AccessibilityHierarchy-" + description.getMethodName() + ".uix");
         try {
-            device.dumpWindowHierarchy(accessibilityHierarchyFile);
+            device.dumpWindowHierarchy(diagFile(description, "AccessibilityHierarchy", "uix"));
         } catch (IOException ex) {
             Log.e(TAG, "Failed to save accessibility hierarchy", ex);
         }
+
+        dumpCommand("logcat -d -s TestRunner", diagFile(description, "FilteredLogcat", "txt"));
     }
 
     private static void dumpStringCommand(String cmd, OutputStream out) throws IOException {
@@ -111,6 +115,14 @@
         dumpCommand(cmd, out);
     }
 
+    private static void dumpCommand(String cmd, File out) {
+        try (BufferedOutputStream buffered = new BufferedOutputStream(
+                new FileOutputStream(out))) {
+            dumpCommand(cmd, buffered);
+        } catch (IOException ex) {
+        }
+    }
+
     private static void dumpCommand(String cmd, OutputStream out) throws IOException {
         try (AutoCloseInputStream in = new AutoCloseInputStream(getInstrumentation()
                 .getUiAutomation().executeShellCommand(cmd))) {