diff --git a/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java b/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
index 6502526dcc..bc38739d57 100644
--- a/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
+++ b/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
@@ -50,7 +50,6 @@ import androidx.annotation.IntDef;
 import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.BaseActivity;
-import com.android.launcher3.BaseDraggingActivity;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.views.ArrowTipView;
@@ -76,7 +75,7 @@ public final class TaskOverlayFactoryGo extends TaskOverlayFactory {
     public static final String ACTIONS_ERROR_CODE = "niu_actions_app_error_code";
     public static final int ERROR_PERMISSIONS_STRUCTURE = 1;
     public static final int ERROR_PERMISSIONS_SCREENSHOT = 2;
-    private static final String NIU_ACTIONS_CONFIRMED = "launcher_go.niu_actions_confirmed";
+    public static final String NIU_ACTIONS_CONFIRMED = "launcher_go.niu_actions_confirmed";
     private static final String ASSIST_SETTINGS_ARGS_BUNDLE = ":settings:show_fragment_args";
     private static final String ASSIST_SETTINGS_ARGS_KEY = ":settings:fragment_args_key";
     private static final String ASSIST_SETTINGS_PREFERENCE_KEY = "default_assist";
@@ -87,10 +86,11 @@ public final class TaskOverlayFactoryGo extends TaskOverlayFactory {
 
     @Retention(SOURCE)
     @IntDef({PRIVACY_CONFIRMATION, ASSISTANT_NOT_SELECTED, ASSISTANT_NOT_SUPPORTED})
-    private @interface DialogType{}
-    private static final int PRIVACY_CONFIRMATION = 0;
-    private static final int ASSISTANT_NOT_SELECTED = 1;
-    private static final int ASSISTANT_NOT_SUPPORTED = 2;
+    @VisibleForTesting
+    public @interface DialogType{}
+    public static final int PRIVACY_CONFIRMATION = 0;
+    public static final int ASSISTANT_NOT_SELECTED = 1;
+    public static final int ASSISTANT_NOT_SUPPORTED = 2;
 
     private AssistContentRequester mContentRequester;
 
@@ -211,7 +211,8 @@ public final class TaskOverlayFactoryGo extends TaskOverlayFactory {
             Intent intent = createNIUIntent(actionType);
             // Only add and send the image if the appropriate permissions are held
             if (mAssistStructurePermitted && mAssistScreenshotPermitted) {
-                mImageApi.shareAsDataWithExplicitIntent(/* crop */ null, intent);
+                mImageApi.shareAsDataWithExplicitIntent(/* crop */ null, intent,
+                        () -> showDialog(actionType, ASSISTANT_NOT_SUPPORTED));
             } else {
                 // If both permissions are disabled, the structure error code takes priority
                 // The user must enable that one before they can enable screenshots
@@ -301,7 +302,6 @@ public final class TaskOverlayFactoryGo extends TaskOverlayFactory {
             mImageApi = imageActionsApi;
         }
 
-        // TODO (b/192406446): Test that these dialogs are shown at the appropriate times
         private void showDialog(String action, @DialogType int type) {
             switch (type) {
                 case PRIVACY_CONFIRMATION:
@@ -334,7 +334,7 @@ public final class TaskOverlayFactoryGo extends TaskOverlayFactory {
                                 int bodyTextID, int button1TextID,
                                 View.OnClickListener button1Callback, int button2TextID,
                                 View.OnClickListener button2Callback) {
-            BaseDraggingActivity activity = BaseActivity.fromContext(getActionsView().getContext());
+            BaseActivity activity = BaseActivity.fromContext(getActionsView().getContext());
             LayoutInflater inflater = LayoutInflater.from(activity);
             View view = inflater.inflate(R.layout.niu_actions_dialog, /* root */ null);
 
@@ -368,6 +368,11 @@ public final class TaskOverlayFactoryGo extends TaskOverlayFactory {
             mDialog.cancel();
         }
 
+        @VisibleForTesting
+        public OverlayDialogGo getDialog() {
+            return mDialog;
+        }
+
         private void onDialogClickSettings(View v) {
             mDialog.dismiss();
 
@@ -401,7 +406,11 @@ public final class TaskOverlayFactoryGo extends TaskOverlayFactory {
         }
     }
 
-    private static final class OverlayDialogGo extends AlertDialog {
+    /**
+     * Basic modal dialog for various user prompts
+     */
+    @VisibleForTesting
+    public static final class OverlayDialogGo extends AlertDialog {
         private final String mAction;
         private final @DialogType int mType;
 
diff --git a/quickstep/res/layout/taskbar_edu.xml b/quickstep/res/layout/taskbar_edu.xml
index b7717b71be..ef57a53690 100644
--- a/quickstep/res/layout/taskbar_edu.xml
+++ b/quickstep/res/layout/taskbar_edu.xml
@@ -17,6 +17,7 @@
 <com.android.launcher3.taskbar.TaskbarEduView xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:androidprv="http://schemas.android.com/apk/prv/res/android"
     xmlns:app="http://schemas.android.com/apk/res-auto"
+    xmlns:launcher="http://schemas.android.com/apk/res-auto"
     android:layout_width="match_parent"
     android:layout_height="wrap_content"
     android:layout_gravity="bottom"
@@ -33,46 +34,121 @@
         android:paddingHorizontal="36dp"
         android:paddingTop="64dp">
 
-        <TextView
-            android:id="@+id/edu_header"
+        <com.android.launcher3.taskbar.TaskbarEduPagedView
+            android:id="@+id/content"
+            android:clipToPadding="false"
             android:layout_width="match_parent"
-            android:layout_height="wrap_content"
-            android:layout_marginHorizontal="16dp"
+            android:layout_height="378dp"
             app:layout_constraintTop_toTopOf="parent"
-            android:gravity="center_horizontal"
-            style="@style/TextHeadline"
-            android:text="@string/taskbar_edu_header_1"
-            android:fontFamily="google-sans"
-            android:textColor="?android:attr/textColorPrimary"
-            android:textSize="24sp"
-            android:maxLines="2"/>
+            launcher:pageIndicator="@+id/content_page_indicator">
 
-        <ImageView
-            android:id="@+id/edu_illustration"
-            android:layout_width="match_parent"
-            android:layout_height="wrap_content"
-            android:layout_marginTop="16dp"
-            app:layout_constraintTop_toBottomOf="@id/edu_header"
-            android:src="@drawable/taskbar_edu_splitscreen"/>
+            <LinearLayout
+                android:id="@+id/page_splitscreen"
+                android:layout_width="match_parent"
+                android:layout_height="wrap_content"
+                android:orientation="vertical"
+                android:gravity="center_horizontal">
+
+                <TextView
+                    android:layout_width="match_parent"
+                    android:layout_height="wrap_content"
+                    android:layout_marginHorizontal="16dp"
+                    android:gravity="center_horizontal"
+                    style="@style/TextHeadline"
+                    android:text="@string/taskbar_edu_splitscreen"
+                    android:fontFamily="google-sans"
+                    android:textColor="?android:attr/textColorPrimary"
+                    android:textSize="24sp"
+                    android:maxLines="2"/>
+
+                <ImageView
+                    android:layout_width="322dp"
+                    android:layout_height="282dp"
+                    android:layout_marginTop="16dp"
+                    android:src="@drawable/taskbar_edu_splitscreen"/>
+            </LinearLayout>
+
+            <LinearLayout
+                android:id="@+id/page_customize"
+                android:layout_width="match_parent"
+                android:layout_height="wrap_content"
+                android:orientation="vertical"
+                android:gravity="center_horizontal">
+
+                <TextView
+                    android:layout_width="match_parent"
+                    android:layout_height="wrap_content"
+                    android:layout_marginHorizontal="16dp"
+                    android:gravity="center_horizontal"
+                    style="@style/TextHeadline"
+                    android:text="@string/taskbar_edu_customize"
+                    android:fontFamily="google-sans"
+                    android:textColor="?android:attr/textColorPrimary"
+                    android:textSize="24sp"
+                    android:maxLines="2"/>
+
+                <ImageView
+                    android:layout_width="322dp"
+                    android:layout_height="282dp"
+                    android:layout_marginTop="16dp"
+                    android:src="@drawable/taskbar_edu_splitscreen"/>
+            </LinearLayout>
+
+            <LinearLayout
+                android:id="@+id/page_dock"
+                android:layout_width="match_parent"
+                android:layout_height="wrap_content"
+                android:orientation="vertical"
+                android:gravity="center_horizontal">
+
+                <TextView
+                    android:layout_width="match_parent"
+                    android:layout_height="wrap_content"
+                    android:layout_marginHorizontal="16dp"
+                    android:gravity="center_horizontal"
+                    style="@style/TextHeadline"
+                    android:text="@string/taskbar_edu_dock"
+                    android:fontFamily="google-sans"
+                    android:textColor="?android:attr/textColorPrimary"
+                    android:textSize="24sp"
+                    android:maxLines="2"/>
+
+                <ImageView
+                    android:layout_width="322dp"
+                    android:layout_height="282dp"
+                    android:layout_marginTop="16dp"
+                    android:src="@drawable/taskbar_edu_splitscreen"/>
+            </LinearLayout>
+        </com.android.launcher3.taskbar.TaskbarEduPagedView>
 
         <Button
-            android:id="@+id/edu_close_button"
+            android:id="@+id/edu_start_button"
             android:layout_width="wrap_content"
             android:layout_height="36dp"
             android:layout_marginBottom="92dp"
-            app:layout_constraintTop_toBottomOf="@id/edu_illustration"
+            app:layout_constraintTop_toBottomOf="@id/content"
             app:layout_constraintBottom_toBottomOf="parent"
             app:layout_constraintStart_toStartOf="parent"
             android:text="@string/taskbar_edu_close"
             style="@style/TaskbarEdu.Button.Close"
             android:textColor="?android:attr/textColorPrimary"/>
 
+        <com.android.launcher3.pageindicators.PageIndicatorDots
+            android:id="@+id/content_page_indicator"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            app:layout_constraintTop_toTopOf="@id/edu_start_button"
+            app:layout_constraintBottom_toBottomOf="@id/edu_start_button"
+            app:layout_constraintStart_toStartOf="parent"
+            app:layout_constraintEnd_toEndOf="parent"
+            android:elevation="1dp" />
+
         <Button
-            android:id="@+id/edu_next_button"
+            android:id="@+id/edu_end_button"
             android:layout_width="wrap_content"
             android:layout_height="36dp"
             android:layout_marginBottom="92dp"
-            app:layout_constraintTop_toBottomOf="@id/edu_illustration"
+            app:layout_constraintTop_toBottomOf="@id/content"
             app:layout_constraintBottom_toBottomOf="parent"
             app:layout_constraintEnd_toEndOf="parent"
             android:text="@string/taskbar_edu_next"
diff --git a/quickstep/res/values/strings.xml b/quickstep/res/values/strings.xml
index da19c95596..2916121c83 100644
--- a/quickstep/res/values/strings.xml
+++ b/quickstep/res/values/strings.xml
@@ -205,11 +205,25 @@
     <string name="gesture_tutorial_action_button_label_skip">Skip</string>
 
     <!-- ******* Taskbar Edu ******* -->
-    <!-- Text in dialog that lets a user know how they can use the taskbar on their device.
-         [CHAR_LIMIT=NONE] -->
-    <string name="taskbar_edu_header_1">Use 2 apps at once and switch apps with the taskbar</string>
-    <!-- Text on button to go to the next screen of a tutorial [CHAR_LIMIT=30] -->
+    <!-- Accessibility text spoken when the taskbar education panel appears [CHAR_LIMIT=NONE] -->
+    <string name="taskbar_edu_opened">Taskbar education appeared</string>
+    <!-- Accessibility text spoken when the taskbar education panel disappears [CHAR_LIMIT=NONE] -->
+    <string name="taskbar_edu_closed">Taskbar education closed</string>
+    <!-- Text in dialog that lets a user know how they can use the taskbar to use multiple apps on their device.
+         [CHAR_LIMIT=60] -->
+    <string name="taskbar_edu_splitscreen" translatable="false">Use 2 apps at once and switch apps with the taskbar</string>
+    <!-- Text in dialog that lets a user know how they can customize the taskbar on their device.
+         [CHAR_LIMIT=60] -->
+    <string name="taskbar_edu_customize" translatable="false">Add your favorite apps and get automatic suggestions</string>
+    <!-- Text in dialog that lets a user know how they can hide the taskbar on their device.
+         [CHAR_LIMIT=60] -->
+    <string name="taskbar_edu_dock">Touch &amp; hold to hide the taskbar anytime</string>
+    <!-- Text on button to go to the next screen of a tutorial [CHAR_LIMIT=16] -->
     <string name="taskbar_edu_next">Next</string>
-    <!-- Text on button to exit a tutorial [CHAR_LIMIT=30] -->
+    <!-- Text on button to go to the previous screen of a tutorial [CHAR_LIMIT=16] -->
+    <string name="taskbar_edu_previous">Back</string>
+    <!-- Text on button to exit a tutorial [CHAR_LIMIT=16] -->
     <string name="taskbar_edu_close">Close</string>
+    <!-- Text on button to finish a tutorial [CHAR_LIMIT=16] -->
+    <string name="taskbar_edu_done">Done</string>
 </resources>
diff --git a/quickstep/res/values/styles.xml b/quickstep/res/values/styles.xml
index c23a3a01a6..ac702794b2 100644
--- a/quickstep/res/values/styles.xml
+++ b/quickstep/res/values/styles.xml
@@ -144,11 +144,13 @@
         <item name="android:background">@drawable/button_taskbar_edu_bordered</item>
         <item name="android:stateListAnimator">@null</item>
         <item name="android:textSize">16sp</item>
+        <item name="android:padding">4dp</item>
     </style>
 
     <style name="TaskbarEdu.Button.Next" parent="@android:style/Widget.Material.Button">
         <item name="android:background">@drawable/button_taskbar_edu_colored</item>
         <item name="android:stateListAnimator">@null</item>
         <item name="android:textSize">16sp</item>
+        <item name="android:padding">4dp</item>
     </style>
 </resources>
\ No newline at end of file
diff --git a/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java b/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
similarity index 65%
rename from quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
rename to quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
index c1b3beb475..f82fbcc15b 100644
--- a/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
+++ b/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
@@ -15,31 +15,24 @@
  */
 package com.android.launcher3.model;
 
-import static android.os.Process.myUserHandle;
-
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
-import static com.android.launcher3.util.WidgetUtils.createAppWidgetProviderInfo;
 
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
-import static org.mockito.Mockito.doReturn;
+import static org.robolectric.Shadows.shadowOf;
 
 import android.app.prediction.AppTarget;
 import android.app.prediction.AppTargetId;
 import android.appwidget.AppWidgetManager;
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ComponentName;
+import android.content.Context;
+import android.os.Process;
 import android.os.UserHandle;
-import android.text.TextUtils;
-import android.util.Log;
-
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
 
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.config.FeatureFlags;
@@ -47,34 +40,38 @@ import com.android.launcher3.icons.ComponentWithLabel;
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
 import com.android.launcher3.model.QuickstepModelDelegate.PredictorState;
+import com.android.launcher3.shadows.ShadowDeviceFlag;
 import com.android.launcher3.util.LauncherModelHelper;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.PendingAddWidgetInfo;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.shadow.api.Shadow;
+import org.robolectric.shadows.ShadowAppWidgetManager;
+import org.robolectric.shadows.ShadowPackageManager;
+import org.robolectric.util.ReflectionHelpers;
 
-import java.util.Arrays;
 import java.util.List;
 import java.util.stream.Collectors;
 
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
 public final class WidgetsPredicationUpdateTaskTest {
 
-    private AppWidgetProviderInfo mApp1Provider1;
-    private AppWidgetProviderInfo mApp1Provider2;
-    private AppWidgetProviderInfo mApp2Provider1;
-    private AppWidgetProviderInfo mApp4Provider1;
-    private AppWidgetProviderInfo mApp4Provider2;
-    private AppWidgetProviderInfo mApp5Provider1;
-    private List<AppWidgetProviderInfo> allWidgets;
+    private AppWidgetProviderInfo mApp1Provider1 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp1Provider2 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp2Provider1 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp4Provider1 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp4Provider2 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp5Provider1 = new AppWidgetProviderInfo();
 
     private FakeBgDataModelCallback mCallback = new FakeBgDataModelCallback();
+    private Context mContext;
     private LauncherModelHelper mModelHelper;
     private UserHandle mUserHandle;
 
@@ -83,56 +80,54 @@ public final class WidgetsPredicationUpdateTaskTest {
 
     @Before
     public void setup() throws Exception {
-        mModelHelper = new LauncherModelHelper();
         MockitoAnnotations.initMocks(this);
         doAnswer(invocation -> {
             ComponentWithLabel componentWithLabel = invocation.getArgument(0);
             return componentWithLabel.getComponent().getShortClassName();
         }).when(mIconCache).getTitleNoCache(any());
 
-        mUserHandle = myUserHandle();
-        mApp1Provider1 = createAppWidgetProviderInfo(
-                ComponentName.createRelative("app1", "provider1"));
-        mApp1Provider2 = createAppWidgetProviderInfo(
-                ComponentName.createRelative("app1", "provider2"));
-        mApp2Provider1 = createAppWidgetProviderInfo(
-                ComponentName.createRelative("app2", "provider1"));
-        mApp4Provider1 = createAppWidgetProviderInfo(
-                ComponentName.createRelative("app4", "provider1"));
-        mApp4Provider2 = createAppWidgetProviderInfo(
-                ComponentName.createRelative("app4", ".provider2"));
-        mApp5Provider1 = createAppWidgetProviderInfo(
-                ComponentName.createRelative("app5", "provider1"));
-        allWidgets = Arrays.asList(mApp1Provider1, mApp1Provider2, mApp2Provider1,
-                mApp4Provider1, mApp4Provider2, mApp5Provider1);
-
-        AppWidgetManager manager = mModelHelper.sandboxContext.spyService(AppWidgetManager.class);
-        doReturn(allWidgets).when(manager).getInstalledProviders();
-        doReturn(allWidgets).when(manager).getInstalledProvidersForProfile(eq(myUserHandle()));
-        doAnswer(i -> {
-            String pkg = i.getArgument(0);
-            Log.e("Hello", "Getting v " + pkg);
-            return TextUtils.isEmpty(pkg) ? allWidgets : allWidgets.stream()
-                    .filter(a -> pkg.equals(a.provider.getPackageName()))
-                    .collect(Collectors.toList());
-        }).when(manager).getInstalledProvidersForPackage(any(), eq(myUserHandle()));
-
+        mContext = RuntimeEnvironment.application;
+        mModelHelper = new LauncherModelHelper();
+        mUserHandle = Process.myUserHandle();
         // 2 widgets, app4/provider1 & app5/provider1, have already been added to the workspace.
-        mModelHelper.initializeData("widgets_predication_update_task_data");
+        mModelHelper.initializeData("/widgets_predication_update_task_data.txt");
+
+        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
+        mApp1Provider1.provider = ComponentName.createRelative("app1", "provider1");
+        ReflectionHelpers.setField(mApp1Provider1, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp1Provider1.provider));
+        mApp1Provider2.provider = ComponentName.createRelative("app1", "provider2");
+        ReflectionHelpers.setField(mApp1Provider2, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp1Provider2.provider));
+        mApp2Provider1.provider = ComponentName.createRelative("app2", "provider1");
+        ReflectionHelpers.setField(mApp2Provider1, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp2Provider1.provider));
+        mApp4Provider1.provider = ComponentName.createRelative("app4", "provider1");
+        ReflectionHelpers.setField(mApp4Provider1, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp4Provider1.provider));
+        mApp4Provider2.provider = ComponentName.createRelative("app4", ".provider2");
+        ReflectionHelpers.setField(mApp4Provider2, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp4Provider2.provider));
+        mApp5Provider1.provider = ComponentName.createRelative("app5", "provider1");
+        ReflectionHelpers.setField(mApp5Provider1, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp5Provider1.provider));
+
+        ShadowAppWidgetManager shadowAppWidgetManager =
+                shadowOf(mContext.getSystemService(AppWidgetManager.class));
+        shadowAppWidgetManager.addInstalledProvider(mApp1Provider1);
+        shadowAppWidgetManager.addInstalledProvider(mApp1Provider2);
+        shadowAppWidgetManager.addInstalledProvider(mApp2Provider1);
+        shadowAppWidgetManager.addInstalledProvider(mApp4Provider1);
+        shadowAppWidgetManager.addInstalledProvider(mApp4Provider2);
+        shadowAppWidgetManager.addInstalledProvider(mApp5Provider1);
+
+        mModelHelper.getModel().addCallbacks(mCallback);
 
-        MAIN_EXECUTOR.submit(() -> mModelHelper.getModel().addCallbacks(mCallback)).get();
         MODEL_EXECUTOR.post(() -> mModelHelper.getBgDataModel().widgetsModel.update(
-                LauncherAppState.getInstance(mModelHelper.sandboxContext),
-                /* packageUser= */ null));
-
-        MODEL_EXECUTOR.submit(() -> { }).get();
-        MAIN_EXECUTOR.submit(() -> { }).get();
+                LauncherAppState.getInstance(mContext), /* packageUser= */ null));
+        waitUntilIdle();
     }
 
-    @After
-    public void tearDown() {
-        mModelHelper.destroy();
-    }
 
     @Test
     public void widgetsRecommendationRan_shouldOnlyReturnNotAddedWidgetsInAppPredictionOrder()
@@ -170,9 +165,9 @@ public final class WidgetsPredicationUpdateTaskTest {
     @Test
     public void widgetsRecommendationRan_localFilterDisabled_shouldReturnWidgetsInPredicationOrder()
             throws Exception {
-        if (FeatureFlags.ENABLE_LOCAL_RECOMMENDED_WIDGETS_FILTER.get()) {
-            return;
-        }
+        ShadowDeviceFlag shadowDeviceFlag = Shadow.extract(
+                FeatureFlags.ENABLE_LOCAL_RECOMMENDED_WIDGETS_FILTER);
+        shadowDeviceFlag.setValue(false);
 
         // WHEN newPredicationTask is executed with 5 predicated widgets.
         AppTarget widget1 = new AppTarget(new AppTargetId("app1"), "app1", "provider1",
@@ -208,8 +203,13 @@ public final class WidgetsPredicationUpdateTaskTest {
         assertThat(actual.getUser()).isEqualTo(expected.getProfile());
     }
 
+    private void waitUntilIdle() {
+        shadowOf(MODEL_EXECUTOR.getLooper()).idle();
+        shadowOf(MAIN_EXECUTOR.getLooper()).idle();
+    }
+
     private WidgetsPredictionUpdateTask newWidgetsPredicationTask(List<AppTarget> appTargets) {
-       return new WidgetsPredictionUpdateTask(
+        return new WidgetsPredictionUpdateTask(
                 new PredictorState(CONTAINER_WIDGETS_PREDICTION, "test_widgets_prediction"),
                 appTargets);
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarEduController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarEduController.java
index c5a0fc1902..ae9592d779 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarEduController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarEduController.java
@@ -32,9 +32,8 @@ public class TaskbarEduController {
         mActivity.getDragLayer().post(() -> {
             mTaskbarEduView = (TaskbarEduView) mActivity.getLayoutInflater().inflate(
                     R.layout.taskbar_edu, mActivity.getDragLayer(), false);
-            mTaskbarEduView.addOnCloseListener(() -> {
-                mTaskbarEduView = null;
-            });
+            mTaskbarEduView.init(new TaskbarEduCallbacks());
+            mTaskbarEduView.addOnCloseListener(() -> mTaskbarEduView = null);
             mTaskbarEduView.show();
         });
     }
@@ -44,4 +43,27 @@ public class TaskbarEduController {
             mTaskbarEduView.close(true /* animate */);
         }
     }
+
+
+    /**
+     * Callbacks for {@link TaskbarEduView} to interact with its controller.
+     */
+    class TaskbarEduCallbacks {
+        void onPageChanged(int currentPage, int pageCount) {
+            if (currentPage == 0) {
+                mTaskbarEduView.updateStartButton(R.string.taskbar_edu_close,
+                        v -> mTaskbarEduView.close(true /* animate */));
+            } else {
+                mTaskbarEduView.updateStartButton(R.string.taskbar_edu_previous,
+                        v -> mTaskbarEduView.snapToPage(currentPage - 1));
+            }
+            if (currentPage == pageCount - 1) {
+                mTaskbarEduView.updateEndButton(R.string.taskbar_edu_done,
+                        v -> mTaskbarEduView.close(true /* animate */));
+            } else {
+                mTaskbarEduView.updateEndButton(R.string.taskbar_edu_next,
+                        v -> mTaskbarEduView.snapToPage(currentPage + 1));
+            }
+        }
+    }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarEduPagedView.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarEduPagedView.java
new file mode 100644
index 0000000000..5efcc4df4c
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarEduPagedView.java
@@ -0,0 +1,79 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.taskbar;
+
+import static com.android.launcher3.AbstractFloatingView.TYPE_ALL;
+import static com.android.launcher3.AbstractFloatingView.TYPE_TASKBAR_EDUCATION_DIALOG;
+
+import android.content.Context;
+import android.util.AttributeSet;
+
+import com.android.launcher3.AbstractFloatingView;
+import com.android.launcher3.PagedView;
+import com.android.launcher3.R;
+import com.android.launcher3.pageindicators.PageIndicatorDots;
+import com.android.launcher3.taskbar.TaskbarEduController.TaskbarEduCallbacks;
+import com.android.launcher3.views.ActivityContext;
+
+/** Horizontal carousel of tutorial screens for Taskbar Edu. */
+public class TaskbarEduPagedView extends PagedView<PageIndicatorDots> {
+
+    private TaskbarEduView mTaskbarEduView;
+    private TaskbarEduCallbacks mControllerCallbacks;
+
+    public TaskbarEduPagedView(Context context, AttributeSet attrs) {
+        super(context, attrs);
+        setImportantForAccessibility(IMPORTANT_FOR_ACCESSIBILITY_YES);
+    }
+
+    void setTaskbarEduView(TaskbarEduView taskbarEduView) {
+        mTaskbarEduView = taskbarEduView;
+        mPageIndicator = taskbarEduView.findViewById(R.id.content_page_indicator);
+        initParentViews(taskbarEduView);
+    }
+
+    void setControllerCallbacks(TaskbarEduCallbacks controllerCallbacks) {
+        mControllerCallbacks = controllerCallbacks;
+        mControllerCallbacks.onPageChanged(getCurrentPage(), getPageCount());
+    }
+
+    @Override
+    protected int getChildGap() {
+        return mTaskbarEduView.getPaddingLeft() + mTaskbarEduView.getPaddingRight();
+    }
+
+    @Override
+    protected void onScrollChanged(int l, int t, int oldl, int oldt) {
+        super.onScrollChanged(l, t, oldl, oldt);
+        if (mMaxScroll > 0) {
+            mPageIndicator.setScroll(l, mMaxScroll);
+        }
+    }
+
+    @Override
+    protected void notifyPageSwitchListener(int prevPage) {
+        super.notifyPageSwitchListener(prevPage);
+        mControllerCallbacks.onPageChanged(getCurrentPage(), getPageCount());
+    }
+
+    @Override
+    protected boolean canScroll(float absVScroll, float absHScroll) {
+        return AbstractFloatingView.getTopOpenViewWithType(
+                ActivityContext.lookupContext(getContext()),
+                TYPE_ALL & ~TYPE_TASKBAR_EDUCATION_DIALOG) == null;
+    }
+}
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarEduView.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarEduView.java
index b5dab7e770..9c4e844fa8 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarEduView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarEduView.java
@@ -21,6 +21,9 @@ import android.animation.PropertyValuesHolder;
 import android.content.Context;
 import android.graphics.Rect;
 import android.util.AttributeSet;
+import android.util.Pair;
+import android.view.View;
+import android.widget.Button;
 
 import com.android.launcher3.Insettable;
 import com.android.launcher3.R;
@@ -34,6 +37,10 @@ public class TaskbarEduView extends AbstractSlideInView<TaskbarActivityContext>
 
     private final Rect mInsets = new Rect();
 
+    private Button mStartButton;
+    private Button mEndButton;
+    private TaskbarEduPagedView mPagedView;
+
     public TaskbarEduView(Context context, AttributeSet attr) {
         this(context, attr, 0);
     }
@@ -43,6 +50,12 @@ public class TaskbarEduView extends AbstractSlideInView<TaskbarActivityContext>
         super(context, attrs, defStyleAttr);
     }
 
+    protected void init(TaskbarEduController.TaskbarEduCallbacks callbacks) {
+        if (mPagedView != null) {
+            mPagedView.setControllerCallbacks(callbacks);
+        }
+    }
+
     @Override
     protected void handleClose(boolean animate) {
         handleClose(animate, DEFAULT_CLOSE_DURATION);
@@ -57,7 +70,10 @@ public class TaskbarEduView extends AbstractSlideInView<TaskbarActivityContext>
     protected void onFinishInflate() {
         super.onFinishInflate();
         mContent = findViewById(R.id.edu_view);
-        findViewById(R.id.edu_close_button).setOnClickListener(v -> close(true /* animate */));
+        mStartButton = findViewById(R.id.edu_start_button);
+        mEndButton = findViewById(R.id.edu_end_button);
+        mPagedView = findViewById(R.id.content);
+        mPagedView.setTaskbarEduView(this);
     }
 
     @Override
@@ -81,6 +97,12 @@ public class TaskbarEduView extends AbstractSlideInView<TaskbarActivityContext>
         animateOpen();
     }
 
+    @Override
+    protected Pair<View, String> getAccessibilityTarget() {
+        return Pair.create(mContent, mIsOpen ? getContext().getString(R.string.taskbar_edu_opened)
+                : getContext().getString(R.string.taskbar_edu_closed));
+    }
+
     @Override
     protected int getScrimColor(Context context) {
         return context.getResources().getColor(R.color.widgets_picker_scrim);
@@ -110,4 +132,18 @@ public class TaskbarEduView extends AbstractSlideInView<TaskbarActivityContext>
         mOpenCloseAnimator.setInterpolator(FAST_OUT_SLOW_IN);
         mOpenCloseAnimator.start();
     }
+
+    void snapToPage(int page) {
+        mPagedView.snapToPage(page);
+    }
+
+    void updateStartButton(int textResId, OnClickListener onClickListener) {
+        mStartButton.setText(textResId);
+        mStartButton.setOnClickListener(onClickListener);
+    }
+
+    void updateEndButton(int textResId, OnClickListener onClickListener) {
+        mEndButton.setText(textResId);
+        mEndButton.setOnClickListener(onClickListener);
+    }
 }
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index fe16b33069..46d58577cd 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -65,6 +65,7 @@ import android.graphics.RectF;
 import android.os.Build;
 import android.os.IBinder;
 import android.os.SystemClock;
+import android.util.Log;
 import android.view.MotionEvent;
 import android.view.View;
 import android.view.View.OnApplyWindowInsetsListener;
@@ -87,6 +88,7 @@ import com.android.launcher3.logging.StatsLogManager;
 import com.android.launcher3.logging.StatsLogManager.StatsLogger;
 import com.android.launcher3.statemanager.BaseState;
 import com.android.launcher3.statemanager.StatefulActivity;
+import com.android.launcher3.testing.TestProtocol;
 import com.android.launcher3.tracing.InputConsumerProto;
 import com.android.launcher3.tracing.SwipeHandlerProto;
 import com.android.launcher3.util.TraceHelper;
@@ -876,6 +878,9 @@ public abstract class AbsSwipeUpHandler<T extends StatefulActivity<S>,
      */
     @UiThread
     public void onGestureEnded(float endVelocity, PointF velocity, PointF downPos) {
+        if (Utilities.IS_RUNNING_IN_TEST_HARNESS) {
+            Log.d(TestProtocol.L3_SWIPE_TO_HOME, "3");
+        }
         float flingThreshold = mContext.getResources()
                 .getDimension(R.dimen.quickstep_fling_threshold_speed);
         boolean isFling = mGestureStarted && !mIsMotionPaused
@@ -1037,6 +1042,9 @@ public abstract class AbsSwipeUpHandler<T extends StatefulActivity<S>,
     @UiThread
     private void handleNormalGestureEnd(float endVelocity, boolean isFling, PointF velocity,
             boolean isCancel) {
+        if (Utilities.IS_RUNNING_IN_TEST_HARNESS) {
+            Log.d(TestProtocol.L3_SWIPE_TO_HOME, "4");
+        }
         long duration = MAX_SWIPE_DURATION;
         float currentShift = mCurrentShift.value;
         final GestureEndTarget endTarget = calculateEndTarget(velocity, endVelocity,
@@ -1155,6 +1163,9 @@ public abstract class AbsSwipeUpHandler<T extends StatefulActivity<S>,
     @UiThread
     private void animateToProgress(float start, float end, long duration, Interpolator interpolator,
             GestureEndTarget target, PointF velocityPxPerMs) {
+        if (Utilities.IS_RUNNING_IN_TEST_HARNESS) {
+            Log.d(TestProtocol.L3_SWIPE_TO_HOME, "5");
+        }
         runOnRecentsAnimationStart(() -> animateToProgressInternal(start, end, duration,
                 interpolator, target, velocityPxPerMs));
     }
@@ -1183,6 +1194,9 @@ public abstract class AbsSwipeUpHandler<T extends StatefulActivity<S>,
     @UiThread
     private void animateToProgressInternal(float start, float end, long duration,
             Interpolator interpolator, GestureEndTarget target, PointF velocityPxPerMs) {
+        if (Utilities.IS_RUNNING_IN_TEST_HARNESS) {
+            Log.d(TestProtocol.L3_SWIPE_TO_HOME, "7");
+        }
         maybeUpdateRecentsAttachedState();
 
         // If we are transitioning to launcher, then listen for the activity to be restarted while
@@ -1812,6 +1826,9 @@ public abstract class AbsSwipeUpHandler<T extends StatefulActivity<S>,
      * be run when it is next started.
      */
     protected void runOnRecentsAnimationStart(Runnable action) {
+        if (Utilities.IS_RUNNING_IN_TEST_HARNESS) {
+            Log.d(TestProtocol.L3_SWIPE_TO_HOME, "6");
+        }
         if (mRecentsAnimationTargets == null) {
             mRecentsAnimationStartCallbacks.add(action);
         } else {
diff --git a/quickstep/src/com/android/quickstep/ImageActionsApi.java b/quickstep/src/com/android/quickstep/ImageActionsApi.java
index 8cb64c24b8..154848d224 100644
--- a/quickstep/src/com/android/quickstep/ImageActionsApi.java
+++ b/quickstep/src/com/android/quickstep/ImageActionsApi.java
@@ -64,20 +64,23 @@ public class ImageActionsApi {
      */
     @UiThread
     public void shareWithExplicitIntent(@Nullable Rect crop, Intent intent) {
-        addImageAndSendIntent(crop, intent, false);
+        addImageAndSendIntent(crop, intent, false, null /* exceptionCallback */);
     }
 
     /**
      * Share the image this api was constructed with using the provided intent. The implementation
      * should set the intent's data field to the URI pointing to the image.
+     * @param exceptionCallback An optional callback to be called when the intent can't be resolved
      */
     @UiThread
-    public void shareAsDataWithExplicitIntent(@Nullable Rect crop, Intent intent) {
-        addImageAndSendIntent(crop, intent, true);
+    public void shareAsDataWithExplicitIntent(@Nullable Rect crop, Intent intent,
+            @Nullable Runnable exceptionCallback) {
+        addImageAndSendIntent(crop, intent, true, exceptionCallback);
     }
 
     @UiThread
-    private void addImageAndSendIntent(@Nullable Rect crop, Intent intent, boolean setData) {
+    private void addImageAndSendIntent(@Nullable Rect crop, Intent intent, boolean setData,
+            @Nullable Runnable exceptionCallback) {
         if (mBitmapSupplier.get() == null) {
             Log.e(TAG, "No snapshot available, not starting share.");
             return;
@@ -92,7 +95,7 @@ public class ImageActionsApi {
                         intentForUri.putExtra(EXTRA_STREAM, uri);
                     }
                     return new Intent[]{intentForUri};
-                }, TAG));
+                }, TAG, exceptionCallback));
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java b/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java
index 35a851ab65..81e69176c2 100644
--- a/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java
+++ b/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java
@@ -462,7 +462,7 @@ class OrientationTouchTransformer {
                             + "mRotation: " + mRotation
                             + " this: " + this);
                 }
-                event.transform(mTmpMatrix);
+                event.applyTransform(mTmpMatrix);
                 return true;
             }
             mTmpPoint[0] = event.getX();
@@ -478,7 +478,7 @@ class OrientationTouchTransformer {
             }
 
             if (contains(mTmpPoint[0], mTmpPoint[1])) {
-                event.transform(mTmpMatrix);
+                event.applyTransform(mTmpMatrix);
                 return true;
             }
             return false;
diff --git a/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java b/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java
index ec51599fbf..1f57e9962a 100644
--- a/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java
+++ b/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java
@@ -278,6 +278,8 @@ public abstract class SwipeUpAnimationLogic implements
         RemoteAnimationTargetCompat primaryTaskTarget;
         RemoteAnimationTargetCompat secondaryTaskTarget;
 
+        // TODO(b/197568823) Determine if we need to exclude assistant as one of the targets we
+        //  animate
         if (!mIsSwipeForStagedSplit) {
             primaryTaskTarget = targets.findTask(mGestureState.getRunningTaskId());
             mRemoteTargetHandles[0].mTransformParams.setTargetSet(targets);
diff --git a/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java b/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
index 725c7c45a5..e6285f2902 100644
--- a/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
+++ b/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
@@ -50,6 +50,7 @@ import android.view.ViewConfiguration;
 import androidx.annotation.UiThread;
 
 import com.android.launcher3.R;
+import com.android.launcher3.Utilities;
 import com.android.launcher3.testing.TestLogging;
 import com.android.launcher3.testing.TestProtocol;
 import com.android.launcher3.tracing.InputConsumerProto;
@@ -356,6 +357,9 @@ public class OtherActivityInputConsumer extends ContextWrapper implements InputC
             }
             case ACTION_CANCEL:
             case ACTION_UP: {
+                if (Utilities.IS_RUNNING_IN_TEST_HARNESS) {
+                    Log.d(TestProtocol.L3_SWIPE_TO_HOME, "1");
+                }
                 if (DEBUG_FAILED_QUICKSWITCH && !mPassedWindowMoveSlop) {
                     float displacementX = mLastPos.x - mDownPos.x;
                     float displacementY = mLastPos.y - mDownPos.y;
@@ -409,6 +413,9 @@ public class OtherActivityInputConsumer extends ContextWrapper implements InputC
      * the animation can still be running.
      */
     private void finishTouchTracking(MotionEvent ev) {
+        if (Utilities.IS_RUNNING_IN_TEST_HARNESS) {
+            Log.d(TestProtocol.L3_SWIPE_TO_HOME, "2");
+        }
         Object traceToken = TraceHelper.INSTANCE.beginSection(UP_EVT,
                 FLAG_CHECK_FOR_RACE_CONDITIONS);
 
diff --git a/quickstep/src/com/android/quickstep/util/ImageActionUtils.java b/quickstep/src/com/android/quickstep/util/ImageActionUtils.java
index de7dbd64f5..51a9915523 100644
--- a/quickstep/src/com/android/quickstep/util/ImageActionUtils.java
+++ b/quickstep/src/com/android/quickstep/util/ImageActionUtils.java
@@ -154,6 +154,18 @@ public class ImageActionUtils {
     @WorkerThread
     public static void persistBitmapAndStartActivity(Context context, Bitmap bitmap, Rect crop,
             Intent intent, BiFunction<Uri, Intent, Intent[]> uriToIntentMap, String tag) {
+        persistBitmapAndStartActivity(context, bitmap, crop, intent, uriToIntentMap, tag,
+                (Runnable) null);
+    }
+
+    /**
+     * Starts activity based on given intent created from image uri.
+     * @param exceptionCallback An optional callback to be called when the intent can't be resolved
+     */
+    @WorkerThread
+    public static void persistBitmapAndStartActivity(Context context, Bitmap bitmap, Rect crop,
+            Intent intent, BiFunction<Uri, Intent, Intent[]> uriToIntentMap, String tag,
+            Runnable exceptionCallback) {
         Intent[] intents = uriToIntentMap.apply(getImageUri(bitmap, crop, context, tag), intent);
 
         try {
@@ -165,6 +177,9 @@ public class ImageActionUtils {
             }
         } catch (ActivityNotFoundException e) {
             Log.e(TAG, "No activity found to receive image intent");
+            if (exceptionCallback != null) {
+                exceptionCallback.run();
+            }
         }
     }
 
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index abeadfd29b..2a19ca2830 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -3314,9 +3314,9 @@ public abstract class RecentsView<ACTIVITY_TYPE extends StatefulActivity<STATE_T
 
     @Override
     protected void onLayout(boolean changed, int left, int top, int right, int bottom) {
-        // If we're going to HOME, avoid unnecessary onLayout that cause TaskViews to re-arrange
-        // during animation to HOME.
-        if (mCurrentGestureEndTarget == GestureState.GestureEndTarget.HOME) {
+        // If we're going to a state without overview panel, avoid unnecessary onLayout that
+        // cause TaskViews to re-arrange during animation to that state.
+        if (!mOverviewStateEnabled && !mFirstLayout) {
             return;
         }
 
@@ -4019,12 +4019,14 @@ public abstract class RecentsView<ACTIVITY_TYPE extends StatefulActivity<STATE_T
             recentsAnimationTargets.addReleaseCheck(mSyncTransactionApplier);
         }
 
-        // TODO Consolidate this shared code with SwipeUpAnimationLogic (or mabe just reuse
-        //  what that class has and pass it into here
-        mRemoteTargetHandles = new RemoteTargetHandle[recentsAnimationTargets.apps.length];
+        RemoteAnimationTargetCompat dividerTarget =
+                recentsAnimationTargets.getNonAppTargetOfType(TYPE_DOCK_DIVIDER);
+        // TODO Consolidate this shared code with SwipeUpAnimationLogic (or maybe just reuse
+        //  what that class has and pass it into here)
+        mRemoteTargetHandles = new RemoteTargetHandle[dividerTarget == null ? 1 : 2];
         TaskViewSimulator primaryTvs = createTaskViewSimulator();
         mRemoteTargetHandles[0] = new RemoteTargetHandle(primaryTvs, new TransformParams());
-        if (recentsAnimationTargets.apps.length == 1) {
+        if (dividerTarget == null) {
             mRemoteTargetHandles[0].mTaskViewSimulator
                     .setPreview(recentsAnimationTargets.apps[0], null);
             mRemoteTargetHandles[0].mTransformParams.setTargetSet(recentsAnimationTargets);
@@ -4034,8 +4036,6 @@ public abstract class RecentsView<ACTIVITY_TYPE extends StatefulActivity<STATE_T
             secondaryTvs.recentsViewScale.value = 1;
 
             mRemoteTargetHandles[1] = new RemoteTargetHandle(secondaryTvs, new TransformParams());
-            RemoteAnimationTargetCompat dividerTarget =
-                    recentsAnimationTargets.getNonAppTargetOfType(TYPE_DOCK_DIVIDER);
             RemoteAnimationTargetCompat primaryTaskTarget = recentsAnimationTargets.apps[0];
             RemoteAnimationTargetCompat secondaryTaskTarget = recentsAnimationTargets.apps[1];
             SplitConfigurationOptions.StagedSplitBounds
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index f84a05ce3d..5986649808 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -674,7 +674,6 @@ public class TaskView extends FrameLayout implements Reusable {
             TestLogging.recordEvent(
                     TestProtocol.SEQUENCE_MAIN, "startActivityFromRecentsAsync", mTask);
             ActivityOptionsWrapper opts =  mActivity.getActivityLaunchOptions(this, null);
-            opts.options.setLaunchDisplayId(getRootViewDisplayId());
             if (ActivityManagerWrapper.getInstance()
                     .startActivityFromRecents(mTask.key, opts.options)) {
                 RecentsView recentsView = getRecentsView();
@@ -715,7 +714,6 @@ public class TaskView extends FrameLayout implements Reusable {
             // Indicate success once the system has indicated that the transition has started
             ActivityOptions opts = ActivityOptionsCompat.makeCustomAnimation(
                     getContext(), 0, 0, () -> callback.accept(true), MAIN_EXECUTOR.getHandler());
-            opts.setLaunchDisplayId(getRootViewDisplayId());
             if (freezeTaskList) {
                 ActivityOptionsCompat.setFreezeRecentTasksList(opts);
             }
@@ -1528,10 +1526,6 @@ public class TaskView extends FrameLayout implements Reusable {
         mDigitalWellBeingToast.setBannerColorTint(tintColor, amount);
     }
 
-    private int getRootViewDisplayId() {
-        return getRootView().getDisplay().getDisplayId();
-    }
-
     /**
      * We update and subsequently draw these in {@link #setFullscreenProgress(float)}.
      */
diff --git a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
index bc8602c940..416c193cfc 100644
--- a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
+++ b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
@@ -30,7 +30,6 @@ import android.content.Context;
 import android.content.pm.PackageManager;
 import android.util.Log;
 
-import androidx.test.uiautomator.By;
 import androidx.test.uiautomator.UiDevice;
 
 import com.android.launcher3.tapl.LauncherInstrumentation;
@@ -217,12 +216,7 @@ public class NavigationModeSwitchRule implements TestRule {
 
     private static void assertTrue(LauncherInstrumentation launcher, String message,
             boolean condition, Description description) {
-        if (launcher.getDevice().hasObject(By.textStartsWith(""))) {
-            // The condition above is "screen is not empty". We are not treating
-            // "Screen is empty" as an anomaly here. It's an acceptable state when
-            // Launcher just starts under instrumentation.
-            launcher.checkForAnomaly();
-        }
+        launcher.checkForAnomaly(true);
         if (!condition) {
             final AssertionError assertionError = new AssertionError(message);
             if (description != null) {
diff --git a/quickstep/robolectric_tests/src/com/android/quickstep/OrientationTouchTransformerTest.java b/quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
similarity index 98%
rename from quickstep/robolectric_tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
rename to quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
index 2b1b57c450..159a51f733 100644
--- a/quickstep/robolectric_tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
+++ b/quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
@@ -19,6 +19,8 @@ package com.android.quickstep;
 
 import static android.view.Display.DEFAULT_DISPLAY;
 
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
 import static com.android.quickstep.SysUINavigationMode.Mode.NO_BUTTON;
 
 import static org.junit.Assert.assertFalse;
@@ -40,6 +42,9 @@ import android.view.Display;
 import android.view.MotionEvent;
 import android.view.Surface;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
 import com.android.launcher3.ResourceUtils;
 import com.android.launcher3.util.DisplayController;
 
@@ -47,10 +52,9 @@ import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.MockitoAnnotations;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
 
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class OrientationTouchTransformerTest {
     static class ScreenSize {
         int mHeight;
@@ -293,7 +297,7 @@ public class OrientationTouchTransformerTest {
     }
 
     private DisplayController.Info createDisplayInfo(ScreenSize screenSize, int rotation) {
-        Context context = RuntimeEnvironment.application;
+        Context context = getApplicationContext();
         Display display = spy(context.getSystemService(DisplayManager.class)
                 .getDisplay(DEFAULT_DISPLAY));
 
diff --git a/tests/res/raw/cache_data_updated_task_data.txt b/robolectric_tests/resources/cache_data_updated_task_data.txt
similarity index 100%
rename from tests/res/raw/cache_data_updated_task_data.txt
rename to robolectric_tests/resources/cache_data_updated_task_data.txt
diff --git a/tests/res/raw/db_schema_v10.json b/robolectric_tests/resources/db_schema_v10.json
similarity index 100%
rename from tests/res/raw/db_schema_v10.json
rename to robolectric_tests/resources/db_schema_v10.json
diff --git a/tests/res/raw/package_install_state_change_task_data.txt b/robolectric_tests/resources/package_install_state_change_task_data.txt
similarity index 100%
rename from tests/res/raw/package_install_state_change_task_data.txt
rename to robolectric_tests/resources/package_install_state_change_task_data.txt
diff --git a/tests/res/raw/widgets_predication_update_task_data.txt b/robolectric_tests/resources/widgets_predication_update_task_data.txt
similarity index 100%
rename from tests/res/raw/widgets_predication_update_task_data.txt
rename to robolectric_tests/resources/widgets_predication_update_task_data.txt
diff --git a/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java b/robolectric_tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
similarity index 93%
rename from tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
rename to robolectric_tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
index 8a4590a388..8aa6f374aa 100644
--- a/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
@@ -13,9 +13,6 @@ import android.content.Intent;
 import android.graphics.Rect;
 import android.util.Pair;
 
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherSettings;
@@ -24,17 +21,19 @@ import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.util.ContentWriter;
-import com.android.launcher3.util.Executors;
 import com.android.launcher3.util.GridOccupancy;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSparseArrayMap;
 import com.android.launcher3.util.LauncherModelHelper;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.annotation.LooperMode.Mode;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -42,8 +41,8 @@ import java.util.List;
 /**
  * Tests for {@link AddWorkspaceItemsTask}
  */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
+@LooperMode(Mode.PAUSED)
 public class AddWorkspaceItemsTaskTest {
 
     private final ComponentName mComponent1 = new ComponentName("a", "b");
@@ -61,7 +60,7 @@ public class AddWorkspaceItemsTaskTest {
     @Before
     public void setup() {
         mModelHelper = new LauncherModelHelper();
-        mTargetContext = mModelHelper.sandboxContext;
+        mTargetContext = RuntimeEnvironment.application;
         mIdp = InvariantDeviceProfile.INSTANCE.get(mTargetContext);
         mIdp.numColumns = mIdp.numRows = 5;
         mAppState = LauncherAppState.getInstance(mTargetContext);
@@ -71,11 +70,6 @@ public class AddWorkspaceItemsTaskTest {
         mNewScreens = new IntArray();
     }
 
-    @After
-    public void tearDown() {
-        mModelHelper.destroy();
-    }
-
     private AddWorkspaceItemsTask newTask(ItemInfo... items) {
         List<Pair<ItemInfo, Object>> list = new ArrayList<>();
         for (ItemInfo item : items) {
@@ -94,7 +88,7 @@ public class AddWorkspaceItemsTaskTest {
 
         int[] spaceFound = newTask().findSpaceForItem(
                 mAppState, mModelHelper.getBgDataModel(), mExistingScreens, mNewScreens, 1, 1);
-        assertEquals(1, spaceFound[0]);
+        assertEquals(2, spaceFound[0]);
         assertTrue(mScreenOccupancy.get(spaceFound[0])
                 .isRegionVacant(spaceFound[1], spaceFound[2], 1, 1));
 
@@ -133,7 +127,7 @@ public class AddWorkspaceItemsTaskTest {
     @Test
     public void testAddItem_some_items_added() throws Exception {
         Callbacks callbacks = mock(Callbacks.class);
-        Executors.MAIN_EXECUTOR.submit(() -> mModelHelper.getModel().addCallbacks(callbacks)).get();
+        mModelHelper.getModel().addCallbacks(callbacks);
 
         WorkspaceItemInfo info = new WorkspaceItemInfo();
         info.intent = new Intent().setComponent(mComponent1);
diff --git a/tests/src/com/android/launcher3/model/BackupRestoreTest.java b/robolectric_tests/src/com/android/launcher3/model/BackupRestoreTest.java
similarity index 62%
rename from tests/src/com/android/launcher3/model/BackupRestoreTest.java
rename to robolectric_tests/src/com/android/launcher3/model/BackupRestoreTest.java
index 41914de1ab..a397db565d 100644
--- a/tests/src/com/android/launcher3/model/BackupRestoreTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/BackupRestoreTest.java
@@ -17,7 +17,6 @@
 package com.android.launcher3.model;
 
 import static android.content.pm.PackageManager.INSTALL_REASON_DEVICE_RESTORE;
-import static android.os.Process.myUserHandle;
 
 import static com.android.launcher3.LauncherSettings.Favorites.BACKUP_TABLE_NAME;
 import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
@@ -27,110 +26,74 @@ import static com.android.launcher3.provider.LauncherDbUtils.tableExists;
 import static com.android.launcher3.util.LauncherModelHelper.APP_ICON;
 import static com.android.launcher3.util.LauncherModelHelper.NO__ICON;
 import static com.android.launcher3.util.LauncherModelHelper.SHORTCUT;
-import static com.android.launcher3.util.ReflectionHelpers.getField;
-import static com.android.launcher3.util.ReflectionHelpers.setField;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.doReturn;
-import static org.mockito.Mockito.spy;
+import static org.robolectric.util.ReflectionHelpers.setField;
 
 import android.app.backup.BackupManager;
 import android.content.pm.PackageInstaller;
 import android.database.Cursor;
 import android.database.sqlite.SQLiteDatabase;
+import android.os.Process;
 import android.os.UserHandle;
-import android.util.ArrayMap;
-import android.util.LongSparseArray;
-
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
+import android.os.UserManager;
 
 import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.provider.RestoreDbTask;
+import com.android.launcher3.shadows.LShadowBackupManager;
 import com.android.launcher3.util.LauncherModelHelper;
-import com.android.launcher3.util.SafeCloseable;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.shadow.api.Shadow;
+import org.robolectric.shadows.ShadowUserManager;
 
 /**
  * Tests to verify backup and restore flow.
  */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
+@LooperMode(LooperMode.Mode.PAUSED)
 public class BackupRestoreTest {
 
-    private static final int PER_USER_RANGE = 200000;
-
-
-    private long mCurrentMyProfileId;
-    private long mOldMyProfileId;
-
-    private long mCurrentWorkProfileId;
-    private long mOldWorkProfileId;
+    private static final long MY_OLD_PROFILE_ID = 1;
+    private static final long MY_PROFILE_ID = 0;
+    private static final long OLD_WORK_PROFILE_ID = 11;
+    private static final int WORK_PROFILE_ID = 10;
 
+    private ShadowUserManager mUserManager;
     private BackupManager mBackupManager;
     private LauncherModelHelper mModelHelper;
     private SQLiteDatabase mDb;
     private InvariantDeviceProfile mIdp;
 
-    private UserHandle mWorkUserHandle;
-
-    private SafeCloseable mUserChangeListener;
-
     @Before
     public void setUp() {
-        mModelHelper = new LauncherModelHelper();
-
-        mCurrentMyProfileId = mModelHelper.defaultProfileId;
-        mOldMyProfileId = mCurrentMyProfileId + 1;
-        mCurrentWorkProfileId = mOldMyProfileId + 1;
-        mOldWorkProfileId = mCurrentWorkProfileId + 1;
-
-        mWorkUserHandle = UserHandle.getUserHandleForUid(PER_USER_RANGE);
-        mUserChangeListener = UserCache.INSTANCE.get(mModelHelper.sandboxContext)
-                .addUserChangeListener(() -> { });
-
         setupUserManager();
         setupBackupManager();
-        RestoreDbTask.setPending(mModelHelper.sandboxContext);
+        mModelHelper = new LauncherModelHelper();
+        RestoreDbTask.setPending(RuntimeEnvironment.application);
         mDb = mModelHelper.provider.getDb();
-        mIdp = InvariantDeviceProfile.INSTANCE.get(mModelHelper.sandboxContext);
-
-    }
-
-    @After
-    public void tearDown() {
-        mUserChangeListener.close();
-        mModelHelper.destroy();
+        mIdp = InvariantDeviceProfile.INSTANCE.get(RuntimeEnvironment.application);
     }
 
     private void setupUserManager() {
-        UserCache cache = UserCache.INSTANCE.get(mModelHelper.sandboxContext);
-        synchronized (cache) {
-            LongSparseArray<UserHandle> users = getField(cache, "mUsers");
-            users.clear();
-            users.put(mCurrentMyProfileId, myUserHandle());
-            users.put(mCurrentWorkProfileId, mWorkUserHandle);
-
-            ArrayMap<UserHandle, Long> userMap = getField(cache, "mUserToSerialMap");
-            userMap.clear();
-            userMap.put(myUserHandle(), mCurrentMyProfileId);
-            userMap.put(mWorkUserHandle, mCurrentWorkProfileId);
-        }
+        final UserManager userManager = RuntimeEnvironment.application.getSystemService(
+                UserManager.class);
+        mUserManager = Shadow.extract(userManager);
+        // sign in to work profile
+        mUserManager.addUser(WORK_PROFILE_ID, "work", ShadowUserManager.FLAG_MANAGED_PROFILE);
     }
 
     private void setupBackupManager() {
-        mBackupManager = spy(new BackupManager(mModelHelper.sandboxContext));
-        doReturn(myUserHandle()).when(mBackupManager)
-                .getUserForAncestralSerialNumber(eq(mOldMyProfileId));
-        doReturn(mWorkUserHandle).when(mBackupManager)
-                .getUserForAncestralSerialNumber(eq(mOldWorkProfileId));
+        mBackupManager = new BackupManager(RuntimeEnvironment.application);
+        final LShadowBackupManager bm = Shadow.extract(mBackupManager);
+        bm.addProfile(MY_OLD_PROFILE_ID, Process.myUserHandle());
+        bm.addProfile(OLD_WORK_PROFILE_ID, UserHandle.of(WORK_PROFILE_ID));
     }
 
     @Test
@@ -155,18 +118,18 @@ public class BackupRestoreTest {
                 { SHORTCUT, SHORTCUT, NO__ICON, NO__ICON},
                 { NO__ICON, NO__ICON, SHORTCUT, SHORTCUT},
                 { APP_ICON, SHORTCUT, SHORTCUT, APP_ICON},
-            }}, 1, mOldMyProfileId);
+            }}, 1, MY_OLD_PROFILE_ID);
         // setup grid for work profile on second screen
         mModelHelper.createGrid(new int[][][]{{
                 { NO__ICON, APP_ICON, SHORTCUT, SHORTCUT},
                 { SHORTCUT, SHORTCUT, NO__ICON, NO__ICON},
                 { NO__ICON, NO__ICON, SHORTCUT, SHORTCUT},
                 { APP_ICON, SHORTCUT, SHORTCUT, NO__ICON},
-            }}, 2, mOldWorkProfileId);
+            }}, 2, OLD_WORK_PROFILE_ID);
         // simulates the creation of backup upon restore
-        new GridBackupTable(mModelHelper.sandboxContext, mDb, mIdp.numDatabaseHotseatIcons,
+        new GridBackupTable(RuntimeEnvironment.application, mDb, mIdp.numDatabaseHotseatIcons,
                 mIdp.numColumns, mIdp.numRows).doBackup(
-                mOldMyProfileId, GridBackupTable.OPTION_REQUIRES_SANITIZATION);
+                        MY_OLD_PROFILE_ID, GridBackupTable.OPTION_REQUIRES_SANITIZATION);
         // reset favorites table
         createTableUsingOldProfileId();
     }
@@ -178,28 +141,28 @@ public class BackupRestoreTest {
     private void verifyTableIsFilled(String tableName, boolean sanitized) {
         assertEquals(sanitized ? 12 : 13, getCount(mDb,
                 "SELECT * FROM " + tableName + " WHERE profileId = "
-                        + (sanitized ? mCurrentMyProfileId : mOldMyProfileId)));
+                        + (sanitized ? MY_PROFILE_ID : MY_OLD_PROFILE_ID)));
         assertEquals(10, getCount(mDb, "SELECT * FROM " + tableName + " WHERE profileId = "
-                + (sanitized ? mCurrentWorkProfileId : mOldWorkProfileId)));
+                + (sanitized ? WORK_PROFILE_ID : OLD_WORK_PROFILE_ID)));
     }
 
     private void createTableUsingOldProfileId() {
         // simulates the creation of favorites table on old device
         dropTable(mDb, TABLE_NAME);
-        addTableToDb(mDb, mOldMyProfileId, false);
+        addTableToDb(mDb, MY_OLD_PROFILE_ID, false);
     }
 
     private void createRestoreSession() throws Exception {
         final PackageInstaller.SessionParams params = new PackageInstaller.SessionParams(
                 PackageInstaller.SessionParams.MODE_FULL_INSTALL);
-        final PackageInstaller installer = mModelHelper.sandboxContext.getPackageManager()
+        final PackageInstaller installer = RuntimeEnvironment.application.getPackageManager()
                 .getPackageInstaller();
         final int sessionId = installer.createSession(params);
         final PackageInstaller.SessionInfo info = installer.getSessionInfo(sessionId);
         setField(info, "installReason", INSTALL_REASON_DEVICE_RESTORE);
         // TODO: (b/148410677) we should verify the following call instead
         //  InstallSessionHelper.INSTANCE.get(getContext()).restoreDbIfApplicable(info);
-        RestoreDbTask.restoreIfPossible(mModelHelper.sandboxContext,
+        RestoreDbTask.restoreIfPossible(RuntimeEnvironment.application,
                 mModelHelper.provider.getHelper(), mBackupManager);
     }
 
diff --git a/tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java b/robolectric_tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
similarity index 92%
rename from tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
rename to robolectric_tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
index dba0a4063f..9ac3fe761e 100644
--- a/tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
@@ -16,8 +16,6 @@ import android.os.UserHandle;
 import android.os.UserManager;
 
 import androidx.annotation.NonNull;
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
 
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.icons.BitmapInfo;
@@ -28,10 +26,13 @@ import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.util.LauncherModelHelper;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.annotation.LooperMode.Mode;
 
 import java.util.Arrays;
 import java.util.HashSet;
@@ -39,8 +40,8 @@ import java.util.HashSet;
 /**
  * Tests for {@link CacheDataUpdatedTask}
  */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
+@LooperMode(Mode.PAUSED)
 public class CacheDataUpdatedTaskTest {
 
     private static final String NEW_LABEL_PREFIX = "new-label-";
@@ -50,10 +51,10 @@ public class CacheDataUpdatedTaskTest {
     @Before
     public void setup() throws Exception {
         mModelHelper = new LauncherModelHelper();
-        mModelHelper.initializeData("cache_data_updated_task_data");
+        mModelHelper.initializeData("/cache_data_updated_task_data.txt");
 
         // Add placeholder entries in the cache to simulate update
-        Context context = mModelHelper.sandboxContext;
+        Context context = RuntimeEnvironment.application;
         IconCache iconCache = LauncherAppState.getInstance(context).getIconCache();
         CachingLogic<ItemInfo> placeholderLogic = new CachingLogic<ItemInfo>() {
             @Override
@@ -85,11 +86,6 @@ public class CacheDataUpdatedTaskTest {
         }
     }
 
-    @After
-    public void tearDown() {
-        mModelHelper.destroy();
-    }
-
     private CacheDataUpdatedTask newTask(int op, String... pkg) {
         return new CacheDataUpdatedTask(op, Process.myUserHandle(),
                 new HashSet<>(Arrays.asList(pkg)));
diff --git a/tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java b/robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
similarity index 91%
rename from tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
rename to robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
index d849c8fff9..be03c7dd4b 100644
--- a/tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
@@ -15,13 +15,12 @@
  */
 package com.android.launcher3.model;
 
-import static androidx.test.InstrumentationRegistry.getContext;
-
 import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertFalse;
 import static junit.framework.Assert.assertNotSame;
 import static junit.framework.Assert.assertTrue;
 
+import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.when;
@@ -33,10 +32,6 @@ import android.database.Cursor;
 import android.database.sqlite.SQLiteDatabase;
 import android.database.sqlite.SQLiteOpenHelper;
 
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-import androidx.test.platform.app.InstrumentationRegistry;
-
 import com.android.launcher3.LauncherProvider;
 import com.android.launcher3.LauncherProvider.DatabaseHelper;
 import com.android.launcher3.LauncherSettings.Favorites;
@@ -45,14 +40,15 @@ import com.android.launcher3.R;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
 
 import java.io.File;
 
 /**
  * Tests for {@link DbDowngradeHelper}
  */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
 public class DbDowngradeHelperTest {
 
     private static final String SCHEMA_FILE = "test_schema.json";
@@ -64,7 +60,7 @@ public class DbDowngradeHelperTest {
 
     @Before
     public void setup() {
-        mContext = InstrumentationRegistry.getInstrumentation().getTargetContext();
+        mContext = RuntimeEnvironment.application;
         mSchemaFile = mContext.getFileStreamPath(SCHEMA_FILE);
         mDbFile = mContext.getDatabasePath(DB_FILE);
     }
@@ -81,10 +77,8 @@ public class DbDowngradeHelperTest {
     public void testUpdateSchemaFile() throws Exception {
         // Setup mock resources
         Resources res = spy(mContext.getResources());
-        Resources myRes = getContext().getResources();
-        doAnswer(i -> myRes.openRawResource(
-                myRes.getIdentifier("db_schema_v10", "raw", getContext().getPackageName())))
-                .when(res).openRawResource(R.raw.downgrade_schema);
+        doAnswer(i ->this.getClass().getResourceAsStream("/db_schema_v10.json"))
+                .when(res).openRawResource(eq(R.raw.downgrade_schema));
         Context context = spy(mContext);
         when(context.getResources()).thenReturn(res);
 
diff --git a/tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java b/robolectric_tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
similarity index 77%
rename from tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
rename to robolectric_tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
index 004ed06b32..655237d08c 100644
--- a/tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
@@ -16,18 +16,18 @@
 
 package com.android.launcher3.model;
 
-import static com.android.launcher3.util.LauncherModelHelper.TEST_ACTIVITY;
 import static com.android.launcher3.util.LauncherModelHelper.TEST_PACKAGE;
 
 import static org.junit.Assert.assertEquals;
+import static org.robolectric.Shadows.shadowOf;
+import static org.robolectric.util.ReflectionHelpers.setField;
 
+import android.content.ComponentName;
 import android.content.Context;
 import android.content.pm.PackageInstaller;
+import android.content.pm.PackageInstaller.SessionInfo;
 import android.content.pm.PackageInstaller.SessionParams;
 
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.icons.BitmapInfo;
 import com.android.launcher3.model.data.FolderInfo;
@@ -35,16 +35,19 @@ import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.util.LauncherLayoutBuilder;
 import com.android.launcher3.util.LauncherModelHelper;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.annotation.LooperMode.Mode;
 
 /**
  * Tests for layout parser for remote layout
  */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
+@LooperMode(Mode.PAUSED)
 public class DefaultLayoutProviderTest {
 
     private LauncherModelHelper mModelHelper;
@@ -53,18 +56,16 @@ public class DefaultLayoutProviderTest {
     @Before
     public void setUp() {
         mModelHelper = new LauncherModelHelper();
-        mTargetContext = mModelHelper.sandboxContext;
-    }
+        mTargetContext = RuntimeEnvironment.application;
 
-    @After
-    public void tearDown() {
-        mModelHelper.destroy();
+        shadowOf(mTargetContext.getPackageManager())
+                .addActivityIfNotPresent(new ComponentName(TEST_PACKAGE, TEST_PACKAGE));
     }
 
     @Test
     public void testCustomProfileLoaded_with_icon_on_hotseat() throws Exception {
         writeLayoutAndLoad(new LauncherLayoutBuilder().atHotseat(0)
-                .putApp(TEST_PACKAGE, TEST_ACTIVITY));
+                .putApp(TEST_PACKAGE, TEST_PACKAGE));
 
         // Verify one item in hotseat
         assertEquals(1, mModelHelper.getBgDataModel().workspaceItems.size());
@@ -76,9 +77,9 @@ public class DefaultLayoutProviderTest {
     @Test
     public void testCustomProfileLoaded_with_folder() throws Exception {
         writeLayoutAndLoad(new LauncherLayoutBuilder().atHotseat(0).putFolder(android.R.string.copy)
-                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
-                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
-                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
+                .addApp(TEST_PACKAGE, TEST_PACKAGE)
+                .addApp(TEST_PACKAGE, TEST_PACKAGE)
+                .addApp(TEST_PACKAGE, TEST_PACKAGE)
                 .build());
 
         // Verify folder
@@ -91,9 +92,9 @@ public class DefaultLayoutProviderTest {
     @Test
     public void testCustomProfileLoaded_with_folder_custom_title() throws Exception {
         writeLayoutAndLoad(new LauncherLayoutBuilder().atHotseat(0).putFolder("CustomFolder")
-                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
-                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
-                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
+                .addApp(TEST_PACKAGE, TEST_PACKAGE)
+                .addApp(TEST_PACKAGE, TEST_PACKAGE)
+                .addApp(TEST_PACKAGE, TEST_PACKAGE)
                 .build());
 
         // Verify folder
@@ -111,10 +112,12 @@ public class DefaultLayoutProviderTest {
         // Add a placeholder session info so that the widget exists
         SessionParams params = new SessionParams(SessionParams.MODE_FULL_INSTALL);
         params.setAppPackageName(pendingAppPkg);
-        params.setAppIcon(BitmapInfo.LOW_RES_ICON);
 
         PackageInstaller installer = mTargetContext.getPackageManager().getPackageInstaller();
-        installer.createSession(params);
+        int sessionId = installer.createSession(params);
+        SessionInfo sessionInfo = installer.getSessionInfo(sessionId);
+        setField(sessionInfo, "installerPackageName", "com.test");
+        setField(sessionInfo, "appIcon", BitmapInfo.LOW_RES_ICON);
 
         writeLayoutAndLoad(new LauncherLayoutBuilder().atWorkspace(0, 1, 0)
                 .putWidget(pendingAppPkg, "PlaceholderWidget", 2, 2));
diff --git a/tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java b/robolectric_tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
similarity index 89%
rename from tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
rename to robolectric_tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
index 005389e5bb..87b08870c8 100644
--- a/tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
+++ b/robolectric_tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
@@ -30,31 +30,26 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
 import android.content.Context;
-import android.content.Intent;
 import android.database.Cursor;
 import android.database.sqlite.SQLiteDatabase;
 import android.graphics.Point;
 import android.os.Process;
 
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.util.LauncherModelHelper;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
 
-import java.util.HashMap;
 import java.util.HashSet;
 
 /** Unit tests for {@link GridSizeMigrationTaskV2} */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
 public class GridSizeMigrationTaskV2Test {
 
     private LauncherModelHelper mModelHelper;
@@ -78,7 +73,7 @@ public class GridSizeMigrationTaskV2Test {
     @Before
     public void setUp() {
         mModelHelper = new LauncherModelHelper();
-        mContext = mModelHelper.sandboxContext;
+        mContext = RuntimeEnvironment.application;
         mDb = mModelHelper.provider.getDb();
 
         mValidPackages = new HashSet<>();
@@ -103,13 +98,8 @@ public class GridSizeMigrationTaskV2Test {
                 LauncherSettings.Favorites.TMP_TABLE);
     }
 
-    @After
-    public void tearDown() {
-        mModelHelper.destroy();
-    }
-
     @Test
-    public void testMigration() throws Exception {
+    public void testMigration() {
         int[] srcHotseatItems = {
                 mModelHelper.addItem(APP_ICON, 0, HOTSEAT, 0, 0, testPackage1, 1, TMP_CONTENT_URI),
                 mModelHelper.addItem(SHORTCUT, 1, HOTSEAT, 0, 0, testPackage2, 2, TMP_CONTENT_URI),
@@ -144,17 +134,17 @@ public class GridSizeMigrationTaskV2Test {
         // Check hotseat items
         Cursor c = mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
                 new String[]{LauncherSettings.Favorites.SCREEN, LauncherSettings.Favorites.INTENT},
-                "container=" + CONTAINER_HOTSEAT, null, LauncherSettings.Favorites.SCREEN, null);
+                "container=" + CONTAINER_HOTSEAT, null, null, null);
         assertEquals(c.getCount(), mIdp.numDatabaseHotseatIcons);
         int screenIndex = c.getColumnIndex(LauncherSettings.Favorites.SCREEN);
         int intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
         c.moveToNext();
-        assertEquals(c.getInt(screenIndex), 0);
-        assertTrue(c.getString(intentIndex).contains(testPackage1));
-        c.moveToNext();
         assertEquals(c.getInt(screenIndex), 1);
         assertTrue(c.getString(intentIndex).contains(testPackage2));
         c.moveToNext();
+        assertEquals(c.getInt(screenIndex), 0);
+        assertTrue(c.getString(intentIndex).contains(testPackage1));
+        c.moveToNext();
         assertEquals(c.getInt(screenIndex), 2);
         assertTrue(c.getString(intentIndex).contains(testPackage3));
         c.moveToNext();
@@ -167,24 +157,35 @@ public class GridSizeMigrationTaskV2Test {
                 new String[]{LauncherSettings.Favorites.CELLX, LauncherSettings.Favorites.CELLY,
                         LauncherSettings.Favorites.INTENT},
                 "container=" + CONTAINER_DESKTOP, null, null, null);
+        assertEquals(c.getCount(), 6);
         intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
         int cellXIndex = c.getColumnIndex(LauncherSettings.Favorites.CELLX);
         int cellYIndex = c.getColumnIndex(LauncherSettings.Favorites.CELLY);
 
-        HashMap<String, Point> locMap = new HashMap<>();
-        while (c.moveToNext()) {
-            locMap.put(
-                    Intent.parseUri(c.getString(intentIndex), 0).getPackage(),
-                    new Point(c.getInt(cellXIndex), c.getInt(cellYIndex)));
-        }
-        c.close();
+        c.moveToNext();
+        assertTrue(c.getString(intentIndex).contains(testPackage7));
+        c.moveToNext();
+        assertTrue(c.getString(intentIndex).contains(testPackage6));
+        assertEquals(c.getInt(cellXIndex), 0);
+        assertEquals(c.getInt(cellYIndex), 3);
+        c.moveToNext();
+        assertTrue(c.getString(intentIndex).contains(testPackage10));
+        assertEquals(c.getInt(cellXIndex), 1);
+        assertEquals(c.getInt(cellYIndex), 3);
+        c.moveToNext();
+        assertTrue(c.getString(intentIndex).contains(testPackage5));
+        assertEquals(c.getInt(cellXIndex), 2);
+        assertEquals(c.getInt(cellYIndex), 3);
+        c.moveToNext();
+        assertTrue(c.getString(intentIndex).contains(testPackage9));
+        assertEquals(c.getInt(cellXIndex), 3);
+        assertEquals(c.getInt(cellYIndex), 3);
+        c.moveToNext();
+        assertTrue(c.getString(intentIndex).contains(testPackage8));
+        assertEquals(c.getInt(cellXIndex), 0);
+        assertEquals(c.getInt(cellYIndex), 2);
 
-        assertEquals(locMap.size(), 6);
-        assertEquals(new Point(0, 2), locMap.get(testPackage8));
-        assertEquals(new Point(0, 3), locMap.get(testPackage6));
-        assertEquals(new Point(1, 3), locMap.get(testPackage10));
-        assertEquals(new Point(2, 3), locMap.get(testPackage5));
-        assertEquals(new Point(3, 3), locMap.get(testPackage9));
+        c.close();
     }
 
     @Test
@@ -211,7 +212,7 @@ public class GridSizeMigrationTaskV2Test {
         // Check hotseat items
         Cursor c = mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
                 new String[]{LauncherSettings.Favorites.SCREEN, LauncherSettings.Favorites.INTENT},
-                "container=" + CONTAINER_HOTSEAT, null, LauncherSettings.Favorites.SCREEN, null);
+                "container=" + CONTAINER_HOTSEAT, null, null, null);
         assertEquals(c.getCount(), numSrcDatabaseHotseatIcons);
         int screenIndex = c.getColumnIndex(LauncherSettings.Favorites.SCREEN);
         int intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
@@ -256,7 +257,7 @@ public class GridSizeMigrationTaskV2Test {
         // Check hotseat items
         Cursor c = mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
                 new String[]{LauncherSettings.Favorites.SCREEN, LauncherSettings.Favorites.INTENT},
-                "container=" + CONTAINER_HOTSEAT, null, LauncherSettings.Favorites.SCREEN, null);
+                "container=" + CONTAINER_HOTSEAT, null, null, null);
         assertEquals(c.getCount(), mIdp.numDatabaseHotseatIcons);
         int screenIndex = c.getColumnIndex(LauncherSettings.Favorites.SCREEN);
         int intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
diff --git a/tests/src/com/android/launcher3/model/LoaderCursorTest.java b/robolectric_tests/src/com/android/launcher3/model/LoaderCursorTest.java
similarity index 93%
rename from tests/src/com/android/launcher3/model/LoaderCursorTest.java
rename to robolectric_tests/src/com/android/launcher3/model/LoaderCursorTest.java
index 6444ef6927..800311afb6 100644
--- a/tests/src/com/android/launcher3/model/LoaderCursorTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/LoaderCursorTest.java
@@ -16,8 +16,6 @@
 
 package com.android.launcher3.model;
 
-import static androidx.test.InstrumentationRegistry.getContext;
-
 import static com.android.launcher3.LauncherSettings.Favorites.CELLX;
 import static com.android.launcher3.LauncherSettings.Favorites.CELLY;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER;
@@ -35,7 +33,7 @@ import static com.android.launcher3.LauncherSettings.Favorites.RESTORED;
 import static com.android.launcher3.LauncherSettings.Favorites.SCREEN;
 import static com.android.launcher3.LauncherSettings.Favorites.TITLE;
 import static com.android.launcher3.LauncherSettings.Favorites._ID;
-import static com.android.launcher3.util.LauncherModelHelper.TEST_ACTIVITY;
+import static com.android.launcher3.util.LauncherModelHelper.TEST_PACKAGE;
 
 import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertFalse;
@@ -43,37 +41,37 @@ import static junit.framework.Assert.assertNotNull;
 import static junit.framework.Assert.assertNull;
 import static junit.framework.Assert.assertTrue;
 
+import static org.robolectric.Shadows.shadowOf;
+
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
 import android.database.MatrixCursor;
 import android.os.Process;
 
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.util.Executors;
-import com.android.launcher3.util.LauncherModelHelper;
 import com.android.launcher3.util.PackageManagerHelper;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.annotation.LooperMode.Mode;
 
 /**
  * Tests for {@link LoaderCursor}
  */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
+@LooperMode(Mode.PAUSED)
 public class LoaderCursorTest {
 
-    private LauncherModelHelper mModelHelper;
     private LauncherAppState mApp;
 
     private MatrixCursor mCursor;
@@ -84,8 +82,7 @@ public class LoaderCursorTest {
 
     @Before
     public void setup() {
-        mModelHelper = new LauncherModelHelper();
-        mContext = mModelHelper.sandboxContext;
+        mContext = RuntimeEnvironment.application;
         mIDP = InvariantDeviceProfile.INSTANCE.get(mContext);
         mApp = LauncherAppState.getInstance(mContext);
 
@@ -100,11 +97,6 @@ public class LoaderCursorTest {
         ums.allUsers.put(0, Process.myUserHandle());
     }
 
-    @After
-    public void tearDown() {
-        mModelHelper.destroy();
-    }
-
     private void initCursor(int itemType, String title) {
         mCursor.newRow()
                 .add(_ID, 1)
@@ -125,7 +117,9 @@ public class LoaderCursorTest {
 
     @Test
     public void getAppShortcutInfo_dontAllowMissing_validComponent() throws Exception {
-        ComponentName cn = new ComponentName(getContext(), TEST_ACTIVITY);
+        ComponentName cn = new ComponentName(TEST_PACKAGE, TEST_PACKAGE);
+        shadowOf(mContext.getPackageManager()).addActivityIfNotPresent(cn);
+
         initCursor(ITEM_TYPE_APPLICATION, "");
         assertTrue(mLoaderCursor.moveToNext());
 
diff --git a/tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java b/robolectric_tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
similarity index 65%
rename from tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
rename to robolectric_tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
index 42c9f11a52..43193553fe 100644
--- a/tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
@@ -15,60 +15,69 @@
  */
 package com.android.launcher3.model;
 
+import static com.android.launcher3.util.Executors.createAndStartNewLooper;
 import static com.android.launcher3.util.LauncherModelHelper.TEST_PACKAGE;
 
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.spy;
+import static org.robolectric.Shadows.shadowOf;
 
 import android.os.Process;
 
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-
 import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.shadows.ShadowLooperExecutor;
 import com.android.launcher3.util.Executors;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.LauncherLayoutBuilder;
 import com.android.launcher3.util.LauncherModelHelper;
+import com.android.launcher3.util.LooperExecutor;
 import com.android.launcher3.util.RunnableList;
-import com.android.launcher3.util.TestUtil;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.annotation.LooperMode.Mode;
+import org.robolectric.shadow.api.Shadow;
+import org.robolectric.shadows.ShadowPackageManager;
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashSet;
 import java.util.List;
-import java.util.Set;
 import java.util.stream.Collectors;
 
 /**
  * Tests to verify multiple callbacks in Loader
  */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
+@LooperMode(Mode.PAUSED)
 public class ModelMultiCallbacksTest {
 
     private LauncherModelHelper mModelHelper;
 
-    @Before
-    public void setUp() {
-        mModelHelper = new LauncherModelHelper();
-    }
+    private ShadowPackageManager mSpm;
+    private LooperExecutor mTempMainExecutor;
 
-    @After
-    public void tearDown() throws Exception {
-        mModelHelper.destroy();
-        TestUtil.uninstallDummyApp();
+    @Before
+    public void setUp() throws Exception {
+        mModelHelper = new LauncherModelHelper();
+        mModelHelper.installApp(TEST_PACKAGE);
+
+        mSpm = shadowOf(RuntimeEnvironment.application.getPackageManager());
+
+        // Since robolectric tests run on main thread, we run the loader-UI calls on a temp thread,
+        // so that we can wait appropriately for the loader to complete.
+        mTempMainExecutor = new LooperExecutor(createAndStartNewLooper("tempMain"));
+        ShadowLooperExecutor sle = Shadow.extract(Executors.MAIN_EXECUTOR);
+        sle.setHandler(mTempMainExecutor.getHandler());
     }
 
     @Test
@@ -76,7 +85,7 @@ public class ModelMultiCallbacksTest {
         setupWorkspacePages(3);
 
         MyCallbacks cb1 = spy(MyCallbacks.class);
-        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().addCallbacksAndLoad(cb1));
+        mModelHelper.getModel().addCallbacksAndLoad(cb1);
 
         waitForLoaderAndTempMainThread();
         cb1.verifySynchronouslyBound(3);
@@ -85,10 +94,10 @@ public class ModelMultiCallbacksTest {
         cb1.reset();
         MyCallbacks cb2 = spy(MyCallbacks.class);
         cb2.mPageToBindSync = IntSet.wrap(2);
-        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().addCallbacksAndLoad(cb2));
+        mModelHelper.getModel().addCallbacksAndLoad(cb2);
 
         waitForLoaderAndTempMainThread();
-        assertFalse(cb1.bindStarted);
+        cb1.verifySynchronouslyBound(3);
         cb2.verifySynchronouslyBound(3);
 
         // Remove callbacks
@@ -96,7 +105,7 @@ public class ModelMultiCallbacksTest {
         cb2.reset();
 
         // No effect on callbacks when removing an callback
-        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().removeCallbacks(cb2));
+        mModelHelper.getModel().removeCallbacks(cb2);
         waitForLoaderAndTempMainThread();
         assertNull(cb1.mPendingTasks);
         assertNull(cb2.mPendingTasks);
@@ -110,48 +119,52 @@ public class ModelMultiCallbacksTest {
 
     @Test
     public void testTwoCallbacks_receiveUpdates() throws Exception {
-        TestUtil.uninstallDummyApp();
-
         setupWorkspacePages(1);
 
         MyCallbacks cb1 = spy(MyCallbacks.class);
         MyCallbacks cb2 = spy(MyCallbacks.class);
-        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().addCallbacksAndLoad(cb1));
-        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().addCallbacksAndLoad(cb2));
+        mModelHelper.getModel().addCallbacksAndLoad(cb1);
+        mModelHelper.getModel().addCallbacksAndLoad(cb2);
         waitForLoaderAndTempMainThread();
 
-        assertTrue(cb1.allApps().contains(TEST_PACKAGE));
-        assertTrue(cb2.allApps().contains(TEST_PACKAGE));
+        cb1.verifyApps(TEST_PACKAGE);
+        cb2.verifyApps(TEST_PACKAGE);
 
         // Install package 1
-        TestUtil.installDummyApp();
-        mModelHelper.getModel().onPackageAdded(TestUtil.DUMMY_PACKAGE, Process.myUserHandle());
+        String pkg1 = "com.test.pkg1";
+        mModelHelper.installApp(pkg1);
+        mModelHelper.getModel().onPackageAdded(pkg1, Process.myUserHandle());
         waitForLoaderAndTempMainThread();
-        assertTrue(cb1.allApps().contains(TestUtil.DUMMY_PACKAGE));
-        assertTrue(cb2.allApps().contains(TestUtil.DUMMY_PACKAGE));
+        cb1.verifyApps(TEST_PACKAGE, pkg1);
+        cb2.verifyApps(TEST_PACKAGE, pkg1);
+
+        // Install package 2
+        String pkg2 = "com.test.pkg2";
+        mModelHelper.installApp(pkg2);
+        mModelHelper.getModel().onPackageAdded(pkg2, Process.myUserHandle());
+        waitForLoaderAndTempMainThread();
+        cb1.verifyApps(TEST_PACKAGE, pkg1, pkg2);
+        cb2.verifyApps(TEST_PACKAGE, pkg1, pkg2);
 
         // Uninstall package 2
-        TestUtil.uninstallDummyApp();
-        mModelHelper.getModel().onPackageRemoved(TestUtil.DUMMY_PACKAGE, Process.myUserHandle());
+        mSpm.removePackage(pkg1);
+        mModelHelper.getModel().onPackageRemoved(pkg1, Process.myUserHandle());
         waitForLoaderAndTempMainThread();
-        assertFalse(cb1.allApps().contains(TestUtil.DUMMY_PACKAGE));
-        assertFalse(cb2.allApps().contains(TestUtil.DUMMY_PACKAGE));
+        cb1.verifyApps(TEST_PACKAGE, pkg2);
+        cb2.verifyApps(TEST_PACKAGE, pkg2);
 
         // Unregister a callback and verify updates no longer received
-        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().removeCallbacks(cb2));
-        TestUtil.installDummyApp();
-        mModelHelper.getModel().onPackageAdded(TestUtil.DUMMY_PACKAGE, Process.myUserHandle());
+        mModelHelper.getModel().removeCallbacks(cb2);
+        mSpm.removePackage(pkg2);
+        mModelHelper.getModel().onPackageRemoved(pkg2, Process.myUserHandle());
         waitForLoaderAndTempMainThread();
-
-        // cb2 didn't get the update
-        assertTrue(cb1.allApps().contains(TestUtil.DUMMY_PACKAGE));
-        assertFalse(cb2.allApps().contains(TestUtil.DUMMY_PACKAGE));
+        cb1.verifyApps(TEST_PACKAGE);
+        cb2.verifyApps(TEST_PACKAGE, pkg2);
     }
 
     private void waitForLoaderAndTempMainThread() throws Exception {
-        Executors.MAIN_EXECUTOR.submit(() -> { }).get();
         Executors.MODEL_EXECUTOR.submit(() -> { }).get();
-        Executors.MAIN_EXECUTOR.submit(() -> { }).get();
+        mTempMainExecutor.submit(() -> { }).get();
     }
 
     private void setupWorkspacePages(int pageCount) throws Exception {
@@ -170,15 +183,9 @@ public class ModelMultiCallbacksTest {
         IntSet mPageBoundSync = new IntSet();
         RunnableList mPendingTasks;
         AppInfo[] mAppInfos;
-        boolean bindStarted;
 
         MyCallbacks() { }
 
-        @Override
-        public void startBinding() {
-            bindStarted = true;
-        }
-
         @Override
         public void onInitialBindComplete(IntSet boundPages, RunnableList pendingTasks) {
             mPageBoundSync = boundPages;
@@ -205,12 +212,10 @@ public class ModelMultiCallbacksTest {
             mPageBoundSync = new IntSet();
             mPendingTasks = null;
             mAppInfos = null;
-            bindStarted = false;
         }
 
         public void verifySynchronouslyBound(int totalItems) {
             // Verify that the requested page is bound synchronously
-            assertTrue(bindStarted);
             assertEquals(mPageToBindSync, mPageBoundSync);
             assertEquals(mItems.size(), 1);
             assertEquals(IntSet.wrap(mItems.get(0).screenId), mPageBoundSync);
@@ -221,14 +226,12 @@ public class ModelMultiCallbacksTest {
             assertEquals(mItems.size(), totalItems);
         }
 
-        public Set<String> allApps() {
-            return Arrays.stream(mAppInfos)
-                    .map(ai -> ai.getTargetComponent().getPackageName())
-                    .collect(Collectors.toSet());
-        }
-
         public void verifyApps(String... apps) {
-            assertTrue(allApps().containsAll(Arrays.asList(apps)));
+            assertEquals(apps.length, mAppInfos.length);
+            assertEquals(Arrays.stream(mAppInfos)
+                    .map(ai -> ai.getTargetComponent().getPackageName())
+                    .collect(Collectors.toSet()),
+                    new HashSet<>(Arrays.asList(apps)));
         }
     }
 }
diff --git a/tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java b/robolectric_tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
similarity index 87%
rename from tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
rename to robolectric_tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
index 519191e251..412ace03ff 100644
--- a/tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
@@ -2,19 +2,18 @@ package com.android.launcher3.model;
 
 import static org.junit.Assert.assertEquals;
 
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.pm.PackageInstallInfo;
 import com.android.launcher3.util.LauncherModelHelper;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.annotation.LooperMode.Mode;
 
 import java.util.Arrays;
 import java.util.HashSet;
@@ -22,8 +21,8 @@ import java.util.HashSet;
 /**
  * Tests for {@link PackageInstallStateChangedTask}
  */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
+@LooperMode(Mode.PAUSED)
 public class PackageInstallStateChangedTaskTest {
 
     private LauncherModelHelper mModelHelper;
@@ -31,12 +30,7 @@ public class PackageInstallStateChangedTaskTest {
     @Before
     public void setup() throws Exception {
         mModelHelper = new LauncherModelHelper();
-        mModelHelper.initializeData("package_install_state_change_task_data");
-    }
-
-    @After
-    public void tearDown() {
-        mModelHelper.destroy();
+        mModelHelper.initializeData("/package_install_state_change_task_data.txt");
     }
 
     private PackageInstallStateChangedTask newTask(String pkg, int progress) {
@@ -72,7 +66,7 @@ public class PackageInstallStateChangedTaskTest {
         HashSet<Integer> updates = new HashSet<>(Arrays.asList(idsUpdated));
         for (ItemInfo info : mModelHelper.getBgDataModel().itemsIdMap) {
             if (info instanceof WorkspaceItemInfo) {
-                assertEquals(updates.contains(info.id) ? progress: 100,
+                assertEquals(updates.contains(info.id) ? progress: 0,
                         ((WorkspaceItemInfo) info).getProgressLevel());
             } else {
                 assertEquals(updates.contains(info.id) ? progress: -1,
diff --git a/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java b/robolectric_tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
similarity index 93%
rename from tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
rename to robolectric_tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
index 48305eebdf..4184d33a68 100644
--- a/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
+++ b/robolectric_tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
@@ -21,21 +21,18 @@ import android.content.ContentValues;
 import android.database.Cursor;
 import android.database.sqlite.SQLiteDatabase;
 
-import androidx.test.InstrumentationRegistry;
-import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
-
 import com.android.launcher3.LauncherProvider.DatabaseHelper;
 import com.android.launcher3.LauncherSettings.Favorites;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
 
 /**
  * Tests for {@link RestoreDbTask}
  */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(RobolectricTestRunner.class)
 public class RestoreDbTaskTest {
 
     @Test
@@ -98,7 +95,7 @@ public class RestoreDbTaskTest {
         private final long mProfileId;
 
         MyDatabaseHelper(long profileId) {
-            super(InstrumentationRegistry.getInstrumentation().getTargetContext(), null, false);
+            super(RuntimeEnvironment.application, null, false);
             mProfileId = profileId;
         }
 
diff --git a/robolectric_tests/src/com/android/launcher3/secondarydisplay/SDWorkModeTest.java b/robolectric_tests/src/com/android/launcher3/secondarydisplay/SDWorkModeTest.java
new file mode 100644
index 0000000000..e3694aeddc
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/secondarydisplay/SDWorkModeTest.java
@@ -0,0 +1,125 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.secondarydisplay;
+
+import static com.android.launcher3.util.LauncherModelHelper.TEST_PACKAGE;
+import static com.android.launcher3.util.LauncherUIHelper.doLayout;
+import static com.android.launcher3.util.Preconditions.assertNotNull;
+
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertTrue;
+
+import android.content.Context;
+import android.os.UserManager;
+import android.provider.Settings;
+
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.allapps.AllAppsPagedView;
+import com.android.launcher3.allapps.AllAppsRecyclerView;
+import com.android.launcher3.util.LauncherLayoutBuilder;
+import com.android.launcher3.util.LauncherModelHelper;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.robolectric.Robolectric;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.annotation.LooperMode.Mode;
+import org.robolectric.shadow.api.Shadow;
+import org.robolectric.shadows.ShadowUserManager;
+
+/**
+ * Tests for {@link SecondaryDisplayLauncher} with work profile
+ */
+@RunWith(RobolectricTestRunner.class)
+@LooperMode(Mode.PAUSED)
+public class SDWorkModeTest {
+
+    private static final int SYSTEM_USER = 0;
+    private static final int FLAG_SYSTEM = 0x00000800;
+    private static final int WORK_PROFILE_ID = 10;
+    private static final int FLAG_PROFILE = 0x00001000;
+
+    private Context mTargetContext;
+    private InvariantDeviceProfile mIdp;
+    private LauncherModelHelper mModelHelper;
+
+    @Before
+    public void setup() throws Exception {
+        mModelHelper = new LauncherModelHelper();
+        mTargetContext = RuntimeEnvironment.application;
+        mIdp = InvariantDeviceProfile.INSTANCE.get(mTargetContext);
+        Settings.Global.putFloat(mTargetContext.getContentResolver(),
+                Settings.Global.WINDOW_ANIMATION_SCALE, 0);
+
+        mModelHelper.installApp(TEST_PACKAGE);
+    }
+
+    @Test
+    public void testAllAppsList_noWorkProfile() throws Exception {
+        SecondaryDisplayLauncher launcher = loadLauncher();
+        launcher.showAppDrawer(true);
+        doLayout(launcher);
+
+        verifyRecyclerViewCount(launcher.getAppsView().getActiveRecyclerView());
+    }
+
+    @Test
+    public void testAllAppsList_workProfile() throws Exception {
+        ShadowUserManager sum = Shadow.extract(mTargetContext.getSystemService(UserManager.class));
+        sum.addUser(SYSTEM_USER, "me", FLAG_SYSTEM);
+        sum.addProfile(SYSTEM_USER, WORK_PROFILE_ID, "work", FLAG_PROFILE);
+
+        SecondaryDisplayLauncher launcher = loadLauncher();
+        launcher.showAppDrawer(true);
+        doLayout(launcher);
+
+        AllAppsRecyclerView rv1 = launcher.getAppsView().getActiveRecyclerView();
+        verifyRecyclerViewCount(rv1);
+
+        assertNotNull(launcher.getAppsView().getWorkModeSwitch());
+        assertTrue(launcher.getAppsView().getRecyclerViewContainer() instanceof AllAppsPagedView);
+
+        AllAppsPagedView pagedView =
+                (AllAppsPagedView) launcher.getAppsView().getRecyclerViewContainer();
+        pagedView.snapToPageImmediately(1);
+        doLayout(launcher);
+
+        AllAppsRecyclerView rv2 = launcher.getAppsView().getActiveRecyclerView();
+        verifyRecyclerViewCount(rv2);
+        assertNotSame(rv1, rv2);
+    }
+
+    private SecondaryDisplayLauncher loadLauncher() throws Exception {
+        // Install 100 apps
+        for (int i = 0; i < 100; i++) {
+            mModelHelper.installApp(TEST_PACKAGE + i);
+        }
+        mModelHelper.setupDefaultLayoutProvider(new LauncherLayoutBuilder()).loadModelSync();
+        SecondaryDisplayLauncher launcher =
+                Robolectric.buildActivity(SecondaryDisplayLauncher.class).setup().get();
+        doLayout(launcher);
+        return launcher;
+    }
+
+    private void verifyRecyclerViewCount(AllAppsRecyclerView rv) {
+        int childCount = rv.getChildCount();
+        assertTrue(childCount > 0);
+        assertTrue(childCount < 100);
+    }
+}
diff --git a/robolectric_tests/src/com/android/launcher3/testing/TestActivity.java b/robolectric_tests/src/com/android/launcher3/testing/TestActivity.java
new file mode 100644
index 0000000000..17d0ac1a9f
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/testing/TestActivity.java
@@ -0,0 +1,46 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.testing;
+
+import com.android.launcher3.BaseActivity;
+import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.views.ActivityContext;
+import com.android.launcher3.views.BaseDragLayer;
+
+/** An empty activity for {@link android.app.Fragment}s, {@link android.view.View}s testing. */
+public class TestActivity extends BaseActivity implements ActivityContext {
+
+    private DeviceProfile mDeviceProfile;
+
+    @Override
+    public BaseDragLayer getDragLayer() {
+        return new BaseDragLayer(this, /* attrs= */ null, /* alphaChannelCount= */ 1) {
+            @Override
+            public void recreateControllers() {
+                // Do nothing.
+            }
+        };
+    }
+
+    @Override
+    public DeviceProfile getDeviceProfile() {
+        return mDeviceProfile;
+    }
+
+    public void setDeviceProfile(DeviceProfile deviceProfile) {
+        mDeviceProfile = deviceProfile;
+    }
+}
diff --git a/tests/src/com/android/launcher3/util/LauncherLayoutBuilder.java b/robolectric_tests/src/com/android/launcher3/util/LauncherLayoutBuilder.java
similarity index 100%
rename from tests/src/com/android/launcher3/util/LauncherLayoutBuilder.java
rename to robolectric_tests/src/com/android/launcher3/util/LauncherLayoutBuilder.java
diff --git a/tests/src/com/android/launcher3/util/LauncherModelHelper.java b/robolectric_tests/src/com/android/launcher3/util/LauncherModelHelper.java
similarity index 60%
rename from tests/src/com/android/launcher3/util/LauncherModelHelper.java
rename to robolectric_tests/src/com/android/launcher3/util/LauncherModelHelper.java
index c9b63aeb11..846e2019a0 100644
--- a/tests/src/com/android/launcher3/util/LauncherModelHelper.java
+++ b/robolectric_tests/src/com/android/launcher3/util/LauncherModelHelper.java
@@ -15,40 +15,26 @@
  */
 package com.android.launcher3.util;
 
-import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
-import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
+import static android.content.Intent.ACTION_CREATE_SHORTCUT;
 
 import static com.android.launcher3.LauncherSettings.Favorites.CONTENT_URI;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 
-import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.atLeast;
-import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
+import static org.robolectric.Shadows.shadowOf;
 
 import android.content.ComponentName;
-import android.content.ContentProvider;
-import android.content.ContentResolver;
 import android.content.ContentValues;
 import android.content.Context;
 import android.content.Intent;
-import android.content.pm.PackageManager;
-import android.content.pm.ProviderInfo;
-import android.content.res.Resources;
+import android.content.IntentFilter;
+import android.content.pm.PackageManager.NameNotFoundException;
 import android.database.sqlite.SQLiteDatabase;
 import android.net.Uri;
-import android.os.ParcelFileDescriptor;
-import android.os.ParcelFileDescriptor.AutoCloseOutputStream;
 import android.os.Process;
 import android.provider.Settings;
-import android.test.mock.MockContentResolver;
-import android.util.ArrayMap;
-
-import androidx.test.core.app.ApplicationProvider;
-import androidx.test.uiautomator.UiDevice;
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
@@ -59,30 +45,27 @@ import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.model.AllAppsList;
 import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.BgDataModel.Callbacks;
-import com.android.launcher3.model.ItemInstallQueue;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfo;
-import com.android.launcher3.pm.InstallSessionHelper;
 import com.android.launcher3.pm.UserCache;
-import com.android.launcher3.testing.TestInformationProvider;
-import com.android.launcher3.uioverrides.plugins.PluginManagerWrapper;
-import com.android.launcher3.util.MainThreadInitializedObject.ObjectProvider;
-import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext;
-import com.android.launcher3.widget.custom.CustomWidgetManager;
+import com.android.launcher3.shadows.ShadowLooperExecutor;
 
 import org.mockito.ArgumentCaptor;
+import org.robolectric.Robolectric;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.shadow.api.Shadow;
+import org.robolectric.shadows.ShadowContentResolver;
+import org.robolectric.shadows.ShadowPackageManager;
+import org.robolectric.util.ReflectionHelpers;
 
 import java.io.BufferedReader;
+import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
-import java.io.File;
-import java.io.FileNotFoundException;
 import java.io.InputStreamReader;
 import java.io.OutputStreamWriter;
 import java.lang.reflect.Field;
 import java.util.HashMap;
 import java.util.List;
-import java.util.UUID;
-import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
 import java.util.function.Function;
@@ -98,9 +81,7 @@ public class LauncherModelHelper {
     public static final int APP_ICON = LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
     public static final int SHORTCUT = LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT;
     public static final int NO__ICON = -1;
-
-    public static final String TEST_PACKAGE = testContext().getPackageName();
-    public static final String TEST_ACTIVITY = "com.android.launcher3.tests.Activity2";
+    public static final String TEST_PACKAGE = "com.android.launcher3.validpackage";
 
     // Authority for providing a test default-workspace-layout data.
     private static final String TEST_PROVIDER_AUTHORITY =
@@ -109,42 +90,21 @@ public class LauncherModelHelper {
     private static final int DEFAULT_GRID_SIZE = 4;
 
     private final HashMap<Class, HashMap<String, Field>> mFieldCache = new HashMap<>();
-    private final MockContentResolver mMockResolver = new MockContentResolver();
     public final TestLauncherProvider provider;
-    public final SanboxModelContext sandboxContext;
-
-    public final long defaultProfileId;
+    private final long mDefaultProfileId;
 
     private BgDataModel mDataModel;
     private AllAppsList mAllAppsList;
 
     public LauncherModelHelper() {
-        Context context = getApplicationContext();
-        // System settings cache content provider. Ensure that they are statically initialized
-        Settings.Secure.getString(context.getContentResolver(), "test");
-        Settings.System.getString(context.getContentResolver(), "test");
-        Settings.Global.getString(context.getContentResolver(), "test");
-
-        provider = new TestLauncherProvider();
-        sandboxContext = new SanboxModelContext();
-        defaultProfileId = UserCache.INSTANCE.get(sandboxContext)
+        provider = Robolectric.setupContentProvider(TestLauncherProvider.class);
+        mDefaultProfileId = UserCache.INSTANCE.get(RuntimeEnvironment.application)
                 .getSerialNumberForUser(Process.myUserHandle());
-        setupProvider(LauncherProvider.AUTHORITY, provider);
-    }
-
-    protected void setupProvider(String authority, ContentProvider provider) {
-        ProviderInfo providerInfo = new ProviderInfo();
-        providerInfo.authority = authority;
-        providerInfo.applicationInfo = sandboxContext.getApplicationInfo();
-        provider.attachInfo(sandboxContext, providerInfo);
-        mMockResolver.addProvider(providerInfo.authority, provider);
-        doReturn(providerInfo)
-                .when(sandboxContext.mPm)
-                .resolveContentProvider(eq(authority), anyInt());
+        ShadowContentResolver.registerProviderInternal(LauncherProvider.AUTHORITY, provider);
     }
 
     public LauncherModel getModel() {
-        return LauncherAppState.getInstance(sandboxContext).getModel();
+        return LauncherAppState.getInstance(RuntimeEnvironment.application).getModel();
     }
 
     public synchronized BgDataModel getBgDataModel() {
@@ -161,28 +121,6 @@ public class LauncherModelHelper {
         return mAllAppsList;
     }
 
-    public void destroy() {
-        // When destroying the context, make sure that the model thread is blocked, so that no
-        // new jobs get posted while we are cleaning up
-        CountDownLatch l1 = new CountDownLatch(1);
-        CountDownLatch l2 = new CountDownLatch(1);
-        MODEL_EXECUTOR.execute(() -> {
-            l1.countDown();
-            waitOrThrow(l2);
-        });
-        waitOrThrow(l1);
-        sandboxContext.onDestroy();
-        l2.countDown();
-    }
-
-    private void waitOrThrow(CountDownLatch latch) {
-        try {
-            latch.await();
-        } catch (Exception e) {
-            throw new RuntimeException(e);
-        }
-    }
-
     /**
      * Synchronously executes the task and returns all the UI callbacks posted.
      */
@@ -223,16 +161,13 @@ public class LauncherModelHelper {
      * Initializes mock data for the test.
      */
     public void initializeData(String resourceName) throws Exception {
+        Context targetContext = RuntimeEnvironment.application;
         BgDataModel bgDataModel = getBgDataModel();
         AllAppsList allAppsList = getAllAppsList();
 
         MODEL_EXECUTOR.submit(() -> {
-            // Copy apk from resources to a local file and install from there.
-            Resources resources = testContext().getResources();
-            int resId = resources.getIdentifier(
-                    resourceName, "raw", testContext().getPackageName());
             try (BufferedReader reader = new BufferedReader(new InputStreamReader(
-                    resources.openRawResource(resId)))) {
+                    this.getClass().getResourceAsStream(resourceName)))) {
                 String line;
                 HashMap<String, Class> classMap = new HashMap<>();
                 while ((line = reader.readLine()) != null) {
@@ -246,7 +181,7 @@ public class LauncherModelHelper {
                             classMap.put(commands[1], Class.forName(commands[2]));
                             break;
                         case "bgItem":
-                            bgDataModel.addItem(sandboxContext,
+                            bgDataModel.addItem(targetContext,
                                     (ItemInfo) initItem(classMap.get(commands[1]), commands, 2),
                                     false);
                             break;
@@ -301,7 +236,7 @@ public class LauncherModelHelper {
     }
 
     public int addItem(int type, int screen, int container, int x, int y) {
-        return addItem(type, screen, container, x, y, defaultProfileId, TEST_PACKAGE);
+        return addItem(type, screen, container, x, y, mDefaultProfileId, TEST_PACKAGE);
     }
 
     public int addItem(int type, int screen, int container, int x, int y, long profileId) {
@@ -309,12 +244,12 @@ public class LauncherModelHelper {
     }
 
     public int addItem(int type, int screen, int container, int x, int y, String packageName) {
-        return addItem(type, screen, container, x, y, defaultProfileId, packageName);
+        return addItem(type, screen, container, x, y, mDefaultProfileId, packageName);
     }
 
     public int addItem(int type, int screen, int container, int x, int y, String packageName,
             int id, Uri contentUri) {
-        addItem(type, screen, container, x, y, defaultProfileId, packageName, id, contentUri);
+        addItem(type, screen, container, x, y, mDefaultProfileId, packageName, id, contentUri);
         return id;
     }
 
@@ -325,7 +260,8 @@ public class LauncherModelHelper {
      */
     public int addItem(int type, int screen, int container, int x, int y, long profileId,
             String packageName) {
-        int id = LauncherSettings.Settings.call(sandboxContext.getContentResolver(),
+        Context context = RuntimeEnvironment.application;
+        int id = LauncherSettings.Settings.call(context.getContentResolver(),
                 LauncherSettings.Settings.METHOD_NEW_ITEM_ID)
                 .getInt(LauncherSettings.Settings.EXTRA_VALUE);
         addItem(type, screen, container, x, y, profileId, packageName, id, CONTENT_URI);
@@ -334,6 +270,8 @@ public class LauncherModelHelper {
 
     public void addItem(int type, int screen, int container, int x, int y, long profileId,
             String packageName, int id, Uri contentUri) {
+        Context context = RuntimeEnvironment.application;
+
         ContentValues values = new ContentValues();
         values.put(LauncherSettings.Favorites._ID, id);
         values.put(LauncherSettings.Favorites.CONTAINER, container);
@@ -357,7 +295,7 @@ public class LauncherModelHelper {
             }
         }
 
-        sandboxContext.getContentResolver().insert(contentUri, values);
+        context.getContentResolver().insert(contentUri, values);
     }
 
     public int[][][] createGrid(int[][][] typeArray) {
@@ -365,11 +303,12 @@ public class LauncherModelHelper {
     }
 
     public int[][][] createGrid(int[][][] typeArray, int startScreen) {
-        LauncherSettings.Settings.call(sandboxContext.getContentResolver(),
+        final Context context = RuntimeEnvironment.application;
+        LauncherSettings.Settings.call(context.getContentResolver(),
                 LauncherSettings.Settings.METHOD_CREATE_EMPTY_DB);
-        LauncherSettings.Settings.call(sandboxContext.getContentResolver(),
+        LauncherSettings.Settings.call(context.getContentResolver(),
                 LauncherSettings.Settings.METHOD_CLEAR_EMPTY_DB_FLAG);
-        return createGrid(typeArray, startScreen, defaultProfileId);
+        return createGrid(typeArray, startScreen, mDefaultProfileId);
     }
 
     /**
@@ -381,13 +320,14 @@ public class LauncherModelHelper {
      * @return the same grid representation where each entry is the corresponding item id.
      */
     public int[][][] createGrid(int[][][] typeArray, int startScreen, long profileId) {
+        Context context = RuntimeEnvironment.application;
         int[][][] ids = new int[typeArray.length][][];
         for (int i = 0; i < typeArray.length; i++) {
             // Add screen to DB
             int screenId = startScreen + i;
 
             // Keep the screen id counter up to date
-            LauncherSettings.Settings.call(sandboxContext.getContentResolver(),
+            LauncherSettings.Settings.call(context.getContentResolver(),
                     LauncherSettings.Settings.METHOD_NEW_SCREEN_ID);
 
             ids[i] = new int[typeArray[i].length][];
@@ -413,45 +353,69 @@ public class LauncherModelHelper {
      */
     public LauncherModelHelper setupDefaultLayoutProvider(LauncherLayoutBuilder builder)
             throws Exception {
-        InvariantDeviceProfile idp = InvariantDeviceProfile.INSTANCE.get(sandboxContext);
+        Context context = RuntimeEnvironment.application;
+        InvariantDeviceProfile idp = InvariantDeviceProfile.INSTANCE.get(context);
         idp.numRows = idp.numColumns = idp.numDatabaseHotseatIcons = DEFAULT_GRID_SIZE;
         idp.iconBitmapSize = DEFAULT_BITMAP_SIZE;
 
-        UiDevice.getInstance(getInstrumentation()).executeShellCommand(
-                "settings put secure launcher3.layout.provider " + TEST_PROVIDER_AUTHORITY);
-        ContentProvider cp = new TestInformationProvider() {
+        Settings.Secure.putString(context.getContentResolver(),
+                "launcher3.layout.provider", TEST_PROVIDER_AUTHORITY);
 
-            @Override
-            public ParcelFileDescriptor openFile(Uri uri, String mode)
-                    throws FileNotFoundException {
-                try {
-                    ParcelFileDescriptor[] pipe = ParcelFileDescriptor.createPipe();
-                    AutoCloseOutputStream outputStream = new AutoCloseOutputStream(pipe[1]);
-                    ByteArrayOutputStream bos = new ByteArrayOutputStream();
-                    builder.build(new OutputStreamWriter(bos));
-                    outputStream.write(bos.toByteArray());
-                    outputStream.flush();
-                    outputStream.close();
-                    return pipe[0];
-                } catch (Exception e) {
-                    throw new FileNotFoundException(e.getMessage());
-                }
-            }
-        };
-        setupProvider(TEST_PROVIDER_AUTHORITY, cp);
+        shadowOf(context.getPackageManager())
+                .addProviderIfNotPresent(new ComponentName("com.test", "Mock")).authority =
+                TEST_PROVIDER_AUTHORITY;
+
+        ByteArrayOutputStream bos = new ByteArrayOutputStream();
+        builder.build(new OutputStreamWriter(bos));
+        Uri layoutUri = LauncherProvider.getLayoutUri(TEST_PROVIDER_AUTHORITY, context);
+        shadowOf(context.getContentResolver()).registerInputStream(layoutUri,
+                new ByteArrayInputStream(bos.toByteArray()));
         return this;
     }
 
+    /**
+     * Simulates an apk install with a default main activity with same class and package name
+     */
+    public void installApp(String component) throws NameNotFoundException {
+        IntentFilter filter = new IntentFilter(Intent.ACTION_MAIN);
+        filter.addCategory(Intent.CATEGORY_LAUNCHER);
+        installApp(component, component, filter);
+    }
+
+    /**
+     * Simulates a custom shortcut install
+     */
+    public void installCustomShortcut(String pkg, String clazz) throws NameNotFoundException {
+        installApp(pkg, clazz, new IntentFilter(ACTION_CREATE_SHORTCUT));
+    }
+
+    private void installApp(String pkg, String clazz, IntentFilter filter)
+            throws NameNotFoundException {
+        ShadowPackageManager spm = shadowOf(RuntimeEnvironment.application.getPackageManager());
+        ComponentName cn = new ComponentName(pkg, clazz);
+        spm.addActivityIfNotPresent(cn);
+
+        filter.addCategory(Intent.CATEGORY_DEFAULT);
+        spm.addIntentFilterForActivity(cn, filter);
+    }
+
     /**
      * Loads the model in memory synchronously
      */
     public void loadModelSync() throws ExecutionException, InterruptedException {
-        Callbacks mockCb = new Callbacks() { };
-        Executors.MAIN_EXECUTOR.submit(() -> getModel().addCallbacksAndLoad(mockCb)).get();
+        // Since robolectric tests run on main thread, we run the loader-UI calls on a temp thread,
+        // so that we can wait appropriately for the loader to complete.
+        ShadowLooperExecutor sle = Shadow.extract(Executors.MAIN_EXECUTOR);
+        sle.setHandler(Executors.UI_HELPER_EXECUTOR.getHandler());
+
+        Callbacks mockCb = mock(Callbacks.class);
+        getModel().addCallbacksAndLoad(mockCb);
 
         Executors.MODEL_EXECUTOR.submit(() -> { }).get();
-        Executors.MAIN_EXECUTOR.submit(() -> { }).get();
-        Executors.MAIN_EXECUTOR.submit(() -> getModel().removeCallbacks(mockCb)).get();
+        Executors.UI_HELPER_EXECUTOR.submit(() -> { }).get();
+
+        sle.setHandler(null);
+        getModel().removeCallbacks(mockCb);
     }
 
     /**
@@ -473,97 +437,4 @@ public class LauncherModelHelper {
             return mOpenHelper;
         }
     }
-
-    public static boolean deleteContents(File dir) {
-        File[] files = dir.listFiles();
-        boolean success = true;
-        if (files != null) {
-            for (File file : files) {
-                if (file.isDirectory()) {
-                    success &= deleteContents(file);
-                }
-                if (!file.delete()) {
-                    success = false;
-                }
-            }
-        }
-        return success;
-    }
-
-    public class SanboxModelContext extends SandboxContext {
-
-        private final ArrayMap<String, Object> mSpiedServices = new ArrayMap<>();
-        private final PackageManager mPm;
-        private final File mDbDir;
-
-        SanboxModelContext() {
-            super(ApplicationProvider.getApplicationContext(),
-                    UserCache.INSTANCE, InstallSessionHelper.INSTANCE,
-                    LauncherAppState.INSTANCE, InvariantDeviceProfile.INSTANCE,
-                    DisplayController.INSTANCE, CustomWidgetManager.INSTANCE,
-                    SettingsCache.INSTANCE, PluginManagerWrapper.INSTANCE,
-                    ItemInstallQueue.INSTANCE);
-            mPm = spy(getBaseContext().getPackageManager());
-            mDbDir = new File(getCacheDir(), UUID.randomUUID().toString());
-        }
-
-        public SanboxModelContext allow(MainThreadInitializedObject object) {
-            mAllowedObjects.add(object);
-            return this;
-        }
-
-        @Override
-        public File getDatabasePath(String name) {
-            if (!mDbDir.exists()) {
-                mDbDir.mkdirs();
-            }
-            return new File(mDbDir, name);
-        }
-
-        @Override
-        public ContentResolver getContentResolver() {
-            return mMockResolver;
-        }
-
-        @Override
-        public void onDestroy() {
-            if (deleteContents(mDbDir)) {
-                mDbDir.delete();
-            }
-            super.onDestroy();
-        }
-
-
-        @Override
-        protected <T> T createObject(ObjectProvider<T> provider) {
-            return spy(provider.get(this));
-        }
-
-        @Override
-        public PackageManager getPackageManager() {
-            return mPm;
-        }
-
-        @Override
-        public Object getSystemService(String name) {
-            Object service = mSpiedServices.get(name);
-            return service != null ? service : super.getSystemService(name);
-        }
-
-        public <T> T spyService(Class<T> tClass) {
-            String name = getSystemServiceName(tClass);
-            Object service = mSpiedServices.get(name);
-            if (service != null) {
-                return (T) service;
-            }
-
-            T result = spy(getSystemService(tClass));
-            mSpiedServices.put(name, result);
-            return result;
-        }
-    }
-
-    private static Context testContext() {
-        return getInstrumentation().getContext();
-    }
 }
diff --git a/robolectric_tests/src/com/android/launcher3/util/LauncherUIHelper.java b/robolectric_tests/src/com/android/launcher3/util/LauncherUIHelper.java
new file mode 100644
index 0000000000..caad40e260
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/util/LauncherUIHelper.java
@@ -0,0 +1,101 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.util;
+
+import static android.view.View.MeasureSpec.EXACTLY;
+import static android.view.View.MeasureSpec.makeMeasureSpec;
+
+import static com.android.launcher3.Utilities.createHomeIntent;
+
+import android.app.Activity;
+import android.content.Context;
+import android.content.Intent;
+import android.content.pm.ResolveInfo;
+import android.graphics.Point;
+import android.view.View;
+import android.view.WindowManager;
+
+import com.android.launcher3.Launcher;
+
+import org.robolectric.Robolectric;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.android.controller.ActivityController;
+import org.robolectric.shadows.ShadowLooper;
+import org.robolectric.util.ReflectionHelpers;
+
+import java.util.List;
+
+/**
+ * Utility class to help manage Launcher UI and related objects for test.
+ */
+public class LauncherUIHelper {
+
+    /**
+     * Returns the class name for the Launcher activity as defined in the manifest
+     */
+    public static String getLauncherClassName() {
+        Context context = RuntimeEnvironment.application;
+        Intent homeIntent = createHomeIntent().setPackage(context.getPackageName());
+
+        List<ResolveInfo> launchers = context.getPackageManager()
+                .queryIntentActivities(homeIntent, 0);
+        if (launchers.size() != 1) {
+            return null;
+        }
+        return launchers.get(0).activityInfo.name;
+    }
+
+    /**
+     * Returns an activity controller for Launcher activity defined in the manifest
+     */
+    public static <T extends Launcher> ActivityController<T> buildLauncher() {
+        try {
+            Class<T> tClass = (Class<T>) Class.forName(getLauncherClassName());
+            return Robolectric.buildActivity(tClass);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    /**
+     * Creates and binds a Launcher activity defined in the manifest.
+     * Note that the model must be bound before calling this
+     */
+    public static <T extends Launcher> T buildAndBindLauncher() {
+        ActivityController<T> controller = buildLauncher();
+
+        T launcher = controller.setup().get();
+        doLayout(launcher);
+        ViewOnDrawExecutor executor = ReflectionHelpers.getField(launcher, "mPendingExecutor");
+        if (executor != null) {
+            executor.markCompleted();
+        }
+        return launcher;
+    }
+
+    /**
+     * Performs a measure and layout pass for the given activity
+     */
+    public static void doLayout(Activity activity) {
+        Point size = new Point();
+        RuntimeEnvironment.application.getSystemService(WindowManager.class)
+                .getDefaultDisplay().getSize(size);
+        View view = activity.getWindow().getDecorView();
+        view.measure(makeMeasureSpec(size.x, EXACTLY), makeMeasureSpec(size.y, EXACTLY));
+        view.layout(0, 0, size.x, size.y);
+        ShadowLooper.idleMainLooper();
+    }
+}
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index 50808244ad..702b73afba 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -50,7 +50,7 @@ import com.android.launcher3.util.SimpleBroadcastReceiver;
 import com.android.launcher3.util.Themes;
 import com.android.launcher3.widget.custom.CustomWidgetManager;
 
-public class LauncherAppState implements SafeCloseable {
+public class LauncherAppState {
 
     public static final String ACTION_FORCE_ROLOAD = "force-reload-launcher";
     private static final String KEY_ICON_STATE = "pref_icon_shape_path";
@@ -158,8 +158,7 @@ public class LauncherAppState implements SafeCloseable {
     /**
      * Call from Application.onTerminate(), which is not guaranteed to ever be called.
      */
-    @Override
-    public void close() {
+    public void onTerminate() {
         mModel.destroy();
         mContext.getSystemService(LauncherApps.class).unregisterCallback(mModel);
         CustomWidgetManager.INSTANCE.get(mContext).setWidgetRefreshCallback(null);
diff --git a/src/com/android/launcher3/LauncherModel.java b/src/com/android/launcher3/LauncherModel.java
index 9ebec0a2ea..7b6a5bfeb8 100644
--- a/src/com/android/launcher3/LauncherModel.java
+++ b/src/com/android/launcher3/LauncherModel.java
@@ -96,10 +96,9 @@ public class LauncherModel extends LauncherApps.Callback implements InstallSessi
     // our monitoring of the package manager provides all updates and we never
     // need to do a requery. This is only ever touched from the loader thread.
     private boolean mModelLoaded;
-    private boolean mModelDestroyed = false;
     public boolean isModelLoaded() {
         synchronized (mLock) {
-            return mModelLoaded && mLoaderTask == null && !mModelDestroyed;
+            return mModelLoaded && mLoaderTask == null;
         }
     }
 
@@ -246,7 +245,6 @@ public class LauncherModel extends LauncherApps.Callback implements InstallSessi
      * Called when the model is destroyed
      */
     public void destroy() {
-        mModelDestroyed = true;
         MODEL_EXECUTOR.execute(mModelDelegate::destroy);
     }
 
@@ -559,9 +557,6 @@ public class LauncherModel extends LauncherApps.Callback implements InstallSessi
     }
 
     public void enqueueModelUpdateTask(ModelUpdateTask task) {
-        if (mModelDestroyed) {
-            return;
-        }
         task.init(mApp, this, mBgDataModel, mBgAllAppsList, MAIN_EXECUTOR);
         MODEL_EXECUTOR.execute(task);
     }
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index 1c8954defa..a96de31241 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -86,7 +86,7 @@ import com.android.launcher3.uioverrides.plugins.PluginManagerWrapper;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSet;
-import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext;
+import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.BaseDragLayer;
 import com.android.launcher3.widget.BaseLauncherAppWidgetHostView;
@@ -97,10 +97,13 @@ import com.android.launcher3.widget.NavigableAppWidgetHostView;
 import com.android.launcher3.widget.custom.CustomWidgetManager;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.ConcurrentLinkedQueue;
 
 /**
@@ -119,16 +122,22 @@ public class LauncherPreviewRenderer extends ContextWrapper
      * Context used just for preview. It also provides a few objects (e.g. UserCache) just for
      * preview purposes.
      */
-    public static class PreviewContext extends SandboxContext {
+    public static class PreviewContext extends ContextWrapper {
+
+        private final Set<MainThreadInitializedObject> mAllowedObjects = new HashSet<>(
+                Arrays.asList(UserCache.INSTANCE, InstallSessionHelper.INSTANCE,
+                        LauncherAppState.INSTANCE, InvariantDeviceProfile.INSTANCE,
+                        CustomWidgetManager.INSTANCE, PluginManagerWrapper.INSTANCE));
 
         private final InvariantDeviceProfile mIdp;
+        private final Map<MainThreadInitializedObject, Object> mObjectMap = new HashMap<>();
         private final ConcurrentLinkedQueue<LauncherIconsForPreview> mIconPool =
                 new ConcurrentLinkedQueue<>();
 
+        private boolean mDestroyed = false;
+
         public PreviewContext(Context base, InvariantDeviceProfile idp) {
-            super(base, UserCache.INSTANCE, InstallSessionHelper.INSTANCE,
-                    LauncherAppState.INSTANCE, InvariantDeviceProfile.INSTANCE,
-                    CustomWidgetManager.INSTANCE, PluginManagerWrapper.INSTANCE);
+            super(base);
             mIdp = idp;
             mObjectMap.put(InvariantDeviceProfile.INSTANCE, idp);
             mObjectMap.put(LauncherAppState.INSTANCE,
@@ -136,6 +145,37 @@ public class LauncherPreviewRenderer extends ContextWrapper
 
         }
 
+        @Override
+        public Context getApplicationContext() {
+            return this;
+        }
+
+        public void onDestroy() {
+            CustomWidgetManager.INSTANCE.get(this).onDestroy();
+            LauncherAppState.INSTANCE.get(this).onTerminate();
+            mDestroyed = true;
+        }
+
+        /**
+         * Find a cached object from mObjectMap if we have already created one. If not, generate
+         * an object using the provider.
+         */
+        public <T> T getObject(MainThreadInitializedObject<T> mainThreadInitializedObject,
+                MainThreadInitializedObject.ObjectProvider<T> provider) {
+            if (FeatureFlags.IS_STUDIO_BUILD && mDestroyed) {
+                throw new RuntimeException("Context already destroyed");
+            }
+            if (!mAllowedObjects.contains(mainThreadInitializedObject)) {
+                throw new IllegalStateException("Leaking unknown objects");
+            }
+            if (mObjectMap.containsKey(mainThreadInitializedObject)) {
+                return (T) mObjectMap.get(mainThreadInitializedObject);
+            }
+            T t = provider.get(this);
+            mObjectMap.put(mainThreadInitializedObject, t);
+            return t;
+        }
+
         public LauncherIcons newLauncherIcons(Context context, boolean shapeDetection) {
             LauncherIconsForPreview launcherIconsForPreview = mIconPool.poll();
             if (launcherIconsForPreview != null) {
diff --git a/src/com/android/launcher3/icons/IconCache.java b/src/com/android/launcher3/icons/IconCache.java
index 1a468aeb88..cd13cd0ec5 100644
--- a/src/com/android/launcher3/icons/IconCache.java
+++ b/src/com/android/launcher3/icons/IconCache.java
@@ -134,9 +134,6 @@ public class IconCache extends BaseIconCache {
      * Closes the cache DB. This will clear any in-memory cache.
      */
     public void close() {
-        // This will clear all pending updates
-        getUpdateHandler();
-
         mIconDb.close();
     }
 
diff --git a/src/com/android/launcher3/util/DisplayController.java b/src/com/android/launcher3/util/DisplayController.java
index 94f29dbbf7..e2c0a32bb0 100644
--- a/src/com/android/launcher3/util/DisplayController.java
+++ b/src/com/android/launcher3/util/DisplayController.java
@@ -55,7 +55,7 @@ import java.util.Set;
  * Utility class to cache properties of default display to avoid a system RPC on every call.
  */
 @SuppressLint("NewApi")
-public class DisplayController implements DisplayListener, ComponentCallbacks, SafeCloseable {
+public class DisplayController implements DisplayListener, ComponentCallbacks {
 
     private static final String TAG = "DisplayController";
 
@@ -79,7 +79,6 @@ public class DisplayController implements DisplayListener, ComponentCallbacks, S
 
     private final ArrayList<DisplayInfoChangeListener> mListeners = new ArrayList<>();
     private Info mInfo;
-    private boolean mDestroyed = false;
 
     private DisplayController(Context context) {
         mContext = context;
@@ -111,17 +110,6 @@ public class DisplayController implements DisplayListener, ComponentCallbacks, S
         mDM.registerDisplayListener(this, UI_HELPER_EXECUTOR.getHandler());
     }
 
-    @Override
-    public void close() {
-        mDestroyed = true;
-        if (mWindowContext != null) {
-            mWindowContext.unregisterComponentCallbacks(this);
-        } else {
-            // TODO: unregister broadcast receiver
-        }
-        mDM.unregisterDisplayListener(this);
-    }
-
     @Override
     public final void onDisplayAdded(int displayId) { }
 
@@ -169,9 +157,6 @@ public class DisplayController implements DisplayListener, ComponentCallbacks, S
      * Only used for pre-S
      */
     private void onConfigChanged(Intent intent) {
-        if (mDestroyed) {
-            return;
-        }
         Configuration config = mContext.getResources().getConfiguration();
         if (mInfo.fontScale != config.fontScale || mInfo.densityDpi != config.densityDpi) {
             Log.d(TAG, "Configuration changed, notifying listeners");
diff --git a/src/com/android/launcher3/util/MainThreadInitializedObject.java b/src/com/android/launcher3/util/MainThreadInitializedObject.java
index ef160b199a..f6003dd7bf 100644
--- a/src/com/android/launcher3/util/MainThreadInitializedObject.java
+++ b/src/com/android/launcher3/util/MainThreadInitializedObject.java
@@ -18,20 +18,13 @@ package com.android.launcher3.util;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
 import android.content.Context;
-import android.content.ContextWrapper;
 import android.os.Looper;
 
-import androidx.annotation.UiThread;
 import androidx.annotation.VisibleForTesting;
 
+import com.android.launcher3.graphics.LauncherPreviewRenderer.PreviewContext;
 import com.android.launcher3.util.ResourceBasedOverride.Overrides;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.ExecutionException;
 
 /**
@@ -47,8 +40,8 @@ public class MainThreadInitializedObject<T> {
     }
 
     public T get(Context context) {
-        if (context instanceof SandboxContext) {
-            return ((SandboxContext) context).getObject(this, mProvider);
+        if (context instanceof PreviewContext) {
+            return ((PreviewContext) context).getObject(this, mProvider);
         }
 
         if (mValue == null) {
@@ -87,78 +80,4 @@ public class MainThreadInitializedObject<T> {
 
         T get(Context context);
     }
-
-    /**
-     * Abstract Context which allows custom implementations for
-     * {@link MainThreadInitializedObject} providers
-     */
-    public static abstract class SandboxContext extends ContextWrapper {
-
-        protected final Set<MainThreadInitializedObject> mAllowedObjects;
-        protected final Map<MainThreadInitializedObject, Object> mObjectMap = new HashMap<>();
-        protected final ArrayList<Object> mOrderedObjects = new ArrayList<>();
-
-        private final Object mDestroyLock = new Object();
-        private boolean mDestroyed = false;
-
-        public SandboxContext(Context base, MainThreadInitializedObject... allowedObjects) {
-            super(base);
-            mAllowedObjects = new HashSet<>(Arrays.asList(allowedObjects));
-        }
-
-        @Override
-        public Context getApplicationContext() {
-            return this;
-        }
-
-        public void onDestroy() {
-            synchronized (mDestroyLock) {
-                // Destroy in reverse order
-                for (int i = mOrderedObjects.size() - 1; i >= 0; i--) {
-                    Object o = mOrderedObjects.get(i);
-                    if (o instanceof SafeCloseable) {
-                        ((SafeCloseable) o).close();
-                    }
-                }
-                mDestroyed = true;
-            }
-        }
-
-        /**
-         * Find a cached object from mObjectMap if we have already created one. If not, generate
-         * an object using the provider.
-         */
-        private <T> T getObject(MainThreadInitializedObject<T> object, ObjectProvider<T> provider) {
-            synchronized (mDestroyLock) {
-                if (mDestroyed) {
-                    throw new RuntimeException("Context already destroyed");
-                }
-                if (!mAllowedObjects.contains(object)) {
-                    throw new IllegalStateException(
-                            "Leaking unknown objects " + object + "  " + provider);
-                }
-                T t = (T) mObjectMap.get(object);
-                if (t != null) {
-                    return t;
-                }
-                if (Looper.myLooper() == Looper.getMainLooper()) {
-                    t = createObject(provider);
-                    mObjectMap.put(object, t);
-                    mOrderedObjects.add(t);
-                    return t;
-                }
-            }
-
-            try {
-                return MAIN_EXECUTOR.submit(() -> getObject(object, provider)).get();
-            } catch (InterruptedException | ExecutionException e) {
-                throw new RuntimeException(e);
-            }
-        }
-
-        @UiThread
-        protected <T> T createObject(ObjectProvider<T> provider) {
-            return provider.get(this);
-        }
-    }
 }
diff --git a/src/com/android/launcher3/util/SettingsCache.java b/src/com/android/launcher3/util/SettingsCache.java
index 0c5b7225d3..10611c7211 100644
--- a/src/com/android/launcher3/util/SettingsCache.java
+++ b/src/com/android/launcher3/util/SettingsCache.java
@@ -47,7 +47,7 @@ import java.util.concurrent.CopyOnWriteArrayList;
  *
  * Cache will also be updated if a key queried is missing (even if it has no listeners registered).
  */
-public class SettingsCache extends ContentObserver implements SafeCloseable {
+public class SettingsCache extends ContentObserver {
 
     /** Hidden field Settings.Secure.NOTIFICATION_BADGING */
     public static final Uri NOTIFICATION_BADGING_URI =
@@ -69,6 +69,7 @@ public class SettingsCache extends ContentObserver implements SafeCloseable {
     private final Map<Uri, CopyOnWriteArrayList<OnChangeListener>> mListenerMap = new HashMap<>();
     protected final ContentResolver mResolver;
 
+
     /**
      * Singleton instance
      */
@@ -80,11 +81,6 @@ public class SettingsCache extends ContentObserver implements SafeCloseable {
         mResolver = context.getContentResolver();
     }
 
-    @Override
-    public void close() {
-        mResolver.unregisterContentObserver(this);
-    }
-
     @Override
     public void onChange(boolean selfChange, Uri uri) {
         // We use default of 1, but if we're getting an onChange call, can assume a non-default
diff --git a/src/com/android/launcher3/util/UiThreadHelper.java b/src/com/android/launcher3/util/UiThreadHelper.java
index ac5368c8b4..0f40179f79 100644
--- a/src/com/android/launcher3/util/UiThreadHelper.java
+++ b/src/com/android/launcher3/util/UiThreadHelper.java
@@ -28,7 +28,7 @@ import android.os.Message;
 import android.view.View;
 import android.view.inputmethod.InputMethodManager;
 
-import com.android.launcher3.BaseActivity;
+import com.android.launcher3.Launcher;
 import com.android.launcher3.views.ActivityContext;
 
 /**
@@ -56,7 +56,7 @@ public class UiThreadHelper {
                 STATS_LOGGER_KEY,
                 Message.obtain(
                         HANDLER.get(root.getContext()),
-                        () -> BaseActivity.fromContext(root.getContext())
+                        () -> Launcher.cast(activityContext)
                                 .getStatsLogManager()
                                 .logger()
                                 .log(LAUNCHER_ALLAPPS_KEYBOARD_CLOSED)
diff --git a/src/com/android/launcher3/widget/custom/CustomWidgetManager.java b/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
index 2e2a968a44..329a44452b 100644
--- a/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
+++ b/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
@@ -33,7 +33,6 @@ import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.uioverrides.plugins.PluginManagerWrapper;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.PackageUserKey;
-import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.widget.LauncherAppWidgetHostView;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.systemui.plugins.CustomWidgetPlugin;
@@ -47,7 +46,7 @@ import java.util.stream.Stream;
 /**
  * CustomWidgetManager handles custom widgets implemented as a plugin.
  */
-public class CustomWidgetManager implements PluginListener<CustomWidgetPlugin>, SafeCloseable {
+public class CustomWidgetManager implements PluginListener<CustomWidgetPlugin> {
 
     public static final MainThreadInitializedObject<CustomWidgetManager> INSTANCE =
             new MainThreadInitializedObject<>(CustomWidgetManager::new);
@@ -72,8 +71,7 @@ public class CustomWidgetManager implements PluginListener<CustomWidgetPlugin>,
                 .addPluginListener(this, CustomWidgetPlugin.class, true);
     }
 
-    @Override
-    public void close() {
+    public void onDestroy() {
         PluginManagerWrapper.INSTANCE.get(mContext).removePluginListener(this);
     }
 
diff --git a/tests/src/com/android/launcher3/secondarydisplay/SDLauncherTest.java b/tests/src/com/android/launcher3/secondarydisplay/SDLauncherTest.java
deleted file mode 100644
index fd86cf1144..0000000000
--- a/tests/src/com/android/launcher3/secondarydisplay/SDLauncherTest.java
+++ /dev/null
@@ -1,53 +0,0 @@
-/*
- * Copyright (C) 2021 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.android.launcher3.secondarydisplay;
-
-import static androidx.test.core.app.ActivityScenario.launch;
-
-import androidx.test.core.app.ActivityScenario;
-import androidx.test.espresso.intent.Intents;
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.MediumTest;
-
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-
-/**
- * Tests for {@link SecondaryDisplayLauncher}
- */
-@MediumTest
-@RunWith(AndroidJUnit4.class)
-public class SDLauncherTest {
-
-    @Before
-    public void setUp() {
-        Intents.init();
-    }
-
-    @After
-    public void tearDown() {
-        Intents.release();
-    }
-
-    @Test
-    public void testAllAppsListOpens() {
-        ActivityScenario<SecondaryDisplayLauncher> launcher =
-                launch(SecondaryDisplayLauncher.class);
-        launcher.onActivity(l -> l.showAppDrawer(true));
-    }
-}
diff --git a/tests/src/com/android/launcher3/util/ReflectionHelpers.java b/tests/src/com/android/launcher3/util/ReflectionHelpers.java
deleted file mode 100644
index d89975de8a..0000000000
--- a/tests/src/com/android/launcher3/util/ReflectionHelpers.java
+++ /dev/null
@@ -1,58 +0,0 @@
-/*
- * Copyright (C) 2021 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.launcher3.util;
-
-import java.lang.reflect.Field;
-
-public class ReflectionHelpers {
-
-    /**
-     * Reflectively get the value of a field.
-     *
-     * @param object Target object.
-     * @param fieldName The field name.
-     * @param <R> The return type.
-     * @return Value of the field on the object.
-     */
-    public static <R> R getField(Object object, String fieldName) {
-        try {
-            Field field = object.getClass().getDeclaredField(fieldName);
-            field.setAccessible(true);
-            return (R) field.get(object);
-        } catch (Exception e) {
-            throw new RuntimeException(e);
-        }
-    }
-
-    /**
-     * Reflectively set the value of a field.
-     *
-     * @param object Target object.
-     * @param fieldName The field name.
-     * @param fieldNewValue New value.
-     */
-    public static void setField(Object object, String fieldName, Object fieldNewValue) {
-        try {
-            Field field = object.getClass().getDeclaredField(fieldName);
-            field.setAccessible(true);
-            field.set(object, fieldNewValue);
-        } catch (Exception e) {
-            throw new RuntimeException(e);
-        }
-    }
-
-}
diff --git a/tests/src/com/android/launcher3/util/Wait.java b/tests/src/com/android/launcher3/util/Wait.java
index fe6143c681..8b5e1970de 100644
--- a/tests/src/com/android/launcher3/util/Wait.java
+++ b/tests/src/com/android/launcher3/util/Wait.java
@@ -52,7 +52,7 @@ public class Wait {
             throw new RuntimeException(t);
         }
         Log.d("Wait", "atMost: timed out: " + SystemClock.uptimeMillis());
-        launcher.checkForAnomaly();
+        launcher.checkForAnomaly(false);
         Assert.fail(message.get());
     }
 
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 23aaa25063..7c73f2b471 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -390,7 +390,7 @@ public final class LauncherInstrumentation {
         }
     }
 
-    private String getSystemAnomalyMessage() {
+    private String getSystemAnomalyMessage(boolean ignoreNavmodeChangeStates) {
         try {
             {
                 final StringBuilder sb = new StringBuilder();
@@ -413,16 +413,18 @@ public final class LauncherInstrumentation {
 
             if (hasSystemUiObject("keyguard_status_view")) return "Phone is locked";
 
-            final String visibleApps = mDevice.findObjects(getAnyObjectSelector())
-                    .stream()
-                    .map(LauncherInstrumentation::getApplicationPackageSafe)
-                    .distinct()
-                    .filter(pkg -> pkg != null)
-                    .collect(Collectors.joining(","));
-            if (SYSTEMUI_PACKAGE.equals(visibleApps)) return "Only System UI views are visible";
+            if (!ignoreNavmodeChangeStates) {
+                final String visibleApps = mDevice.findObjects(getAnyObjectSelector())
+                        .stream()
+                        .map(LauncherInstrumentation::getApplicationPackageSafe)
+                        .distinct()
+                        .filter(pkg -> pkg != null)
+                        .collect(Collectors.joining(","));
+                if (SYSTEMUI_PACKAGE.equals(visibleApps)) return "Only System UI views are visible";
 
-            if (!mDevice.wait(Until.hasObject(getAnyObjectSelector()), WAIT_TIME_MS)) {
-                return "Screen is empty";
+                if (!mDevice.wait(Until.hasObject(getAnyObjectSelector()), WAIT_TIME_MS)) {
+                    return "Screen is empty";
+                }
             }
 
             final String navigationModeError = getNavigationModeMismatchError(true);
@@ -434,8 +436,12 @@ public final class LauncherInstrumentation {
         return null;
     }
 
-    public void checkForAnomaly() {
-        final String systemAnomalyMessage = getSystemAnomalyMessage();
+    private void checkForAnomaly() {
+        checkForAnomaly(false);
+    }
+
+    public void checkForAnomaly(boolean ignoreNavmodeChangeStates) {
+        final String systemAnomalyMessage = getSystemAnomalyMessage(ignoreNavmodeChangeStates);
         if (systemAnomalyMessage != null) {
             Assert.fail(formatSystemHealthMessage(formatErrorWithEvents(
                     "http://go/tapl : Tests are broken by a non-Launcher system error: "
@@ -767,7 +773,7 @@ public final class LauncherInstrumentation {
 
                 if (hasLauncherObject(CONTEXT_MENU_RES_ID)) {
                     GestureScope gestureScope = gestureStartFromLauncher
-                            ? (isTablet()? GestureScope.INSIDE : GestureScope.INSIDE_TO_OUTSIDE)
+                            ? (isTablet() ? GestureScope.INSIDE : GestureScope.INSIDE_TO_OUTSIDE)
                             : GestureScope.OUTSIDE_WITH_PILFER;
                     linearGesture(
                             displaySize.x / 2, displaySize.y - 1,
@@ -1027,6 +1033,15 @@ public final class LauncherInstrumentation {
         }
     }
 
+    List<UiObject2> getChildren(UiObject2 container) {
+        try {
+            return container.getChildren();
+        } catch (StaleObjectException e) {
+            fail("The container disappeared from screen");
+            return null;
+        }
+    }
+
     private boolean hasLauncherObject(String resId) {
         return mDevice.hasObject(getLauncherObjectSelector(resId));
     }
diff --git a/tests/tapl/com/android/launcher3/tapl/Widgets.java b/tests/tapl/com/android/launcher3/tapl/Widgets.java
index 99d988926a..6e7264ac0b 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widgets.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widgets.java
@@ -116,9 +116,9 @@ public final class Widgets extends LauncherInstrumentation.VisibleContainer {
                     "widget_preview");
             int i = 0;
             for (; ; ) {
-                final Collection<UiObject2> tableRows = widgetsContainer.getChildren();
+                final Collection<UiObject2> tableRows = mLauncher.getChildren(widgetsContainer);
                 for (UiObject2 row : tableRows) {
-                    final Collection<UiObject2> widgetCells = row.getChildren();
+                    final Collection<UiObject2> widgetCells = mLauncher.getChildren(row);
                     for (UiObject2 widget : widgetCells) {
                         final UiObject2 label = mLauncher.findObjectInContainer(widget,
                                 labelSelector);