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/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/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/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/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/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));
                 }
         );