diff --git a/quickstep/res/interpolator/app_open_x.xml b/quickstep/res/interpolator/app_open_x.xml
new file mode 100644
index 0000000..5fa0bcb
--- /dev/null
+++ b/quickstep/res/interpolator/app_open_x.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+/*
+** Copyright 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.
+*/
+-->
+
+<pathInterpolator xmlns:android="http://schemas.android.com/apk/res/android"
+    android:pathData="M 0, 0 C 0.1217, 0.0462, 0.15, 0.4686, 0.1667, 0.66 C 0.1834, 0.8878, 0.1667, 1, 1, 1"/>
diff --git a/quickstep/res/interpolator/three_point_fast_out_extra_slow_in.xml b/quickstep/res/interpolator/three_point_fast_out_extra_slow_in.xml
new file mode 100644
index 0000000..70c4231
--- /dev/null
+++ b/quickstep/res/interpolator/three_point_fast_out_extra_slow_in.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+/*
+** Copyright 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.
+*/
+-->
+
+<pathInterpolator xmlns:android="http://schemas.android.com/apk/res/android"
+    android:pathData="M 0,0 C 0.05, 0, 0.133333, 0.06, 0.166666, 0.4 C 0.208333, 0.82, 0.25, 1, 1, 1"/>
diff --git a/quickstep/robolectric_tests/src/com/android/quickstep/RecentsActivityTest.java b/quickstep/robolectric_tests/src/com/android/quickstep/RecentsActivityTest.java
deleted file mode 100644
index 82e8903..0000000
--- a/quickstep/robolectric_tests/src/com/android/quickstep/RecentsActivityTest.java
+++ /dev/null
@@ -1,75 +0,0 @@
-/*
- * Copyright (C) 2020 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.android.quickstep;
-
-import static com.android.launcher3.util.LauncherUIHelper.doLayout;
-
-import android.app.ActivityManager.RunningTaskInfo;
-import android.graphics.Bitmap;
-import android.graphics.Bitmap.Config;
-
-import com.android.quickstep.fallback.FallbackRecentsView;
-import com.android.systemui.shared.recents.model.ThumbnailData;
-
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.robolectric.Robolectric;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.android.controller.ActivityController;
-import org.robolectric.annotation.LooperMode;
-import org.robolectric.annotation.LooperMode.Mode;
-import org.robolectric.shadows.ShadowLooper;
-import org.robolectric.util.ReflectionHelpers;
-
-
-@RunWith(RobolectricTestRunner.class)
-@LooperMode(Mode.PAUSED)
-@org.junit.Ignore
-public class RecentsActivityTest {
-
-    @Test
-    public void testRecentsActivityCreates() {
-        ActivityController<RecentsActivity> controller =
-                Robolectric.buildActivity(RecentsActivity.class);
-
-        RecentsActivity launcher = controller.setup().get();
-        doLayout(launcher);
-
-        // TODO: Ensure that LauncherAppState is not created
-    }
-
-    @Test
-    public void testRecents_showCurrentTask() {
-        ActivityController<RecentsActivity> controller =
-                Robolectric.buildActivity(RecentsActivity.class);
-
-        RecentsActivity activity = controller.setup().get();
-        doLayout(activity);
-
-        FallbackRecentsView frv = activity.getOverviewPanel();
-
-        RunningTaskInfo placeholderTask = new RunningTaskInfo();
-        placeholderTask.taskId = 22;
-        frv.showCurrentTask(new RunningTaskInfo[]{placeholderTask});
-        doLayout(activity);
-
-        ThumbnailData thumbnailData = new ThumbnailData();
-        ReflectionHelpers.setField(thumbnailData, "thumbnail",
-                Bitmap.createBitmap(300, 500, Config.ARGB_8888));
-        frv.switchToScreenshot(thumbnailData, () -> { });
-        ShadowLooper.idleMainLooper();
-    }
-}
diff --git a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
index 433ae3c..f06447b 100644
--- a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
+++ b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
@@ -32,7 +32,6 @@
 import static com.android.launcher3.anim.Interpolators.AGGRESSIVE_EASE;
 import static com.android.launcher3.anim.Interpolators.DEACCEL_1_5;
 import static com.android.launcher3.anim.Interpolators.DEACCEL_1_7;
-import static com.android.launcher3.anim.Interpolators.EXAGGERATED_EASE;
 import static com.android.launcher3.anim.Interpolators.LINEAR;
 import static com.android.launcher3.config.FeatureFlags.ENABLE_SCRIM_FOR_APP_LAUNCH;
 import static com.android.launcher3.config.FeatureFlags.KEYGUARD_ANIMATION;
@@ -71,6 +70,7 @@
 import android.view.View;
 import android.view.ViewRootImpl;
 import android.view.ViewTreeObserver;
+import android.view.animation.AnimationUtils;
 import android.view.animation.Interpolator;
 import android.view.animation.PathInterpolator;
 
@@ -142,21 +142,11 @@
     private static final String CONTROL_REMOTE_APP_TRANSITION_PERMISSION =
             "android.permission.CONTROL_REMOTE_APP_TRANSITION_ANIMATIONS";
 
-    private static final long APP_LAUNCH_DURATION = 450;
-    // Use a shorter duration for x or y translation to create a curve effect
-    private static final long APP_LAUNCH_CURVED_DURATION = 250;
+    private static final long APP_LAUNCH_DURATION = 500;
+
     private static final long APP_LAUNCH_ALPHA_DURATION = 50;
     private static final long APP_LAUNCH_ALPHA_START_DELAY = 25;
 
-    // We scale the durations for the downward app launch animations (minus the scale animation).
-    private static final float APP_LAUNCH_DOWN_DUR_SCALE_FACTOR = 0.8f;
-    private static final long APP_LAUNCH_DOWN_DURATION =
-            (long) (APP_LAUNCH_DURATION * APP_LAUNCH_DOWN_DUR_SCALE_FACTOR);
-    private static final long APP_LAUNCH_DOWN_CURVED_DURATION =
-            (long) (APP_LAUNCH_CURVED_DURATION * APP_LAUNCH_DOWN_DUR_SCALE_FACTOR);
-    private static final long APP_LAUNCH_ALPHA_DOWN_DURATION =
-            (long) (APP_LAUNCH_ALPHA_DURATION * APP_LAUNCH_DOWN_DUR_SCALE_FACTOR);
-
     public static final int ANIMATION_NAV_FADE_IN_DURATION = 266;
     public static final int ANIMATION_NAV_FADE_OUT_DURATION = 133;
     public static final long ANIMATION_DELAY_NAV_FADE_IN =
@@ -166,9 +156,6 @@
     public static final Interpolator NAV_FADE_OUT_INTERPOLATOR =
             new PathInterpolator(0.2f, 0f, 1f, 1f);
 
-    private static final long CROP_DURATION = 375;
-    private static final long RADIUS_DURATION = 375;
-
     public static final int RECENTS_LAUNCH_DURATION = 336;
     private static final int LAUNCHER_RESUME_START_DELAY = 100;
     private static final int CLOSING_TRANSITION_DURATION_MS = 250;
@@ -222,6 +209,9 @@
     // Will never be larger than MAX_NUM_TASKS
     private LinkedHashMap<Integer, Pair<Integer, Integer>> mTaskStartParams;
 
+    private final Interpolator mOpeningXInterpolator;
+    private final Interpolator mOpeningInterpolator;
+
     public QuickstepTransitionManager(Context context) {
         mLauncher = Launcher.cast(Launcher.getLauncher(context));
         mDragLayer = mLauncher.getDragLayer();
@@ -248,6 +238,10 @@
             SystemUiProxy.INSTANCE.get(mLauncher).setStartingWindowListener(
                     mStartingWindowListener);
         }
+
+        mOpeningXInterpolator = AnimationUtils.loadInterpolator(context, R.interpolator.app_open_x);
+        mOpeningInterpolator = AnimationUtils.loadInterpolator(context,
+                R.interpolator.three_point_fast_out_extra_slow_in);
     }
 
     @Override
@@ -686,27 +680,29 @@
         final float finalShadowRadius = appTargetsAreTranslucent ? 0 : mMaxShadowRadius;
 
         MultiValueUpdateListener listener = new MultiValueUpdateListener() {
-            FloatProp mDx = new FloatProp(0, prop.dX, 0, prop.xDuration, AGGRESSIVE_EASE);
-            FloatProp mDy = new FloatProp(0, prop.dY, 0, prop.yDuration, AGGRESSIVE_EASE);
+            FloatProp mDx = new FloatProp(0, prop.dX, 0, APP_LAUNCH_DURATION,
+                    mOpeningXInterpolator);
+            FloatProp mDy = new FloatProp(0, prop.dY, 0, APP_LAUNCH_DURATION,
+                    mOpeningInterpolator);
 
             FloatProp mIconScaleToFitScreen = new FloatProp(prop.initialAppIconScale,
-                    prop.finalAppIconScale, 0, APP_LAUNCH_DURATION, EXAGGERATED_EASE);
+                    prop.finalAppIconScale, 0, APP_LAUNCH_DURATION, mOpeningInterpolator);
             FloatProp mIconAlpha = new FloatProp(prop.iconAlphaStart, 0f,
-                    APP_LAUNCH_ALPHA_START_DELAY, prop.alphaDuration, LINEAR);
+                    APP_LAUNCH_ALPHA_START_DELAY, APP_LAUNCH_ALPHA_DURATION, LINEAR);
 
             FloatProp mWindowRadius = new FloatProp(initialWindowRadius, finalWindowRadius, 0,
-                    RADIUS_DURATION, EXAGGERATED_EASE);
+                    APP_LAUNCH_DURATION, mOpeningInterpolator);
             FloatProp mShadowRadius = new FloatProp(0, finalShadowRadius, 0,
-                    APP_LAUNCH_DURATION, EXAGGERATED_EASE);
+                    APP_LAUNCH_DURATION, mOpeningInterpolator);
 
             FloatProp mCropRectCenterX = new FloatProp(prop.cropCenterXStart, prop.cropCenterXEnd,
-                    0, CROP_DURATION, EXAGGERATED_EASE);
+                    0, APP_LAUNCH_DURATION, mOpeningInterpolator);
             FloatProp mCropRectCenterY = new FloatProp(prop.cropCenterYStart, prop.cropCenterYEnd,
-                    0, CROP_DURATION, EXAGGERATED_EASE);
+                    0, APP_LAUNCH_DURATION, mOpeningInterpolator);
             FloatProp mCropRectWidth = new FloatProp(prop.cropWidthStart, prop.cropWidthEnd, 0,
-                    CROP_DURATION, EXAGGERATED_EASE);
+                    APP_LAUNCH_DURATION, mOpeningInterpolator);
             FloatProp mCropRectHeight = new FloatProp(prop.cropHeightStart, prop.cropHeightEnd, 0,
-                    CROP_DURATION, EXAGGERATED_EASE);
+                    APP_LAUNCH_DURATION, mOpeningInterpolator);
 
             FloatProp mNavFadeOut = new FloatProp(1f, 0f, 0, ANIMATION_NAV_FADE_OUT_DURATION,
                     NAV_FADE_OUT_INTERPOLATOR);
@@ -908,22 +904,23 @@
                     WIDGET_CROSSFADE_DURATION_MILLIS / 2 /* delay */,
                     WIDGET_CROSSFADE_DURATION_MILLIS / 2 /* duration */, LINEAR);
             final FloatProp mWindowRadius = new FloatProp(initialWindowRadius, finalWindowRadius,
-                    0 /* start */, RADIUS_DURATION, LINEAR);
-            final FloatProp mCornerRadiusProgress = new FloatProp(0, 1, 0, RADIUS_DURATION, LINEAR);
+                    0 /* start */, APP_LAUNCH_DURATION, mOpeningInterpolator);
+            final FloatProp mCornerRadiusProgress = new FloatProp(0, 1, 0, APP_LAUNCH_DURATION,
+                    mOpeningInterpolator);
 
             // Window & widget background positioning bounds
             final FloatProp mDx = new FloatProp(widgetBackgroundBounds.centerX(),
-                    windowTargetBounds.centerX(), 0 /* delay */, APP_LAUNCH_CURVED_DURATION,
-                    EXAGGERATED_EASE);
+                    windowTargetBounds.centerX(), 0 /* delay */, APP_LAUNCH_DURATION,
+                    mOpeningXInterpolator);
             final FloatProp mDy = new FloatProp(widgetBackgroundBounds.centerY(),
                     windowTargetBounds.centerY(), 0 /* delay */, APP_LAUNCH_DURATION,
-                    EXAGGERATED_EASE);
+                    mOpeningInterpolator);
             final FloatProp mWidth = new FloatProp(widgetBackgroundBounds.width(),
                     windowTargetBounds.width(), 0 /* delay */, APP_LAUNCH_DURATION,
-                    EXAGGERATED_EASE);
+                    mOpeningInterpolator);
             final FloatProp mHeight = new FloatProp(widgetBackgroundBounds.height(),
                     windowTargetBounds.height(), 0 /* delay */, APP_LAUNCH_DURATION,
-                    EXAGGERATED_EASE);
+                    mOpeningInterpolator);
 
             final FloatProp mNavFadeOut = new FloatProp(1f, 0f, 0, ANIMATION_NAV_FADE_OUT_DURATION,
                     NAV_FADE_OUT_INTERPOLATOR);
@@ -1467,10 +1464,6 @@
         public final float dX;
         public final float dY;
 
-        public final long xDuration;
-        public final long yDuration;
-        public final long alphaDuration;
-
         public final float initialAppIconScale;
         public final float finalAppIconScale;
 
@@ -1502,14 +1495,6 @@
             dX = centerX - launcherIconBounds.centerX();
             dY = centerY - launcherIconBounds.centerY();
 
-            boolean useUpwardAnimation = launcherIconBounds.top > centerY
-                    || Math.abs(dY) < dp.cellHeightPx;
-            xDuration = useUpwardAnimation ? APP_LAUNCH_CURVED_DURATION
-                    : APP_LAUNCH_DOWN_DURATION;
-            yDuration = useUpwardAnimation ? APP_LAUNCH_DURATION
-                    : APP_LAUNCH_DOWN_CURVED_DURATION;
-            alphaDuration = useUpwardAnimation ? APP_LAUNCH_ALPHA_DURATION
-                    : APP_LAUNCH_ALPHA_DOWN_DURATION;
             iconAlphaStart = hasSplashScreen && !hasDifferentAppIcon ? 0 : 1f;
 
             final int windowIconSize = ResourceUtils.getDimenByName("starting_surface_icon_size",
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index 069af96..14bc380 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -84,9 +84,9 @@
 import com.android.quickstep.util.QuickstepOnboardingPrefs;
 import com.android.quickstep.views.RecentsView;
 import com.android.quickstep.views.TaskView;
-import com.android.unfold.UnfoldTransitionFactory;
-import com.android.unfold.UnfoldTransitionProgressProvider;
-import com.android.unfold.config.UnfoldTransitionConfig;
+import com.android.systemui.unfold.UnfoldTransitionFactory;
+import com.android.systemui.unfold.UnfoldTransitionProgressProvider;
+import com.android.systemui.unfold.config.UnfoldTransitionConfig;
 
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
diff --git a/quickstep/src/com/android/quickstep/util/LauncherUnfoldAnimationController.java b/quickstep/src/com/android/quickstep/util/LauncherUnfoldAnimationController.java
index d0fb9e5..c5ab84d 100644
--- a/quickstep/src/com/android/quickstep/util/LauncherUnfoldAnimationController.java
+++ b/quickstep/src/com/android/quickstep/util/LauncherUnfoldAnimationController.java
@@ -24,8 +24,8 @@
 import com.android.launcher3.Hotseat;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.util.HorizontalInsettableView;
-import com.android.unfold.UnfoldTransitionProgressProvider;
-import com.android.unfold.UnfoldTransitionProgressProvider.TransitionProgressListener;
+import com.android.systemui.unfold.UnfoldTransitionProgressProvider;
+import com.android.systemui.unfold.UnfoldTransitionProgressProvider.TransitionProgressListener;
 
 /**
  * Controls animations that are happening during unfolding foldable devices
diff --git a/quickstep/src/com/android/quickstep/util/ProxyScreenStatusProvider.java b/quickstep/src/com/android/quickstep/util/ProxyScreenStatusProvider.java
index 9eda8f4..3777c65 100644
--- a/quickstep/src/com/android/quickstep/util/ProxyScreenStatusProvider.java
+++ b/quickstep/src/com/android/quickstep/util/ProxyScreenStatusProvider.java
@@ -17,7 +17,7 @@
 
 import androidx.annotation.NonNull;
 
-import com.android.unfold.updates.screen.ScreenStatusProvider;
+import com.android.systemui.unfold.updates.screen.ScreenStatusProvider;
 
 import java.util.ArrayList;
 import java.util.List;
diff --git a/quickstep/src/com/android/quickstep/util/UnfoldMoveFromCenterWorkspaceAnimator.java b/quickstep/src/com/android/quickstep/util/UnfoldMoveFromCenterWorkspaceAnimator.java
index 126e044..482092d 100644
--- a/quickstep/src/com/android/quickstep/util/UnfoldMoveFromCenterWorkspaceAnimator.java
+++ b/quickstep/src/com/android/quickstep/util/UnfoldMoveFromCenterWorkspaceAnimator.java
@@ -29,7 +29,7 @@
 import com.android.launcher3.widget.NavigableAppWidgetHostView;
 import com.android.systemui.shared.animation.UnfoldMoveFromCenterAnimator;
 import com.android.systemui.shared.animation.UnfoldMoveFromCenterAnimator.TranslationApplier;
-import com.android.unfold.UnfoldTransitionProgressProvider;
+import com.android.systemui.unfold.UnfoldTransitionProgressProvider;
 
 import java.util.HashMap;
 import java.util.Map;
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index ef52b41..853a3aa 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -4062,6 +4062,7 @@
         setCurrentTask(-1);
         mRecentsAnimationController = null;
         executeSideTaskLaunchCallback();
+        mRemoteTargetHandles = null;
     }
 
     public void setDisallowScrollToClearAll(boolean disallowScrollToClearAll) {
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index 174e1b3..cb379f7 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -654,6 +654,7 @@
             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();
@@ -694,6 +695,7 @@
             // 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);
             }
@@ -1506,6 +1508,10 @@
         mDigitalWellBeingToast.setBannerColorTint(tintColor, amount);
     }
 
+    private int getRootViewDisplayId() {
+        return getRootView().getDisplay().getDisplayId();
+    }
+
     /**
      * We update and subsequently draw these in {@link #setFullscreenProgress(float)}.
      */
diff --git a/robolectric_tests/src/com/android/launcher3/ui/LauncherUIScrollTest.java b/robolectric_tests/src/com/android/launcher3/ui/LauncherUIScrollTest.java
deleted file mode 100644
index ea75548..0000000
--- a/robolectric_tests/src/com/android/launcher3/ui/LauncherUIScrollTest.java
+++ /dev/null
@@ -1,181 +0,0 @@
-/*
- * Copyright (C) 2020 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */package com.android.launcher3.ui;
-
-import static com.android.launcher3.util.LauncherModelHelper.TEST_PACKAGE;
-import static com.android.launcher3.util.LauncherUIHelper.buildAndBindLauncher;
-import static com.android.launcher3.util.LauncherUIHelper.doLayout;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotEquals;
-
-import android.content.Context;
-import android.os.SystemClock;
-import android.provider.Settings;
-import android.view.InputDevice;
-import android.view.MotionEvent;
-import android.view.MotionEvent.PointerProperties;
-
-import com.android.launcher3.DeviceProfile;
-import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.Launcher;
-import com.android.launcher3.LauncherState;
-import com.android.launcher3.folder.Folder;
-import com.android.launcher3.folder.FolderIcon;
-import com.android.launcher3.folder.FolderPagedView;
-import com.android.launcher3.util.LauncherLayoutBuilder;
-import com.android.launcher3.util.LauncherLayoutBuilder.FolderBuilder;
-import com.android.launcher3.util.LauncherModelHelper;
-import com.android.launcher3.widget.picker.WidgetsFullSheet;
-
-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.shadows.ShadowLooper;
-
-/**
- * Tests scroll behavior at various Launcher UI components
- */
-@RunWith(RobolectricTestRunner.class)
-@LooperMode(Mode.PAUSED)
-@org.junit.Ignore
-public class LauncherUIScrollTest {
-
-    private Context mTargetContext;
-    private InvariantDeviceProfile mIdp;
-    private LauncherModelHelper mModelHelper;
-
-    private LauncherLayoutBuilder mLayoutBuilder;
-
-    @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);
-        // LayoutBuilder with 3 workspace pages
-        mLayoutBuilder = new LauncherLayoutBuilder()
-                .atWorkspace(0,  mIdp.numRows - 1, 0).putApp(TEST_PACKAGE, TEST_PACKAGE)
-                .atWorkspace(0,  mIdp.numRows - 1, 1).putApp(TEST_PACKAGE, TEST_PACKAGE)
-                .atWorkspace(0,  mIdp.numRows - 1, 2).putApp(TEST_PACKAGE, TEST_PACKAGE);
-    }
-
-    @Test
-    public void testWorkspacePagesBound() throws Exception {
-        // Verify that the workspace if bound synchronously
-        Launcher launcher = loadLauncher();
-        assertEquals(3, launcher.getWorkspace().getPageCount());
-        assertEquals(0, launcher.getWorkspace().getCurrentPage());
-
-        launcher.dispatchGenericMotionEvent(createScrollEvent(-1));
-        assertNotEquals("Workspace was not scrolled",
-                0, launcher.getWorkspace().getNextPage());
-    }
-
-    @Test
-    public void testAllAppsScroll() throws Exception {
-        // Install 100 apps
-        for (int i = 0; i < 100; i++) {
-            mModelHelper.installApp(TEST_PACKAGE + i);
-        }
-
-        // Bind and open all-apps
-        Launcher launcher = loadLauncher();
-        launcher.getStateManager().goToState(LauncherState.ALL_APPS, false);
-        doLayout(launcher);
-
-        int currentScroll = launcher.getAppsView().getActiveRecyclerView().getCurrentScrollY();
-        launcher.dispatchGenericMotionEvent(createScrollEvent(-1));
-        int newScroll = launcher.getAppsView().getActiveRecyclerView().getCurrentScrollY();
-
-        assertNotEquals("All Apps was not scrolled", currentScroll, newScroll);
-        assertEquals("Workspace was scrolled", 0, launcher.getWorkspace().getNextPage());
-    }
-
-    @Test
-    public void testWidgetsListScroll() throws Exception {
-        // Install 100 widgets
-        for (int i = 0; i < 100; i++) {
-            mModelHelper.installCustomShortcut(TEST_PACKAGE + i, "shortcutProvider");
-        }
-
-        // Bind and open widgets
-        Launcher launcher = loadLauncher();
-        WidgetsFullSheet widgets = WidgetsFullSheet.show(launcher, false);
-        doLayout(launcher);
-
-        int currentScroll = widgets.getRecyclerView().getCurrentScrollY();
-        launcher.dispatchGenericMotionEvent(createScrollEvent(-1));
-        int newScroll = widgets.getRecyclerView().getCurrentScrollY();
-        assertNotEquals("Widgets was not scrolled", currentScroll, newScroll);
-        assertEquals("Workspace was scrolled", 0, launcher.getWorkspace().getNextPage());
-    }
-
-    @Test
-    public void testFolderPageScroll() throws Exception {
-        // Add a folder with multiple icons
-        FolderBuilder fb = mLayoutBuilder.atWorkspace(mIdp.numColumns / 2, mIdp.numRows / 2, 0)
-                .putFolder(0);
-        for (int i = 0; i < 100; i++) {
-            fb.addApp(TEST_PACKAGE, TEST_PACKAGE);
-        }
-
-        // Bind and open folder
-        Launcher launcher = loadLauncher();
-        doLayout(launcher);
-        launcher.getWorkspace().getFirstMatch((i, v) -> v instanceof FolderIcon).performClick();
-        ShadowLooper.idleMainLooper();
-        doLayout(launcher);
-        FolderPagedView folderPages = Folder.getOpen(launcher).getContent();
-
-        assertEquals(0, folderPages.getNextPage());
-        launcher.dispatchGenericMotionEvent(createScrollEvent(-1));
-        assertNotEquals("Folder page was not scrolled", 0, folderPages.getNextPage());
-        assertEquals("Workspace was scrolled", 0, launcher.getWorkspace().getNextPage());
-    }
-
-    private Launcher loadLauncher() throws Exception {
-        mModelHelper.setupDefaultLayoutProvider(mLayoutBuilder).loadModelSync();
-        return buildAndBindLauncher();
-    }
-
-    private static MotionEvent createScrollEvent(int scroll) {
-        DeviceProfile dp = InvariantDeviceProfile.INSTANCE
-                .get(RuntimeEnvironment.application).supportedProfiles.get(0);
-
-        final PointerProperties[] pointerProperties = new PointerProperties[1];
-        pointerProperties[0] = new PointerProperties();
-        pointerProperties[0].id = 0;
-        final MotionEvent.PointerCoords[] coords = new MotionEvent.PointerCoords[1];
-        coords[0] = new MotionEvent.PointerCoords();
-        coords[0].setAxisValue(MotionEvent.AXIS_VSCROLL, scroll);
-        coords[0].x = dp.widthPx / 2;
-        coords[0].y = dp.heightPx / 2;
-
-        final long time = SystemClock.uptimeMillis();
-        return MotionEvent.obtain(time, time, MotionEvent.ACTION_SCROLL, 1,
-                pointerProperties, coords, 0, 0, 1.0f, 1.0f, 0, 0,
-                InputDevice.SOURCE_CLASS_POINTER, 0);
-    }
-
-}
diff --git a/src/com/android/launcher3/BaseActivity.java b/src/com/android/launcher3/BaseActivity.java
index 9778b61..ec96c6d 100644
--- a/src/com/android/launcher3/BaseActivity.java
+++ b/src/com/android/launcher3/BaseActivity.java
@@ -145,6 +145,7 @@
     /**
      * Returns {@link StatsLogManager} for user event logging.
      */
+    @Override
     public StatsLogManager getStatsLogManager() {
         if (mStatsLogManager == null) {
             mStatsLogManager = StatsLogManager.newInstance(this);
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index d162abd..38beeea 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -683,8 +683,9 @@
         CellLayout finalScreen = mWorkspaceScreens.get(finalScreenId);
 
         // If the final screen is empty, convert it to the extra empty screen
-        if (finalScreen.getShortcutsAndWidgets().getChildCount() == 0 &&
-                !finalScreen.isDropPending()) {
+        if (finalScreen != null
+                && finalScreen.getShortcutsAndWidgets().getChildCount() == 0
+                && !finalScreen.isDropPending()) {
             mWorkspaceScreens.remove(finalScreenId);
             mScreenOrder.removeValue(finalScreenId);
 
diff --git a/src/com/android/launcher3/views/ActivityContext.java b/src/com/android/launcher3/views/ActivityContext.java
index b95904e..c822213 100644
--- a/src/com/android/launcher3/views/ActivityContext.java
+++ b/src/com/android/launcher3/views/ActivityContext.java
@@ -27,6 +27,7 @@
 import com.android.launcher3.dot.DotInfo;
 import com.android.launcher3.dragndrop.DragController;
 import com.android.launcher3.folder.FolderIcon;
+import com.android.launcher3.logging.StatsLogManager;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.util.ViewCache;
 
@@ -109,6 +110,10 @@
         return null;
     }
 
+    default StatsLogManager getStatsLogManager() {
+        return StatsLogManager.newInstance((Context) this);
+    }
+
     /**
      * Returns the ActivityContext associated with the given Context.
      */
diff --git a/src/com/android/launcher3/widget/DatabaseWidgetPreviewLoader.java b/src/com/android/launcher3/widget/DatabaseWidgetPreviewLoader.java
index 95c3e1e..aacb9c5 100644
--- a/src/com/android/launcher3/widget/DatabaseWidgetPreviewLoader.java
+++ b/src/com/android/launcher3/widget/DatabaseWidgetPreviewLoader.java
@@ -17,6 +17,7 @@
 
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
+import android.content.Context;
 import android.content.res.Resources;
 import android.graphics.Bitmap;
 import android.graphics.Canvas;
@@ -34,7 +35,6 @@
 
 import androidx.annotation.NonNull;
 
-import com.android.launcher3.BaseActivity;
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.R;
@@ -46,6 +46,7 @@
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.pm.ShortcutConfigActivityInfo;
 import com.android.launcher3.util.Executors;
+import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.widget.util.WidgetSizes;
 
 import java.util.concurrent.ExecutionException;
@@ -56,10 +57,10 @@
 
     private static final String TAG = "WidgetPreviewLoader";
 
-    private final BaseActivity mContext;
+    private final Context mContext;
     private final float mPreviewBoxCornerRadius;
 
-    public DatabaseWidgetPreviewLoader(BaseActivity context) {
+    public DatabaseWidgetPreviewLoader(Context context) {
         mContext = context;
         float previewCornerRadius = RoundedCornerEnforcement.computeEnforcedRadius(context);
         mPreviewBoxCornerRadius = previewCornerRadius > 0
@@ -139,12 +140,13 @@
         int previewWidth;
         int previewHeight;
 
+        DeviceProfile dp = ActivityContext.lookupContext(mContext).getDeviceProfile();
+
         if (widgetPreviewExists && drawable.getIntrinsicWidth() > 0
                 && drawable.getIntrinsicHeight() > 0) {
             previewWidth = drawable.getIntrinsicWidth();
             previewHeight = drawable.getIntrinsicHeight();
         } else {
-            DeviceProfile dp = mContext.getDeviceProfile();
             Size widgetSize = WidgetSizes.getWidgetPaddedSizePx(mContext, info.provider, dp, spanX,
                     spanY);
             previewWidth = widgetSize.getWidth();
@@ -215,7 +217,7 @@
                     Drawable icon = LauncherAppState.getInstance(mContext).getIconCache()
                             .getFullResIcon(info.provider.getPackageName(), info.icon);
                     if (icon != null) {
-                        int appIconSize = mContext.getDeviceProfile().iconSizePx;
+                        int appIconSize = dp.iconSizePx;
                         int iconSize = (int) Math.min(appIconSize * scale,
                                 Math.min(boxRect.width(), boxRect.height()));
 
@@ -248,7 +250,7 @@
 
     private Bitmap generateShortcutPreview(
             ShortcutConfigActivityInfo info, int maxWidth, int maxHeight) {
-        int iconSize = mContext.getDeviceProfile().allAppsIconSizePx;
+        int iconSize = ActivityContext.lookupContext(mContext).getDeviceProfile().allAppsIconSizePx;
         int padding = mContext.getResources()
                 .getDimensionPixelSize(R.dimen.widget_preview_shortcut_padding);
 
diff --git a/src/com/android/launcher3/widget/WidgetCell.java b/src/com/android/launcher3/widget/WidgetCell.java
index 423c66a..f1ac656 100644
--- a/src/com/android/launcher3/widget/WidgetCell.java
+++ b/src/com/android/launcher3/widget/WidgetCell.java
@@ -43,7 +43,6 @@
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
-import com.android.launcher3.BaseActivity;
 import com.android.launcher3.CheckLongPressHelper;
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.Launcher;
@@ -123,7 +122,7 @@
     protected HandlerRunnable mActiveRequest;
     private boolean mAnimatePreview = true;
 
-    protected final BaseActivity mActivity;
+    protected final ActivityContext mActivity;
     private final CheckLongPressHelper mLongPressHelper;
     private final float mEnforcedCornerRadius;
 
@@ -143,8 +142,8 @@
     public WidgetCell(Context context, AttributeSet attrs, int defStyle) {
         super(context, attrs, defStyle);
 
-        mActivity = BaseActivity.fromContext(context);
-        mWidgetPreviewLoader = new DatabaseWidgetPreviewLoader(mActivity);
+        mActivity = ActivityContext.lookupContext(context);
+        mWidgetPreviewLoader = new DatabaseWidgetPreviewLoader(context);
         mLongPressHelper = new CheckLongPressHelper(this);
         mLongPressHelper.setLongPressTimeoutFactor(1);
 
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
index de0d8b8..d52134c 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
@@ -39,13 +39,13 @@
 import androidx.recyclerview.widget.RecyclerView.LayoutParams;
 import androidx.recyclerview.widget.RecyclerView.ViewHolder;
 
-import com.android.launcher3.Launcher;
 import com.android.launcher3.R;
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.recyclerview.ViewHolderBinder;
 import com.android.launcher3.util.LabelComparator;
 import com.android.launcher3.util.PackageUserKey;
+import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.widget.model.WidgetListSpaceEntry;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
@@ -85,7 +85,7 @@
     private static final int VIEW_TYPE_WIDGETS_HEADER = R.id.view_type_widgets_header;
     private static final int VIEW_TYPE_WIDGETS_SEARCH_HEADER = R.id.view_type_widgets_search_header;
 
-    private final Launcher mLauncher;
+    private final Context mContext;
     private final WidgetsDiffReporter mDiffReporter;
     private final SparseArray<ViewHolderBinder> mViewHolderBinders = new SparseArray<>();
     private final WidgetListBaseRowEntryComparator mRowComparator =
@@ -109,7 +109,7 @@
     public WidgetsListAdapter(Context context, LayoutInflater layoutInflater,
             IconCache iconCache, IntSupplier emptySpaceHeightProvider,
             OnClickListener iconClickListener, OnLongClickListener iconLongClickListener) {
-        mLauncher = Launcher.getLauncher(context);
+        mContext = context;
         mDiffReporter = new WidgetsDiffReporter(iconCache, this);
         WidgetsListDrawableFactory listDrawableFactory = new WidgetsListDrawableFactory(context);
 
@@ -340,7 +340,8 @@
 
         if (showWidgets) {
             mWidgetsContentVisiblePackageUserKey = packageUserKey;
-            mLauncher.getStatsLogManager().logger().log(LAUNCHER_WIDGETSTRAY_APP_EXPANDED);
+            ActivityContext.lookupContext(mContext)
+                    .getStatsLogManager().logger().log(LAUNCHER_WIDGETSTRAY_APP_EXPANDED);
         } else {
             mWidgetsContentVisiblePackageUserKey = null;
         }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
index feeb0fe..8c9ff09 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
@@ -110,10 +110,13 @@
 
                 // When preview loads, notify adapter to rebind the item and possibly animate
                 widget.applyFromCellItem(widgetItem, 1f,
-                        bitmap -> holder.getBindingAdapter().notifyItemChanged(
-                                holder.getBindingAdapterPosition(),
-                                Pair.create(widgetItem, bitmap)),
-                        holder.previewCache.get(widgetItem));
+                        bitmap -> {
+                        if (holder.getBindingAdapter() != null) {
+                            holder.getBindingAdapter().notifyItemChanged(
+                                    holder.getBindingAdapterPosition(),
+                                    Pair.create(widgetItem, bitmap));
+                            }
+                        }, holder.previewCache.get(widgetItem));
             }
         }
     }
diff --git a/tests/Android.bp b/tests/Android.bp
index 37231f9..aeddc4c 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -56,6 +56,7 @@
     resource_dirs: ["res"],
     static_libs: [
         "launcher-aosp-tapl",
+        "androidx.test.core",
         "androidx.test.runner",
         "androidx.test.rules",
         "androidx.test.ext.junit",
diff --git a/robolectric_tests/src/com/android/launcher3/folder/FolderNameProviderTest.java b/tests/src/com/android/launcher3/folder/FolderNameProviderTest.java
similarity index 90%
rename from robolectric_tests/src/com/android/launcher3/folder/FolderNameProviderTest.java
rename to tests/src/com/android/launcher3/folder/FolderNameProviderTest.java
index 2a94d9b..23e6235 100644
--- a/robolectric_tests/src/com/android/launcher3/folder/FolderNameProviderTest.java
+++ b/tests/src/com/android/launcher3/folder/FolderNameProviderTest.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.folder;
 
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
@@ -23,6 +25,9 @@
 import android.content.Intent;
 import android.os.UserHandle;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.util.Executors;
@@ -30,15 +35,11 @@
 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.ArrayList;
 
-@RunWith(RobolectricTestRunner.class)
-@LooperMode(Mode.PAUSED)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public final class FolderNameProviderTest {
     private Context mContext;
     private WorkspaceItemInfo mItem1;
@@ -46,7 +47,7 @@
 
     @Before
     public void setUp() {
-        mContext = RuntimeEnvironment.application;
+        mContext = getApplicationContext();
         mItem1 = new WorkspaceItemInfo(new AppInfo(
                 new ComponentName("a.b.c", "a.b.c/a.b.c.d"),
                 "title1",
diff --git a/robolectric_tests/src/com/android/launcher3/logging/FileLogTest.java b/tests/src/com/android/launcher3/logging/FileLogTest.java
similarity index 89%
rename from robolectric_tests/src/com/android/launcher3/logging/FileLogTest.java
rename to tests/src/com/android/launcher3/logging/FileLogTest.java
index 01b23ba..e5f8cec 100644
--- a/robolectric_tests/src/com/android/launcher3/logging/FileLogTest.java
+++ b/tests/src/com/android/launcher3/logging/FileLogTest.java
@@ -1,16 +1,17 @@
 package com.android.launcher3.logging;
 
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
 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.io.File;
 import java.io.PrintWriter;
@@ -20,8 +21,8 @@
 /**
  * Tests for {@link FileLog}
  */
-@RunWith(RobolectricTestRunner.class)
-@LooperMode(Mode.PAUSED)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class FileLogTest {
 
     private File mTempDir;
@@ -29,7 +30,7 @@
     public void setUp() {
         int count = 0;
         do {
-            mTempDir = new File(RuntimeEnvironment.application.getCacheDir(),
+            mTempDir = new File(getApplicationContext().getCacheDir(),
                     "log-test-" + (count++));
         } while (!mTempDir.mkdir());
 
diff --git a/robolectric_tests/src/com/android/launcher3/popup/PopupPopulatorTest.java b/tests/src/com/android/launcher3/popup/PopupPopulatorTest.java
similarity index 95%
rename from robolectric_tests/src/com/android/launcher3/popup/PopupPopulatorTest.java
rename to tests/src/com/android/launcher3/popup/PopupPopulatorTest.java
index 83bf7da..6764e09 100644
--- a/robolectric_tests/src/com/android/launcher3/popup/PopupPopulatorTest.java
+++ b/tests/src/com/android/launcher3/popup/PopupPopulatorTest.java
@@ -16,6 +16,8 @@
 
 package com.android.launcher3.popup;
 
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
 import static com.android.launcher3.popup.PopupPopulator.MAX_SHORTCUTS;
 import static com.android.launcher3.popup.PopupPopulator.NUM_DYNAMIC;
 
@@ -27,10 +29,11 @@
 
 import android.content.pm.ShortcutInfo;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -39,7 +42,8 @@
 /**
  * Tests the sorting and filtering of shortcuts in {@link PopupPopulator}.
  */
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class PopupPopulatorTest {
 
     @Test
@@ -137,7 +141,7 @@
 
     private ShortcutInfo createInfo(boolean isStatic, int rank) {
         ShortcutInfo info = spy(new ShortcutInfo.Builder(
-                RuntimeEnvironment.application, generateId(isStatic, rank))
+                getApplicationContext(), generateId(isStatic, rank))
                 .setRank(rank)
                 .build());
         doReturn(isStatic).when(info).isDeclaredInManifest();
diff --git a/tests/src/com/android/launcher3/util/ActivityContextWrapper.java b/tests/src/com/android/launcher3/util/ActivityContextWrapper.java
new file mode 100644
index 0000000..2618a2e
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/ActivityContextWrapper.java
@@ -0,0 +1,62 @@
+/*
+ * 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 android.content.Context;
+import android.content.ContextWrapper;
+import android.view.ContextThemeWrapper;
+
+import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.views.ActivityContext;
+import com.android.launcher3.views.BaseDragLayer;
+
+/**
+ * {@link ContextWrapper} with internal Launcher interface for testing
+ */
+public class ActivityContextWrapper extends ContextThemeWrapper implements ActivityContext {
+
+    private final DeviceProfile mProfile;
+    private final MyDragLayer mMyDragLayer;
+
+    public ActivityContextWrapper(Context base) {
+        super(base, android.R.style.Theme_DeviceDefault);
+        mProfile = InvariantDeviceProfile.INSTANCE.get(base).getDeviceProfile(base).copy(base);
+        mMyDragLayer = new MyDragLayer(this);
+    }
+
+    @Override
+    public BaseDragLayer getDragLayer() {
+        return mMyDragLayer;
+    }
+
+    @Override
+    public DeviceProfile getDeviceProfile() {
+        return mProfile;
+    }
+
+    private static class MyDragLayer extends BaseDragLayer<ActivityContextWrapper> {
+
+        MyDragLayer(Context context) {
+            super(context, null, 1);
+        }
+
+        @Override
+        public void recreateControllers() {
+            mControllers = new TouchController[0];
+        }
+    }
+}
diff --git a/robolectric_tests/src/com/android/launcher3/util/GridOccupancyTest.java b/tests/src/com/android/launcher3/util/GridOccupancyTest.java
similarity index 92%
rename from robolectric_tests/src/com/android/launcher3/util/GridOccupancyTest.java
rename to tests/src/com/android/launcher3/util/GridOccupancyTest.java
index 2f3fc37..b498077 100644
--- a/robolectric_tests/src/com/android/launcher3/util/GridOccupancyTest.java
+++ b/tests/src/com/android/launcher3/util/GridOccupancyTest.java
@@ -4,14 +4,17 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.robolectric.RobolectricTestRunner;
 
 /**
  * Unit tests for {@link GridOccupancy}
  */
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class GridOccupancyTest {
 
     @Test
diff --git a/robolectric_tests/src/com/android/launcher3/util/IntArrayTest.java b/tests/src/com/android/launcher3/util/IntArrayTest.java
similarity index 86%
rename from robolectric_tests/src/com/android/launcher3/util/IntArrayTest.java
rename to tests/src/com/android/launcher3/util/IntArrayTest.java
index c08e198..a3c7007 100644
--- a/robolectric_tests/src/com/android/launcher3/util/IntArrayTest.java
+++ b/tests/src/com/android/launcher3/util/IntArrayTest.java
@@ -17,14 +17,17 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.robolectric.RobolectricTestRunner;
 
 /**
- * Robolectric unit tests for {@link IntArray}
+ * Unit tests for {@link IntArray}
  */
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class IntArrayTest {
 
     @Test
diff --git a/robolectric_tests/src/com/android/launcher3/util/IntSetTest.java b/tests/src/com/android/launcher3/util/IntSetTest.java
similarity index 91%
rename from robolectric_tests/src/com/android/launcher3/util/IntSetTest.java
rename to tests/src/com/android/launcher3/util/IntSetTest.java
index 7a8c00b..cdb2891 100644
--- a/robolectric_tests/src/com/android/launcher3/util/IntSetTest.java
+++ b/tests/src/com/android/launcher3/util/IntSetTest.java
@@ -21,14 +21,17 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.robolectric.RobolectricTestRunner;
 
 /**
- * Robolectric unit tests for {@link IntSet}
+ * Unit tests for {@link IntSet}
  */
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class IntSetTest {
 
     @Test
diff --git a/tests/src/com/android/launcher3/util/WidgetUtils.java b/tests/src/com/android/launcher3/util/WidgetUtils.java
index 7bc752e..6fc8491 100644
--- a/tests/src/com/android/launcher3/util/WidgetUtils.java
+++ b/tests/src/com/android/launcher3/util/WidgetUtils.java
@@ -15,12 +15,19 @@
  */
 package com.android.launcher3.util;
 
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
+
 import static com.android.launcher3.WorkspaceLayoutManager.FIRST_SCREEN_ID;
 
 import android.appwidget.AppWidgetHost;
+import android.appwidget.AppWidgetManager;
+import android.appwidget.AppWidgetProviderInfo;
+import android.content.ComponentName;
 import android.content.ContentResolver;
 import android.content.Context;
 import android.os.Bundle;
+import android.os.Process;
 
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.model.data.ItemInfo;
@@ -101,4 +108,17 @@
         resolver.insert(LauncherSettings.Favorites.CONTENT_URI,
                 writer.getValues(targetContext));
     }
+
+
+    /**
+     * Creates a {@link AppWidgetProviderInfo} for the provided component name
+     */
+    public static AppWidgetProviderInfo createAppWidgetProviderInfo(ComponentName cn) {
+        AppWidgetProviderInfo info = AppWidgetManager.getInstance(getApplicationContext())
+                .getInstalledProvidersForPackage(
+                        getInstrumentation().getContext().getPackageName(), Process.myUserHandle())
+                .get(0);
+        info.provider = cn;
+        return info;
+    }
 }
diff --git a/tests/src/com/android/launcher3/util/rule/FailureWatcher.java b/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
index 0b60ffc..fae55cc 100644
--- a/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
+++ b/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
@@ -30,6 +30,7 @@
     public FailureWatcher(UiDevice device, LauncherInstrumentation launcher) {
         mDevice = device;
         mLauncher = launcher;
+        Log.d("b/196820244", "FailureWatcher.ctor", new Exception());
     }
 
     @Override
@@ -44,7 +45,9 @@
             @Override
             public void evaluate() throws Throwable {
                 try {
+                    Log.d("b/196820244", "Before evaluate");
                     FailureWatcher.super.apply(base, description).evaluate();
+                    Log.d("b/196820244", "After evaluate");
                 } finally {
                     if (mLauncher.hadNontestEvents()) {
                         throw new AssertionError(
@@ -64,7 +67,9 @@
     }
 
     public static void onError(UiDevice device, Description description, Throwable e) {
+        Log.d("b/196820244", "onError 1");
         if (device == null) return;
+        Log.d("b/196820244", "onError 2");
         final File parentFile = getInstrumentation().getTargetContext().getFilesDir();
         final File sceenshot = new File(parentFile,
                 "TestScreenshot-" + description.getMethodName() + ".png");
diff --git a/robolectric_tests/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfoTest.java b/tests/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfoTest.java
similarity index 96%
rename from robolectric_tests/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfoTest.java
rename to tests/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfoTest.java
index a6f892c..24ae583 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfoTest.java
+++ b/tests/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfoTest.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.widget;
 
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
@@ -25,6 +27,9 @@
 import android.graphics.Point;
 import android.graphics.Rect;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.InvariantDeviceProfile;
 
@@ -32,15 +37,13 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mockito;
-import org.mockito.MockitoAnnotations;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
 
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public final class LauncherAppWidgetProviderInfoTest {
 
     private static final int CELL_SIZE = 50;
@@ -51,8 +54,7 @@
 
     @Before
     public void setUp() {
-        MockitoAnnotations.initMocks(this);
-        mContext = RuntimeEnvironment.application;
+        mContext = getApplicationContext();
     }
 
     @Test
@@ -260,6 +262,7 @@
             return null;
         }).when(profile).getCellSize(any(Point.class));
         Mockito.when(profile.getCellSize()).thenReturn(new Point(CELL_SIZE, CELL_SIZE));
+        Mockito.when(profile.shouldInsetWidgets()).thenReturn(true);
 
         InvariantDeviceProfile idp = new InvariantDeviceProfile();
         List<DeviceProfile> supportedProfiles = new ArrayList<>(idp.supportedProfiles);
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java b/tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
similarity index 95%
rename from robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
rename to tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
index b9f183c..6232938 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
@@ -15,13 +15,16 @@
  */
 package com.android.launcher3.widget.picker;
 
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
+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.Mockito.doAnswer;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyZeroInteractions;
-import static org.robolectric.Shadows.shadowOf;
 
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ComponentName;
@@ -30,6 +33,8 @@
 import android.os.UserHandle;
 
 import androidx.recyclerview.widget.RecyclerView;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.icons.BitmapInfo;
@@ -48,15 +53,12 @@
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
-import org.robolectric.shadows.ShadowPackageManager;
-import org.robolectric.util.ReflectionHelpers;
 
 import java.util.ArrayList;
 import java.util.List;
 
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public final class WidgetsDiffReporterTest {
     private static final String TEST_PACKAGE_PREFIX = "com.android.test";
     private static final WidgetListBaseRowEntryComparator COMPARATOR =
@@ -87,7 +89,7 @@
                 .getComponent().getPackageName())
                 .when(mIconCache).getTitleNoCache(any());
 
-        mContext = RuntimeEnvironment.application;
+        mContext = getApplicationContext();
         mWidgetsDiffReporter = new WidgetsDiffReporter(mIconCache, mAdapter);
         mHeaderA = createWidgetsHeaderEntry(TEST_PACKAGE_PREFIX + "A",
                 /* appName= */ "A", /* numOfWidgets= */ 3);
@@ -294,14 +296,10 @@
     }
 
     private List<WidgetItem> generateWidgetItems(String packageName, int numOfWidgets) {
-        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
         ArrayList<WidgetItem> widgetItems = new ArrayList<>();
         for (int i = 0; i < numOfWidgets; i++) {
             ComponentName cn = ComponentName.createRelative(packageName, ".SampleWidget" + i);
-            AppWidgetProviderInfo widgetInfo = new AppWidgetProviderInfo();
-            widgetInfo.provider = cn;
-            ReflectionHelpers.setField(widgetInfo, "providerInfo",
-                    packageManager.addReceiverIfNotPresent(cn));
+            AppWidgetProviderInfo widgetInfo = createAppWidgetProviderInfo(cn);
 
             WidgetItem widgetItem = new WidgetItem(
                     LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, widgetInfo),
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java b/tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
similarity index 92%
rename from robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
rename to tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
index 12aac8b..44d6964 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
@@ -15,12 +15,13 @@
  */
 package com.android.launcher3.widget.picker;
 
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Matchers.eq;
 import static org.mockito.Matchers.isNull;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.verify;
-import static org.robolectric.Shadows.shadowOf;
 
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ComponentName;
@@ -31,6 +32,8 @@
 import android.view.LayoutInflater;
 
 import androidx.recyclerview.widget.RecyclerView;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.icons.BitmapInfo;
@@ -38,7 +41,9 @@
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
+import com.android.launcher3.util.ActivityContextWrapper;
 import com.android.launcher3.util.PackageUserKey;
+import com.android.launcher3.util.WidgetUtils;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
@@ -50,15 +55,16 @@
 import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
-import org.robolectric.shadows.ShadowPackageManager;
-import org.robolectric.util.ReflectionHelpers;
 
 import java.util.ArrayList;
 import java.util.List;
 
-@RunWith(RobolectricTestRunner.class)
+/**
+ * Unit tests for WidgetsListAdapter
+ * Note that all indices matching are shifted by 1 to account for the empty space at the start.
+ */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public final class WidgetsListAdapterTest {
     private static final String TEST_PACKAGE_PLACEHOLDER = "com.google.test";
 
@@ -74,7 +80,7 @@
     @Before
     public void setup() {
         MockitoAnnotations.initMocks(this);
-        mContext = RuntimeEnvironment.application;
+        mContext = new ActivityContextWrapper(getApplicationContext());
         mTestProfile = new InvariantDeviceProfile();
         mTestProfile.numRows = 5;
         mTestProfile.numColumns = 5;
@@ -100,7 +106,7 @@
         mAdapter.setWidgets(generateSampleMap(1));
         mAdapter.setWidgets(generateSampleMap(2));
 
-        verify(mListener).onItemRangeInserted(eq(1), eq(1));
+        verify(mListener).onItemRangeInserted(eq(2), eq(1));
     }
 
     @Test
@@ -108,7 +114,7 @@
         mAdapter.setWidgets(generateSampleMap(2));
         mAdapter.setWidgets(generateSampleMap(1));
 
-        verify(mListener).onItemRangeRemoved(eq(1), eq(1));
+        verify(mListener).onItemRangeRemoved(eq(2), eq(1));
     }
 
     @Test
@@ -116,7 +122,7 @@
         mAdapter.setWidgets(generateSampleMap(1));
         mAdapter.setWidgets(generateSampleMap(1));
 
-        verify(mListener).onItemRangeChanged(eq(0), eq(1), isNull());
+        verify(mListener).onItemRangeChanged(eq(1), eq(1), isNull());
     }
 
     @Test
@@ -135,7 +141,7 @@
         // THEN the visible entries list becomes:
         // [com.google.test0, com.google.test1, com.google.test1 content, com.google.test2]
         // com.google.test.1 content is inserted into position 2.
-        verify(mListener).onItemRangeInserted(eq(2), eq(1));
+        verify(mListener).onItemRangeInserted(eq(3), eq(1));
     }
 
     @Test
@@ -163,7 +169,7 @@
         mAdapter.setWidgets(allEntries);
 
         // THEN the onItemRangeChanged is invoked for "com.google.test1 content" at index 2.
-        verify(mListener).onItemRangeChanged(eq(2), eq(1), isNull());
+        verify(mListener).onItemRangeChanged(eq(3), eq(1), isNull());
     }
 
     @Test
@@ -194,15 +200,16 @@
                 allAppsWithWidgets.get(6), allAppsWithWidgets.get(7));
         mAdapter.setWidgets(newList);
 
+        // Account for 1st items as empty space
         // Computation logic                           | [Intermediate list during computation]
         // THEN B <> C < 0, removed B from index 1     | [A, E]
-        verify(mListener).onItemRangeRemoved(/* positionStart= */ 1, /* itemCount= */ 1);
+        verify(mListener).onItemRangeRemoved(/* positionStart= */ 2, /* itemCount= */ 1);
         // THEN E <> C > 0, C inserted to index 1      | [A, C, E]
-        verify(mListener).onItemRangeInserted(/* positionStart= */ 1, /* itemCount= */ 1);
-        // THEN E <> D > 0, D inserted to index 2      | [A, C, D, E]
         verify(mListener).onItemRangeInserted(/* positionStart= */ 2, /* itemCount= */ 1);
+        // THEN E <> D > 0, D inserted to index 2      | [A, C, D, E]
+        verify(mListener).onItemRangeInserted(/* positionStart= */ 3, /* itemCount= */ 1);
         // THEN E <> null = -1, E deleted from index 3 | [A, C, D]
-        verify(mListener).onItemRangeRemoved(/* positionStart= */ 3, /* itemCount= */ 1);
+        verify(mListener).onItemRangeRemoved(/* positionStart= */ 4, /* itemCount= */ 1);
     }
 
     @Test
@@ -225,8 +232,8 @@
 
         // THEN expanded app is reset and the visible entries list becomes:
         // [com.google.test0, com.google.test1, com.google.test2]
-        verify(mListener).onItemRangeChanged(eq(1), eq(1), isNull());
-        verify(mListener).onItemRangeRemoved(/* positionStart= */ 2, /* itemCount= */ 1);
+        verify(mListener).onItemRangeChanged(eq(2), eq(1), isNull());
+        verify(mListener).onItemRangeRemoved(/* positionStart= */ 3, /* itemCount= */ 1);
     }
 
     /**
@@ -263,14 +270,10 @@
     }
 
     private List<WidgetItem> generateWidgetItems(String packageName, int numOfWidgets) {
-        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
         ArrayList<WidgetItem> widgetItems = new ArrayList<>();
         for (int i = 0; i < numOfWidgets; i++) {
             ComponentName cn = ComponentName.createRelative(packageName, ".SampleWidget" + i);
-            AppWidgetProviderInfo widgetInfo = new AppWidgetProviderInfo();
-            widgetInfo.provider = cn;
-            ReflectionHelpers.setField(widgetInfo, "providerInfo",
-                    packageManager.addReceiverIfNotPresent(cn));
+            AppWidgetProviderInfo widgetInfo = WidgetUtils.createAppWidgetProviderInfo(cn);
 
             widgetItems.add(new WidgetItem(
                     LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, widgetInfo),
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java b/tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
similarity index 76%
rename from robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
rename to tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
index fa000c0..969c12a 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
@@ -15,13 +15,14 @@
  */
 package com.android.launcher3.widget.picker;
 
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
 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.verify;
-import static org.robolectric.Shadows.shadowOf;
 
 import static java.util.Collections.EMPTY_LIST;
 
@@ -33,7 +34,9 @@
 import android.widget.FrameLayout;
 import android.widget.TextView;
 
-import com.android.launcher3.DeviceProfile;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.R;
 import com.android.launcher3.icons.BitmapInfo;
@@ -41,28 +44,23 @@
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
-import com.android.launcher3.testing.TestActivity;
+import com.android.launcher3.util.ActivityContextWrapper;
 import com.android.launcher3.util.PackageUserKey;
+import com.android.launcher3.util.WidgetUtils;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
 
-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.Robolectric;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
-import org.robolectric.android.controller.ActivityController;
-import org.robolectric.shadows.ShadowPackageManager;
-import org.robolectric.util.ReflectionHelpers;
 
 import java.util.ArrayList;
 import java.util.List;
 
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public final class WidgetsListHeaderViewHolderBinderTest {
     private static final String TEST_PACKAGE = "com.google.test";
     private static final String APP_NAME = "Test app";
@@ -70,49 +68,35 @@
     private Context mContext;
     private WidgetsListHeaderViewHolderBinder mViewHolderBinder;
     private InvariantDeviceProfile mTestProfile;
-    // Replace ActivityController with ActivityScenario, which is the recommended way for activity
-    // testing.
-    private ActivityController<TestActivity> mActivityController;
-    private TestActivity mTestActivity;
 
     @Mock
     private IconCache mIconCache;
     @Mock
-    private DeviceProfile mDeviceProfile;
-    @Mock
     private OnHeaderClickListener mOnHeaderClickListener;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mContext = RuntimeEnvironment.application;
+
+        mContext = new ActivityContextWrapper(getApplicationContext());
         mTestProfile = new InvariantDeviceProfile();
         mTestProfile.numRows = 5;
         mTestProfile.numColumns = 5;
 
-        mActivityController = Robolectric.buildActivity(TestActivity.class);
-        mTestActivity = mActivityController.setup().get();
-        mTestActivity.setDeviceProfile(mDeviceProfile);
-
         doAnswer(invocation -> {
             ComponentWithLabel componentWithLabel = (ComponentWithLabel) invocation.getArgument(0);
             return componentWithLabel.getComponent().getShortClassName();
         }).when(mIconCache).getTitleNoCache(any());
         mViewHolderBinder = new WidgetsListHeaderViewHolderBinder(
-                LayoutInflater.from(mTestActivity),
+                LayoutInflater.from(mContext),
                 mOnHeaderClickListener,
-                new WidgetsListDrawableFactory(mTestActivity));
-    }
-
-    @After
-    public void tearDown() {
-        mActivityController.destroy();
+                new WidgetsListDrawableFactory(mContext));
     }
 
     @Test
     public void bindViewHolder_appWith3Widgets_shouldShowTheCorrectAppNameAndSubtitle() {
         WidgetsListHeaderHolder viewHolder = mViewHolderBinder.newViewHolder(
-                new FrameLayout(mTestActivity));
+                new FrameLayout(mContext));
         WidgetsListHeader widgetsListHeader = viewHolder.mWidgetsListHeader;
         WidgetsListHeaderEntry entry = generateSampleAppHeader(
                 APP_NAME,
@@ -129,7 +113,7 @@
     @Test
     public void bindViewHolder_shouldAttachOnHeaderClickListener() {
         WidgetsListHeaderHolder viewHolder = mViewHolderBinder.newViewHolder(
-                new FrameLayout(mTestActivity));
+                new FrameLayout(mContext));
         WidgetsListHeader widgetsListHeader = viewHolder.mWidgetsListHeader;
         WidgetsListHeaderEntry entry = generateSampleAppHeader(
                 APP_NAME,
@@ -155,14 +139,10 @@
     }
 
     private List<WidgetItem> generateWidgetItems(String packageName, int numOfWidgets) {
-        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
         ArrayList<WidgetItem> widgetItems = new ArrayList<>();
         for (int i = 0; i < numOfWidgets; i++) {
             ComponentName cn = ComponentName.createRelative(packageName, ".SampleWidget" + i);
-            AppWidgetProviderInfo widgetInfo = new AppWidgetProviderInfo();
-            widgetInfo.provider = cn;
-            ReflectionHelpers.setField(widgetInfo, "providerInfo",
-                    packageManager.addReceiverIfNotPresent(cn));
+            AppWidgetProviderInfo widgetInfo = WidgetUtils.createAppWidgetProviderInfo(cn);
 
             widgetItems.add(new WidgetItem(
                     LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, widgetInfo),
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java b/tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java
similarity index 77%
rename from robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java
rename to tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java
index b18c8b7..453f4fb 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java
@@ -15,13 +15,14 @@
  */
 package com.android.launcher3.widget.picker;
 
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
 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.verify;
-import static org.robolectric.Shadows.shadowOf;
 
 import static java.util.Collections.EMPTY_LIST;
 
@@ -33,7 +34,9 @@
 import android.widget.FrameLayout;
 import android.widget.TextView;
 
-import com.android.launcher3.DeviceProfile;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.R;
 import com.android.launcher3.icons.BitmapInfo;
@@ -41,28 +44,23 @@
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
-import com.android.launcher3.testing.TestActivity;
+import com.android.launcher3.util.ActivityContextWrapper;
 import com.android.launcher3.util.PackageUserKey;
+import com.android.launcher3.util.WidgetUtils;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.model.WidgetsListSearchHeaderEntry;
 
-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.Robolectric;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
-import org.robolectric.android.controller.ActivityController;
-import org.robolectric.shadows.ShadowPackageManager;
-import org.robolectric.util.ReflectionHelpers;
 
 import java.util.ArrayList;
 import java.util.List;
 
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public final class WidgetsListSearchHeaderViewHolderBinderTest {
     private static final String TEST_PACKAGE = "com.google.test";
     private static final String APP_NAME = "Test app";
@@ -70,49 +68,34 @@
     private Context mContext;
     private WidgetsListSearchHeaderViewHolderBinder mViewHolderBinder;
     private InvariantDeviceProfile mTestProfile;
-    // Replace ActivityController with ActivityScenario, which is the recommended way for activity
-    // testing.
-    private ActivityController<TestActivity> mActivityController;
-    private TestActivity mTestActivity;
 
     @Mock
     private IconCache mIconCache;
     @Mock
-    private DeviceProfile mDeviceProfile;
-    @Mock
     private OnHeaderClickListener mOnHeaderClickListener;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mContext = RuntimeEnvironment.application;
+        mContext = new ActivityContextWrapper(getApplicationContext());
         mTestProfile = new InvariantDeviceProfile();
         mTestProfile.numRows = 5;
         mTestProfile.numColumns = 5;
 
-        mActivityController = Robolectric.buildActivity(TestActivity.class);
-        mTestActivity = mActivityController.setup().get();
-        mTestActivity.setDeviceProfile(mDeviceProfile);
-
         doAnswer(invocation -> {
             ComponentWithLabel componentWithLabel = (ComponentWithLabel) invocation.getArgument(0);
             return componentWithLabel.getComponent().getShortClassName();
         }).when(mIconCache).getTitleNoCache(any());
         mViewHolderBinder = new WidgetsListSearchHeaderViewHolderBinder(
-                LayoutInflater.from(mTestActivity),
+                LayoutInflater.from(mContext),
                 mOnHeaderClickListener,
-                new WidgetsListDrawableFactory(mTestActivity));
-    }
-
-    @After
-    public void tearDown() {
-        mActivityController.destroy();
+                new WidgetsListDrawableFactory(mContext));
     }
 
     @Test
     public void bindViewHolder_appWith3Widgets_shouldShowTheCorrectAppNameAndSubtitle() {
         WidgetsListSearchHeaderHolder viewHolder = mViewHolderBinder.newViewHolder(
-                new FrameLayout(mTestActivity));
+                new FrameLayout(mContext));
         WidgetsListHeader widgetsListHeader = viewHolder.mWidgetsListHeader;
         WidgetsListSearchHeaderEntry entry = generateSampleSearchHeader(
                 APP_NAME,
@@ -130,7 +113,7 @@
     @Test
     public void bindViewHolder_shouldAttachOnHeaderClickListener() {
         WidgetsListSearchHeaderHolder viewHolder = mViewHolderBinder.newViewHolder(
-                new FrameLayout(mTestActivity));
+                new FrameLayout(mContext));
         WidgetsListHeader widgetsListHeader = viewHolder.mWidgetsListHeader;
         WidgetsListSearchHeaderEntry entry = generateSampleSearchHeader(
                 APP_NAME,
@@ -156,14 +139,10 @@
     }
 
     private List<WidgetItem> generateWidgetItems(String packageName, int numOfWidgets) {
-        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
         ArrayList<WidgetItem> widgetItems = new ArrayList<>();
         for (int i = 0; i < numOfWidgets; i++) {
             ComponentName cn = ComponentName.createRelative(packageName, ".SampleWidget" + i);
-            AppWidgetProviderInfo widgetInfo = new AppWidgetProviderInfo();
-            widgetInfo.provider = cn;
-            ReflectionHelpers.setField(widgetInfo, "providerInfo",
-                    packageManager.addReceiverIfNotPresent(cn));
+            AppWidgetProviderInfo widgetInfo = WidgetUtils.createAppWidgetProviderInfo(cn);
 
             widgetItems.add(new WidgetItem(
                     LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, widgetInfo),
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java b/tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
similarity index 74%
rename from robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
rename to tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
index cb38c6f..5816b77 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
@@ -15,13 +15,12 @@
  */
 package com.android.launcher3.widget.picker;
 
-import static android.os.Looper.getMainLooper;
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
 
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
-import static org.robolectric.Shadows.shadowOf;
 
 import static java.util.Collections.EMPTY_LIST;
 
@@ -37,7 +36,9 @@
 import android.widget.TableRow;
 import android.widget.TextView;
 
-import com.android.launcher3.DeviceProfile;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.R;
 import com.android.launcher3.icons.BitmapInfo;
@@ -45,29 +46,24 @@
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
-import com.android.launcher3.testing.TestActivity;
-import com.android.launcher3.widget.DatabaseWidgetPreviewLoader;
+import com.android.launcher3.util.ActivityContextWrapper;
+import com.android.launcher3.util.Executors;
+import com.android.launcher3.util.WidgetUtils;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.WidgetCell;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
 
-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.Robolectric;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
-import org.robolectric.android.controller.ActivityController;
-import org.robolectric.shadows.ShadowPackageManager;
-import org.robolectric.util.ReflectionHelpers;
 
 import java.util.ArrayList;
 import java.util.List;
 
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public final class WidgetsListTableViewHolderBinderTest {
     private static final String TEST_PACKAGE = "com.google.test";
     private static final String APP_NAME = "Test app";
@@ -75,10 +71,6 @@
     private Context mContext;
     private WidgetsListTableViewHolderBinder mViewHolderBinder;
     private InvariantDeviceProfile mTestProfile;
-    // Replace ActivityController with ActivityScenario, which is the recommended way for activity
-    // testing.
-    private ActivityController<TestActivity> mActivityController;
-    private TestActivity mTestActivity;
 
     @Mock
     private OnLongClickListener mOnLongClickListener;
@@ -86,50 +78,37 @@
     private OnClickListener mOnIconClickListener;
     @Mock
     private IconCache mIconCache;
-    @Mock
-    private DatabaseWidgetPreviewLoader mWidgetPreviewLoader;
-    @Mock
-    private DeviceProfile mDeviceProfile;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mContext = RuntimeEnvironment.application;
+        mContext = new ActivityContextWrapper(getApplicationContext());
         mTestProfile = new InvariantDeviceProfile();
         mTestProfile.numRows = 5;
         mTestProfile.numColumns = 5;
 
-        mActivityController = Robolectric.buildActivity(TestActivity.class);
-        mTestActivity = mActivityController.setup().get();
-        mTestActivity.setDeviceProfile(mDeviceProfile);
-
         doAnswer(invocation -> {
             ComponentWithLabel componentWithLabel = (ComponentWithLabel) invocation.getArgument(0);
             return componentWithLabel.getComponent().getShortClassName();
         }).when(mIconCache).getTitleNoCache(any());
 
         mViewHolderBinder = new WidgetsListTableViewHolderBinder(
-                LayoutInflater.from(mTestActivity),
+                LayoutInflater.from(mContext),
                 mOnIconClickListener,
                 mOnLongClickListener,
-                new WidgetsListDrawableFactory(mTestActivity));
-    }
-
-    @After
-    public void tearDown() {
-        mActivityController.destroy();
+                new WidgetsListDrawableFactory(mContext));
     }
 
     @Test
-    public void bindViewHolder_appWith3Widgets_shouldHave3Widgets() {
+    public void bindViewHolder_appWith3Widgets_shouldHave3Widgets() throws Exception {
         WidgetsRowViewHolder viewHolder = mViewHolderBinder.newViewHolder(
-                new FrameLayout(mTestActivity));
+                new FrameLayout(mContext));
         WidgetsListContentEntry entry = generateSampleAppWithWidgets(
                 APP_NAME,
                 TEST_PACKAGE,
                 /* numOfWidgets= */ 3);
         mViewHolderBinder.bindViewHolder(viewHolder, entry, /* position= */ 0, EMPTY_LIST);
-        shadowOf(getMainLooper()).idle();
+        Executors.MAIN_EXECUTOR.submit(() -> { }).get();
 
         // THEN the table container has one row, which contains 3 widgets.
         // View:  .SampleWidget0 | .SampleWidget1 | .SampleWidget2
@@ -152,18 +131,15 @@
 
         return new WidgetsListContentEntry(appInfo,
                 /* titleSectionName= */ "",
-                generateWidgetItems(packageName, numOfWidgets));
+                generateWidgetItems(packageName, numOfWidgets),
+                Integer.MAX_VALUE);
     }
 
     private List<WidgetItem> generateWidgetItems(String packageName, int numOfWidgets) {
-        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
         ArrayList<WidgetItem> widgetItems = new ArrayList<>();
         for (int i = 0; i < numOfWidgets; i++) {
             ComponentName cn = ComponentName.createRelative(packageName, ".SampleWidget" + i);
-            AppWidgetProviderInfo widgetInfo = new AppWidgetProviderInfo();
-            widgetInfo.provider = cn;
-            ReflectionHelpers.setField(widgetInfo, "providerInfo",
-                    packageManager.addReceiverIfNotPresent(cn));
+            AppWidgetProviderInfo widgetInfo = WidgetUtils.createAppWidgetProviderInfo(cn);
 
             widgetItems.add(new WidgetItem(
                     LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, widgetInfo),
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java b/tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
similarity index 92%
rename from robolectric_tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
rename to tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
index 106cac0..4b61b2c 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
@@ -15,15 +15,20 @@
  */
 package com.android.launcher3.widget.picker.model;
 
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
+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.Mockito.doAnswer;
-import static org.robolectric.Shadows.shadowOf;
 
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ComponentName;
-import android.content.Context;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.icons.ComponentWithLabel;
@@ -38,16 +43,13 @@
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
-import org.robolectric.shadows.ShadowPackageManager;
-import org.robolectric.util.ReflectionHelpers;
 
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public final class WidgetsListContentEntryTest {
     private static final String PACKAGE_NAME = "com.android.test";
     private static final String PACKAGE_NAME_2 = "com.android.test2";
@@ -60,7 +62,6 @@
 
     @Mock private IconCache mIconCache;
 
-    private Context mContext;
     private InvariantDeviceProfile mTestProfile;
 
     @Before
@@ -71,7 +72,6 @@
         mWidgetsToLabels.put(mWidget2, "Dog");
         mWidgetsToLabels.put(mWidget3, "Bird");
 
-        mContext = RuntimeEnvironment.application;
         mTestProfile = new InvariantDeviceProfile();
         mTestProfile.numRows = 5;
         mTestProfile.numColumns = 5;
@@ -242,17 +242,12 @@
         assertThat(widgetsListRowEntry1.equals(widgetsListRowEntry2)).isTrue();
     }
 
-
     private WidgetItem createWidgetItem(ComponentName componentName, int spanX, int spanY) {
         String label = mWidgetsToLabels.get(componentName);
-        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
-        AppWidgetProviderInfo widgetInfo = new AppWidgetProviderInfo();
-        widgetInfo.provider = componentName;
-        ReflectionHelpers.setField(widgetInfo, "providerInfo",
-                packageManager.addReceiverIfNotPresent(componentName));
+        AppWidgetProviderInfo widgetInfo = createAppWidgetProviderInfo(componentName);
 
         LauncherAppWidgetProviderInfo launcherAppWidgetProviderInfo =
-                LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, widgetInfo);
+                LauncherAppWidgetProviderInfo.fromProviderInfo(getApplicationContext(), widgetInfo);
         launcherAppWidgetProviderInfo.spanX = spanX;
         launcherAppWidgetProviderInfo.spanY = spanY;
         launcherAppWidgetProviderInfo.label = label;
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java b/tests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
similarity index 90%
rename from robolectric_tests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
rename to tests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
index 36b6f01..c862d6b 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
@@ -16,7 +16,10 @@
 
 package com.android.launcher3.widget.picker.search;
 
-import static android.os.Looper.getMainLooper;
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
+import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
+import static com.android.launcher3.util.WidgetUtils.createAppWidgetProviderInfo;
 
 import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
@@ -25,7 +28,6 @@
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.verify;
-import static org.robolectric.Shadows.shadowOf;
 
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ComponentName;
@@ -33,6 +35,9 @@
 import android.graphics.Bitmap;
 import android.os.UserHandle;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.icons.BitmapInfo;
 import com.android.launcher3.icons.ComponentWithLabel;
@@ -52,16 +57,13 @@
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
-import org.robolectric.shadows.ShadowPackageManager;
-import org.robolectric.util.ReflectionHelpers;
 
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class SimpleWidgetsSearchAlgorithmTest {
 
     @Mock private IconCache mIconCache;
@@ -82,7 +84,7 @@
     private SearchCallback<WidgetsListBaseEntry> mSearchCallback;
 
     @Before
-    public void setUp() {
+    public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
         doAnswer(invocation -> {
             ComponentWithLabel componentWithLabel = (ComponentWithLabel) invocation.getArgument(0);
@@ -91,7 +93,7 @@
         mTestProfile = new InvariantDeviceProfile();
         mTestProfile.numRows = 5;
         mTestProfile.numColumns = 5;
-        mContext = RuntimeEnvironment.application;
+        mContext = getApplicationContext();
 
         mCalendarHeaderEntry =
                 createWidgetsHeaderEntry("com.example.android.Calendar", "Calendar", 2);
@@ -102,8 +104,8 @@
         mClockHeaderEntry = createWidgetsHeaderEntry("com.example.android.Clock", "Clock", 3);
         mClockContentEntry = createWidgetsContentEntry("com.example.android.Clock", "Clock", 3);
 
-
-        mSimpleWidgetsSearchAlgorithm = new SimpleWidgetsSearchAlgorithm(mDataProvider);
+        mSimpleWidgetsSearchAlgorithm = MAIN_EXECUTOR.submit(
+                () -> new SimpleWidgetsSearchAlgorithm(mDataProvider)).get();
         doReturn(Collections.EMPTY_LIST).when(mDataProvider).getAllWidgets();
     }
 
@@ -156,13 +158,13 @@
     }
 
     @Test
-    public void doSearch_shouldInformCallback() {
+    public void doSearch_shouldInformCallback() throws Exception {
         doReturn(List.of(mCalendarHeaderEntry, mCalendarContentEntry, mCameraHeaderEntry,
                 mCameraContentEntry, mClockHeaderEntry, mClockContentEntry))
                 .when(mDataProvider)
                 .getAllWidgets();
         mSimpleWidgetsSearchAlgorithm.doSearch("Ca", mSearchCallback);
-        shadowOf(getMainLooper()).idle();
+        MAIN_EXECUTOR.submit(() -> { }).get();
         verify(mSearchCallback).onSearchResult(
                 matches("Ca"), argThat(a -> a != null && !a.isEmpty()));
     }
@@ -195,14 +197,10 @@
     }
 
     private List<WidgetItem> generateWidgetItems(String packageName, int numOfWidgets) {
-        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
         ArrayList<WidgetItem> widgetItems = new ArrayList<>();
         for (int i = 0; i < numOfWidgets; i++) {
             ComponentName cn = ComponentName.createRelative(packageName, ".SampleWidget" + i);
-            AppWidgetProviderInfo widgetInfo = new AppWidgetProviderInfo();
-            widgetInfo.provider = cn;
-            ReflectionHelpers.setField(widgetInfo, "providerInfo",
-                    packageManager.addReceiverIfNotPresent(cn));
+            AppWidgetProviderInfo widgetInfo = createAppWidgetProviderInfo(cn);
 
             WidgetItem widgetItem = new WidgetItem(
                     LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, widgetInfo),
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/search/WidgetsSearchBarControllerTest.java b/tests/src/com/android/launcher3/widget/picker/search/WidgetsSearchBarControllerTest.java
similarity index 83%
rename from robolectric_tests/src/com/android/launcher3/widget/picker/search/WidgetsSearchBarControllerTest.java
rename to tests/src/com/android/launcher3/widget/picker/search/WidgetsSearchBarControllerTest.java
index 7ac879a..583d37f 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/search/WidgetsSearchBarControllerTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/search/WidgetsSearchBarControllerTest.java
@@ -16,39 +16,38 @@
 
 package com.android.launcher3.widget.picker.search;
 
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
 import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 
+import android.content.Context;
 import android.view.View;
 import android.widget.ImageButton;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
 import com.android.launcher3.ExtendedEditText;
 import com.android.launcher3.search.SearchAlgorithm;
-import com.android.launcher3.testing.TestActivity;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 
-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.Robolectric;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.android.controller.ActivityController;
 
 import java.util.ArrayList;
 
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class WidgetsSearchBarControllerTest {
 
     private WidgetsSearchBarController mController;
-    // TODO: Replace ActivityController with ActivityScenario, which is the recommended way for
-    // activity testing.
-    private ActivityController<TestActivity> mActivityController;
     private ExtendedEditText mEditText;
     private ImageButton mCancelButton;
     @Mock
@@ -59,20 +58,14 @@
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mActivityController = Robolectric.buildActivity(TestActivity.class);
-        TestActivity testActivity = mActivityController.setup().get();
+        Context context = getApplicationContext();
 
-        mEditText = new ExtendedEditText(testActivity);
-        mCancelButton = new ImageButton(testActivity);
+        mEditText = new ExtendedEditText(context);
+        mCancelButton = new ImageButton(context);
         mController = new WidgetsSearchBarController(
                 mSearchAlgorithm, mEditText, mCancelButton, mSearchModeListener);
     }
 
-    @After
-    public void tearDown() {
-        mActivityController.destroy();
-    }
-
     @Test
     public void onSearchResult_shouldInformSearchModeListener() {
         ArrayList<WidgetsListBaseEntry> entries = new ArrayList<>();
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java b/tests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java
similarity index 89%
rename from robolectric_tests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java
rename to tests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java
index 56d7d68..d6da776 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java
@@ -15,11 +15,14 @@
  */
 package com.android.launcher3.widget.picker.util;
 
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
+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.Mockito.doAnswer;
-import static org.robolectric.Shadows.shadowOf;
 
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ComponentName;
@@ -29,6 +32,9 @@
 import android.graphics.drawable.Drawable;
 import android.os.UserHandle;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.icons.ComponentWithLabel;
 import com.android.launcher3.icons.IconCache;
@@ -42,15 +48,12 @@
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
-import org.robolectric.shadows.ShadowPackageManager;
-import org.robolectric.util.ReflectionHelpers;
 
 import java.util.ArrayList;
 import java.util.List;
 
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public final class WidgetsTableUtilsTest {
     private static final String TEST_PACKAGE = "com.google.test";
 
@@ -73,7 +76,7 @@
     public void setUp() {
         MockitoAnnotations.initMocks(this);
 
-        mContext = RuntimeEnvironment.application;
+        mContext = getApplicationContext();
 
         mTestProfile = new InvariantDeviceProfile();
         mTestProfile.numRows = 5;
@@ -152,16 +155,14 @@
         ArrayList<WidgetItem> widgetItems = new ArrayList<>();
         widgetSizes.stream().forEach(
                 widgetSize -> {
-                    ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
-                    AppWidgetProviderInfo info = new AppWidgetProviderInfo();
-                    info.provider = ComponentName.createRelative(TEST_PACKAGE,
-                            ".WidgetProvider_" + widgetSize.x + "x" + widgetSize.y);
+                    AppWidgetProviderInfo info = createAppWidgetProviderInfo(
+                            ComponentName.createRelative(
+                                    TEST_PACKAGE,
+                                    ".WidgetProvider_" + widgetSize.x + "x" + widgetSize.y));
                     LauncherAppWidgetProviderInfo widgetInfo =
                             LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, info);
                     widgetInfo.spanX = widgetSize.x;
                     widgetInfo.spanY = widgetSize.y;
-                    ReflectionHelpers.setField(widgetInfo, "providerInfo",
-                            packageManager.addReceiverIfNotPresent(widgetInfo.provider));
                     widgetItems.add(new WidgetItem(widgetInfo, mTestProfile, mIconCache));
                 }
         );