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/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/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/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