diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
index eb4c136..85d10b4 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
@@ -819,8 +819,11 @@
                 .setDuration(isStashed ? duration / 2 : duration));
     }
 
-    private static void play(AnimatorSet as, Animator a, long startDelay, long duration,
+    private static void play(AnimatorSet as, @Nullable Animator a, long startDelay, long duration,
             Interpolator interpolator) {
+        if (a == null) {
+            return;
+        }
         a.setDuration(duration);
         a.setStartDelay(startDelay);
         a.setInterpolator(interpolator);
diff --git a/quickstep/tests/src/com/android/quickstep/ViewInflationDuringSwipeUp.java b/quickstep/tests/src/com/android/quickstep/ViewInflationDuringSwipeUp.java
index 7e408a8..acb2db5 100644
--- a/quickstep/tests/src/com/android/quickstep/ViewInflationDuringSwipeUp.java
+++ b/quickstep/tests/src/com/android/quickstep/ViewInflationDuringSwipeUp.java
@@ -187,7 +187,7 @@
                     LauncherSettings.Settings.METHOD_CREATE_EMPTY_DB);
             LauncherSettings.Settings.call(mResolver,
                     LauncherSettings.Settings.METHOD_CLEAR_EMPTY_DB_FLAG);
-            LauncherAppWidgetProviderInfo info = TestViewHelpers.findWidgetProvider(this, false);
+            LauncherAppWidgetProviderInfo info = TestViewHelpers.findWidgetProvider(false);
             // Make sure the widget is big enough to show a list of items
             info.minSpanX = 2;
             info.minSpanY = 2;
diff --git a/src/com/android/launcher3/model/DeviceGridState.java b/src/com/android/launcher3/model/DeviceGridState.java
index edc8c1b..f24d1d2 100644
--- a/src/com/android/launcher3/model/DeviceGridState.java
+++ b/src/com/android/launcher3/model/DeviceGridState.java
@@ -33,6 +33,7 @@
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.logging.StatsLogManager.LauncherEvent;
+import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext;
 
 import java.util.Locale;
 import java.util.Objects;
@@ -92,6 +93,9 @@
      * Stores the device state to shared preferences
      */
     public void writeToPrefs(Context context) {
+        if (context instanceof SandboxContext) {
+            return;
+        }
         LauncherPrefs.get(context).put(
                 WORKSPACE_SIZE.to(mGridSizeString),
                 HOTSEAT_COUNT.to(mNumHotseat),
diff --git a/src/com/android/launcher3/model/GridSizeMigrationUtil.java b/src/com/android/launcher3/model/GridSizeMigrationUtil.java
index 9d16610..c233872 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationUtil.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationUtil.java
@@ -45,7 +45,6 @@
 import com.android.launcher3.provider.LauncherDbUtils.SQLiteTransaction;
 import com.android.launcher3.util.GridOccupancy;
 import com.android.launcher3.util.IntArray;
-import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.WidgetManagerHelper;
 
@@ -133,10 +132,8 @@
             Log.v(TAG, "Workspace migration completed in "
                     + (System.currentTimeMillis() - migrationStartTime));
 
-            if (!(context instanceof SandboxContext)) {
-                // Save current configuration, so that the migration does not run again.
-                destDeviceState.writeToPrefs(context);
-            }
+            // Save current configuration, so that the migration does not run again.
+            destDeviceState.writeToPrefs(context);
         }
     }
 
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 91fd65d..0767426 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -367,18 +367,9 @@
         final boolean isSdCardReady = Utilities.isBootCompleted();
         final WidgetManagerHelper widgetHelper = new WidgetManagerHelper(context);
 
-        boolean clearDb = false;
-        if (!mApp.getModel().getModelDbController().migrateGridIfNeeded()) {
-            // Migration failed. Clear workspace.
-            clearDb = true;
-        }
-
-        if (clearDb) {
-            Log.d(TAG, "loadWorkspace: resetting launcher database");
-            Settings.call(contentResolver, Settings.METHOD_CREATE_EMPTY_DB);
-        }
-
+        mApp.getModel().getModelDbController().tryMigrateDB();
         Log.d(TAG, "loadWorkspace: loading default favorites");
+        mApp.getModel().getModelDbController().loadDefaultFavoritesIfNecessary();
         Settings.call(contentResolver, Settings.METHOD_LOAD_DEFAULT_FAVORITES);
 
         synchronized (mBgDataModel) {
diff --git a/src/com/android/launcher3/model/ModelDbController.java b/src/com/android/launcher3/model/ModelDbController.java
index f0e5ef6..d5cc82b 100644
--- a/src/com/android/launcher3/model/ModelDbController.java
+++ b/src/com/android/launcher3/model/ModelDbController.java
@@ -280,13 +280,29 @@
                 mOpenHelper.getReadableDatabase(), Favorites.HYBRID_HOTSEAT_BACKUP_TABLE);
     }
 
+
+    /**
+     * Migrates the DB if needed. If the migration failed, it clears the DB.
+     */
+    public void tryMigrateDB() {
+        if (!migrateGridIfNeeded()) {
+            Log.d(TAG, "Migration failed: resetting launcher database");
+            createEmptyDB();
+            LauncherPrefs.get(mContext).putSync(
+                    getEmptyDbCreatedKey(mOpenHelper.getDatabaseName()).to(true));
+
+            // Write the grid state to avoid another migration
+            new DeviceGridState(LauncherAppState.getIDP(mContext)).writeToPrefs(mContext);
+        }
+    }
+
     /**
      * Migrates the DB if needed, and returns false if the migration failed
      * and DB needs to be cleared.
      * @return true if migration was success or ignored, false if migration failed
      * and the DB should be reset.
      */
-    public boolean migrateGridIfNeeded() {
+    private boolean migrateGridIfNeeded() {
         createDbIfNotExists();
         InvariantDeviceProfile idp = LauncherAppState.getIDP(mContext);
         if (!GridSizeMigrationUtil.needsToMigrate(mContext, idp)) {
diff --git a/tests/src/com/android/launcher3/celllayout/FavoriteItemsTransaction.java b/tests/src/com/android/launcher3/celllayout/FavoriteItemsTransaction.java
index 8ce932d..80cc5aa 100644
--- a/tests/src/com/android/launcher3/celllayout/FavoriteItemsTransaction.java
+++ b/tests/src/com/android/launcher3/celllayout/FavoriteItemsTransaction.java
@@ -15,62 +15,80 @@
  */
 package com.android.launcher3.celllayout;
 
+import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
+import static com.android.launcher3.util.TestUtil.runOnExecutorSync;
 
-import android.content.ContentResolver;
 import android.content.ContentValues;
 import android.content.Context;
 
 import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherModel;
 import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.model.BgDataModel.Callbacks;
+import com.android.launcher3.model.ModelDbController;
 import com.android.launcher3.model.data.ItemInfo;
-import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.util.ContentWriter;
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.concurrent.ExecutionException;
+import java.util.function.Supplier;
 
 public class FavoriteItemsTransaction {
-    private ArrayList<ItemInfo> mItemsToSubmit;
+    private ArrayList<Supplier<ItemInfo>> mItemsToSubmit;
     private Context mContext;
-    private ContentResolver mResolver;
-    public AbstractLauncherUiTest mTest;
 
-    public FavoriteItemsTransaction(Context context, AbstractLauncherUiTest test) {
+    public FavoriteItemsTransaction(Context context) {
         mItemsToSubmit = new ArrayList<>();
         mContext = context;
-        mResolver = mContext.getContentResolver();
-        mTest = test;
     }
 
-    public FavoriteItemsTransaction addItem(ItemInfo itemInfo) {
+    public FavoriteItemsTransaction addItem(Supplier<ItemInfo> itemInfo) {
         this.mItemsToSubmit.add(itemInfo);
         return this;
     }
 
-    public FavoriteItemsTransaction removeLast() {
-        this.mItemsToSubmit.remove(this.mItemsToSubmit.size() - 1);
-        return this;
-    }
-
     /**
      * Commits all the ItemInfo into the database of Favorites
      **/
-    public void commit() throws ExecutionException, InterruptedException {
-        List<ContentValues> values = new ArrayList<>();
-        for (ItemInfo item : this.mItemsToSubmit) {
-            ContentWriter writer = new ContentWriter(mContext);
-            item.onAddToDatabase(writer);
-            writer.put(LauncherSettings.Favorites._ID, item.id);
-            values.add(writer.getValues(mContext));
-        }
-        // Submit the icons to the database in the model thread to prevent race conditions
-        MODEL_EXECUTOR.submit(() -> mResolver.bulkInsert(LauncherSettings.Favorites.CONTENT_URI,
-                values.toArray(new ContentValues[0]))).get();
-        // Reload the state of the Launcher
-        MAIN_EXECUTOR.submit(() -> LauncherAppState.getInstance(
-                mContext).getModel().forceReload()).get();
+    public void commit() {
+        LauncherModel model = LauncherAppState.getInstance(mContext).getModel();
+        // Load the model once so that there is no pending migration:
+        loadModelSync(model);
+
+        runOnExecutorSync(MODEL_EXECUTOR, () -> {
+            ModelDbController controller = model.getModelDbController();
+            // Migrate any previous data so that the DB state is correct
+            controller.tryMigrateDB();
+
+            // Create DB again to load fresh data
+            controller.createEmptyDB();
+            controller.clearEmptyDbFlag();
+
+            // Add new data
+            List<ContentValues> values = new ArrayList<>();
+            int count = mItemsToSubmit.size();
+            for (int i = 0; i < count; i++) {
+                ContentWriter writer = new ContentWriter(mContext);
+                mItemsToSubmit.get(i).get().onAddToDatabase(writer);
+                writer.put(LauncherSettings.Favorites._ID, i);
+                values.add(writer.getValues(mContext));
+            }
+            controller.bulkInsert(TABLE_NAME, values.toArray(new ContentValues[0]));
+        });
+
+        // Reload model
+        runOnExecutorSync(MAIN_EXECUTOR, model::forceReload);
+        loadModelSync(model);
+    }
+
+    private void loadModelSync(LauncherModel model) {
+        Callbacks mockCb = new Callbacks() { };
+        runOnExecutorSync(MAIN_EXECUTOR, () -> model.addCallbacksAndLoad(mockCb));
+        runOnExecutorSync(MODEL_EXECUTOR, () -> { });
+
+        runOnExecutorSync(MAIN_EXECUTOR, () -> { });
+        runOnExecutorSync(MAIN_EXECUTOR, () -> model.removeCallbacks(mockCb));
     }
 }
diff --git a/tests/src/com/android/launcher3/celllayout/ReorderWidgets.java b/tests/src/com/android/launcher3/celllayout/ReorderWidgets.java
index c4aa42b..c426480 100644
--- a/tests/src/com/android/launcher3/celllayout/ReorderWidgets.java
+++ b/tests/src/com/android/launcher3/celllayout/ReorderWidgets.java
@@ -60,7 +60,7 @@
 
     @Before
     public void setup() throws Throwable {
-        mWorkspaceBuilder = new TestWorkspaceBuilder(this, mTargetContext);
+        mWorkspaceBuilder = new TestWorkspaceBuilder(mTargetContext);
         TaplTestsLauncher3.initialize(this);
         clearHomescreen();
     }
@@ -109,7 +109,7 @@
                 testCase.mStart);
 
         FavoriteItemsTransaction transaction =
-                new FavoriteItemsTransaction(mTargetContext, this);
+                new FavoriteItemsTransaction(mTargetContext);
         transaction = buildWorkspaceFromBoards(testCase.mStart, transaction);
         transaction.commit();
         // resetLoaderState triggers the launcher to start loading the workspace which allows
diff --git a/tests/src/com/android/launcher3/celllayout/TestWorkspaceBuilder.java b/tests/src/com/android/launcher3/celllayout/TestWorkspaceBuilder.java
index 5945605..cde733e 100644
--- a/tests/src/com/android/launcher3/celllayout/TestWorkspaceBuilder.java
+++ b/tests/src/com/android/launcher3/celllayout/TestWorkspaceBuilder.java
@@ -15,6 +15,9 @@
  */
 package com.android.launcher3.celllayout;
 
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
+import static com.android.launcher3.ui.TestViewHelpers.findWidgetProvider;
 import static com.android.launcher3.util.WidgetUtils.createWidgetInfo;
 
 import android.content.ComponentName;
@@ -25,33 +28,29 @@
 import android.os.UserHandle;
 import android.util.Log;
 
-import androidx.test.core.app.ApplicationProvider;
-
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
-import com.android.launcher3.ui.AbstractLauncherUiTest;
-import com.android.launcher3.ui.TestViewHelpers;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 
+import java.util.function.Supplier;
+import java.util.stream.IntStream;
+
 public class TestWorkspaceBuilder {
 
     private static final String TAG = "CellLayoutBoardBuilder";
     private static final ComponentName APP_COMPONENT_NAME = new ComponentName(
             "com.google.android.calculator", "com.android.calculator2.Calculator");
 
-    public AbstractLauncherUiTest mTest;
-
     private UserHandle mMyUser;
 
     private Context mContext;
     private ContentResolver mResolver;
 
-    public TestWorkspaceBuilder(AbstractLauncherUiTest test, Context context) {
-        mTest = test;
+    public TestWorkspaceBuilder(Context context) {
         mMyUser = Process.myUserHandle();
         mContext = context;
         mResolver = mContext.getContentResolver();
@@ -68,10 +67,9 @@
             for (int y = initY; y < initY + widgetRect.getSpanY(); y++) {
                 try {
                     // this widgets are filling, we don't care if we can't place them
-                    ItemInfo item = createWidgetInCell(
+                    transaction.addItem(createWidgetInCell(
                             new CellLayoutBoard.WidgetRect(CellLayoutBoard.CellType.IGNORE,
-                                    new Rect(x, y, x, y)), screenId);
-                    transaction.addItem(item);
+                                    new Rect(x, y, x, y)), screenId));
                 } catch (Exception e) {
                     Log.d(TAG, "Unable to place filling widget at " + x + "," + y);
                 }
@@ -108,7 +106,7 @@
         board.getWidgets().forEach(
                 (widgetRect) -> addCorrespondingWidgetRect(widgetRect, transaction, screenId));
         board.getIcons().forEach((iconPoint) ->
-                transaction.addItem(createIconInCell(iconPoint, screenId))
+                transaction.addItem(() -> createIconInCell(iconPoint, screenId))
         );
         return transaction;
     }
@@ -118,24 +116,25 @@
      * be clean otherwise this doesn't overrides the existing icons.
      */
     public FavoriteItemsTransaction fillHotseatIcons(FavoriteItemsTransaction transaction) {
-        int hotseatCount = InvariantDeviceProfile.INSTANCE.get(mContext).numDatabaseHotseatIcons;
-        for (int i = 0; i < hotseatCount; i++) {
-            transaction.addItem(getHotseatValues(i));
-        }
+        IntStream.range(0, InvariantDeviceProfile.INSTANCE.get(mContext).numDatabaseHotseatIcons)
+                .forEach(i -> transaction.addItem(() -> getHotseatValues(i)));
         return transaction;
     }
 
-    private ItemInfo createWidgetInCell(CellLayoutBoard.WidgetRect widgetRect, int screenId) {
-        LauncherAppWidgetProviderInfo info = TestViewHelpers.findWidgetProvider(mTest, false);
-        LauncherAppWidgetInfo item = createWidgetInfo(info,
-                ApplicationProvider.getApplicationContext(), true);
-        item.id = getID();
-        item.cellX = widgetRect.getCellX();
-        item.cellY = widgetRect.getCellY();
-        item.spanX = widgetRect.getSpanX();
-        item.spanY = widgetRect.getSpanY();
-        item.screenId = screenId;
-        return item;
+    private Supplier<ItemInfo> createWidgetInCell(
+            CellLayoutBoard.WidgetRect widgetRect, int screenId) {
+        // Create the widget lazily since the appWidgetId can get lost during setup
+        return () -> {
+            LauncherAppWidgetProviderInfo info = findWidgetProvider(false);
+            LauncherAppWidgetInfo item = createWidgetInfo(info, getApplicationContext(), true);
+            item.id = getID();
+            item.cellX = widgetRect.getCellX();
+            item.cellY = widgetRect.getCellY();
+            item.spanX = widgetRect.getSpanX();
+            item.spanY = widgetRect.getSpanY();
+            item.screenId = screenId;
+            return item;
+        };
     }
 
     private ItemInfo createIconInCell(CellLayoutBoard.IconPoint iconPoint, int screenId) {
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 09c4b1f..4506fa3 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -19,6 +19,7 @@
 
 import static com.android.launcher3.ui.TaplTestsLauncher3.getAppPackageName;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
+import static com.android.launcher3.util.TestUtil.getOnUiThread;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
@@ -39,7 +40,6 @@
 import android.os.RemoteException;
 import android.os.UserHandle;
 import android.os.UserManager;
-import android.system.OsConstants;
 import android.util.Log;
 
 import androidx.test.InstrumentationRegistry;
@@ -64,6 +64,7 @@
 import com.android.launcher3.testing.shared.TestProtocol;
 import com.android.launcher3.util.LooperExecutor;
 import com.android.launcher3.util.SimpleBroadcastReceiver;
+import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.Wait;
 import com.android.launcher3.util.WidgetUtils;
 import com.android.launcher3.util.rule.FailureWatcher;
@@ -82,10 +83,8 @@
 import org.junit.rules.TestRule;
 
 import java.io.IOException;
-import java.util.concurrent.Callable;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -98,7 +97,7 @@
     public static final long DEFAULT_ACTIVITY_TIMEOUT = TimeUnit.SECONDS.toMillis(10);
     public static final long DEFAULT_BROADCAST_TIMEOUT_SECS = 5;
 
-    public static final long DEFAULT_UI_TIMEOUT = 10000;
+    public static final long DEFAULT_UI_TIMEOUT = TestUtil.DEFAULT_UI_TIMEOUT;
     private static final String TAG = "AbstractLauncherUiTest";
 
     private static boolean sDumpWasGenerated = false;
@@ -332,22 +331,6 @@
         mLauncher.waitForLauncherInitialized();
     }
 
-    /**
-     * Runs the callback on the UI thread and returns the result.
-     */
-    protected <T> T getOnUiThread(final Callable<T> callback) {
-        try {
-            return mMainThreadExecutor.submit(callback).get(DEFAULT_UI_TIMEOUT,
-                    TimeUnit.MILLISECONDS);
-        } catch (TimeoutException e) {
-            Log.e(TAG, "Timeout in getOnUiThread, sending SIGABRT", e);
-            Process.sendSignal(Process.myPid(), OsConstants.SIGABRT);
-            throw new RuntimeException(e);
-        } catch (Throwable e) {
-            throw new RuntimeException(e);
-        }
-    }
-
     protected <T> T getFromLauncher(Function<Launcher, T> f) {
         if (!TestHelpers.isInLauncherProcess()) return null;
         return getOnUiThread(() -> f.apply(mActivityMonitor.getActivity()));
diff --git a/tests/src/com/android/launcher3/ui/TestViewHelpers.java b/tests/src/com/android/launcher3/ui/TestViewHelpers.java
index 083f580..4b2bade 100644
--- a/tests/src/com/android/launcher3/ui/TestViewHelpers.java
+++ b/tests/src/com/android/launcher3/ui/TestViewHelpers.java
@@ -15,11 +15,14 @@
  */
 package com.android.launcher3.ui;
 
-import static androidx.test.InstrumentationRegistry.getInstrumentation;
-import static androidx.test.InstrumentationRegistry.getTargetContext;
+import static android.os.Process.myUserHandle;
 
+import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
+
+import static com.android.launcher3.util.TestUtil.getOnUiThread;
+
+import android.app.Instrumentation;
 import android.content.ComponentName;
-import android.os.Process;
 import android.util.Log;
 import android.view.View;
 import android.view.ViewGroup;
@@ -29,7 +32,6 @@
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.WidgetManagerHelper;
 
-import java.util.concurrent.Callable;
 import java.util.function.Function;
 
 public class TestViewHelpers {
@@ -38,23 +40,16 @@
     /**
      * Finds a widget provider which can fit on the home screen.
      *
-     * @param test               test suite.
      * @param hasConfigureScreen if true, a provider with a config screen is returned.
      */
-    public static LauncherAppWidgetProviderInfo findWidgetProvider(AbstractLauncherUiTest test,
-            final boolean hasConfigureScreen) {
-        LauncherAppWidgetProviderInfo info =
-                test.getOnUiThread(new Callable<LauncherAppWidgetProviderInfo>() {
-                    @Override
-                    public LauncherAppWidgetProviderInfo call() throws Exception {
-                        ComponentName cn = new ComponentName(getInstrumentation().getContext(),
-                                hasConfigureScreen ? AppWidgetWithConfig.class
-                                        : AppWidgetNoConfig.class);
-                        Log.d(TAG, "findWidgetProvider componentName=" + cn.flattenToString());
-                        return new WidgetManagerHelper(getTargetContext())
-                                .findProvider(cn, Process.myUserHandle());
-                    }
-                });
+    public static LauncherAppWidgetProviderInfo findWidgetProvider(boolean hasConfigureScreen) {
+        LauncherAppWidgetProviderInfo info = getOnUiThread(() -> {
+            Instrumentation i = getInstrumentation();
+            ComponentName cn = new ComponentName(i.getContext(),
+                    hasConfigureScreen ? AppWidgetWithConfig.class : AppWidgetNoConfig.class);
+            Log.d(TAG, "findWidgetProvider componentName=" + cn.flattenToString());
+            return new WidgetManagerHelper(i.getTargetContext()).findProvider(cn, myUserHandle());
+        });
         if (info == null) {
             throw new IllegalArgumentException("No valid widget provider");
         }
diff --git a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
index 3806648..fddbc37 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
@@ -64,7 +64,7 @@
     @Before
     public void setUp() throws Exception {
         super.setUp();
-        mWidgetInfo = TestViewHelpers.findWidgetProvider(this, true /* hasConfigureScreen */);
+        mWidgetInfo = TestViewHelpers.findWidgetProvider(true /* hasConfigureScreen */);
         mAppWidgetManager = AppWidgetManager.getInstance(mTargetContext);
     }
 
diff --git a/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
index 028b9b6..5cdc886 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
@@ -61,7 +61,7 @@
         waitForLauncherCondition("Workspace didn't finish loading", l -> !l.isWorkspaceLoading());
 
         final LauncherAppWidgetProviderInfo widgetInfo =
-                TestViewHelpers.findWidgetProvider(this, false /* hasConfigureScreen */);
+                TestViewHelpers.findWidgetProvider(false /* hasConfigureScreen */);
 
         WidgetResizeFrame resizeFrame = mLauncher
                 .getWorkspace()
diff --git a/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java
index 0f861eb..99d2201 100644
--- a/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java
@@ -17,6 +17,7 @@
 
 import static androidx.test.InstrumentationRegistry.getTargetContext;
 
+import static com.android.launcher3.util.TestUtil.getOnUiThread;
 import static com.android.launcher3.util.WidgetUtils.createWidgetInfo;
 
 import static org.junit.Assert.assertEquals;
@@ -103,7 +104,7 @@
 
     @Test
     public void testBindNormalWidget_withConfig() {
-        LauncherAppWidgetProviderInfo info = TestViewHelpers.findWidgetProvider(this, true);
+        LauncherAppWidgetProviderInfo info = TestViewHelpers.findWidgetProvider(true);
         LauncherAppWidgetInfo item = createWidgetInfo(info, getTargetContext(), true);
 
         addItemToScreen(item);
@@ -112,7 +113,7 @@
 
     @Test
     public void testBindNormalWidget_withoutConfig() {
-        LauncherAppWidgetProviderInfo info = TestViewHelpers.findWidgetProvider(this, false);
+        LauncherAppWidgetProviderInfo info = TestViewHelpers.findWidgetProvider(false);
         LauncherAppWidgetInfo item = createWidgetInfo(info, getTargetContext(), true);
 
         addItemToScreen(item);
@@ -121,7 +122,7 @@
 
     @Test
     public void testUnboundWidget_removed() {
-        LauncherAppWidgetProviderInfo info = TestViewHelpers.findWidgetProvider(this, false);
+        LauncherAppWidgetProviderInfo info = TestViewHelpers.findWidgetProvider(false);
         LauncherAppWidgetInfo item = createWidgetInfo(info, getTargetContext(), false);
         item.appWidgetId = -33;
 
@@ -140,7 +141,7 @@
     @Test
     public void testPendingWidget_autoRestored() {
         // A non-restored widget with no config screen gets restored automatically.
-        LauncherAppWidgetProviderInfo info = TestViewHelpers.findWidgetProvider(this, false);
+        LauncherAppWidgetProviderInfo info = TestViewHelpers.findWidgetProvider(false);
 
         // Do not bind the widget
         LauncherAppWidgetInfo item = createWidgetInfo(info, getTargetContext(), false);
@@ -153,7 +154,7 @@
     @Test
     public void testPendingWidget_withConfigScreen() {
         // A non-restored widget with config screen get bound and shows a 'Click to setup' UI.
-        LauncherAppWidgetProviderInfo info = TestViewHelpers.findWidgetProvider(this, true);
+        LauncherAppWidgetProviderInfo info = TestViewHelpers.findWidgetProvider(true);
 
         // Do not bind the widget
         LauncherAppWidgetInfo item = createWidgetInfo(info, getTargetContext(), false);
diff --git a/tests/src/com/android/launcher3/util/TestUtil.java b/tests/src/com/android/launcher3/util/TestUtil.java
index 14bc8b9..7332b83 100644
--- a/tests/src/com/android/launcher3/util/TestUtil.java
+++ b/tests/src/com/android/launcher3/util/TestUtil.java
@@ -18,14 +18,13 @@
 import static android.util.Base64.NO_PADDING;
 import static android.util.Base64.NO_WRAP;
 
-import static androidx.test.InstrumentationRegistry.getContext;
-import static androidx.test.InstrumentationRegistry.getInstrumentation;
-import static androidx.test.InstrumentationRegistry.getTargetContext;
+import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 
 import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_KEY;
 import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_LABEL;
 import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_TAG;
 
+import android.app.Instrumentation;
 import android.app.blob.BlobHandle;
 import android.app.blob.BlobStoreManager;
 import android.content.Context;
@@ -35,9 +34,12 @@
 import android.os.Handler;
 import android.os.Looper;
 import android.os.ParcelFileDescriptor.AutoCloseOutputStream;
+import android.os.Process;
 import android.os.UserHandle;
 import android.provider.Settings;
+import android.system.OsConstants;
 import android.util.Base64;
+import android.util.Log;
 
 import androidx.test.uiautomator.UiDevice;
 
@@ -52,27 +54,35 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.security.MessageDigest;
+import java.util.concurrent.Callable;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
+import java.util.concurrent.FutureTask;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.function.Predicate;
 import java.util.function.ToIntFunction;
 
 public class TestUtil {
+    private static final String TAG = "TestUtil";
+
     public static final String DUMMY_PACKAGE = "com.example.android.aardwolf";
     public static final int DEFAULT_USER_ID = 0;
+    public static final long DEFAULT_UI_TIMEOUT = 10000;
 
     public static void installDummyApp() throws IOException {
         installDummyAppForUser(DEFAULT_USER_ID);
     }
 
     public static void installDummyAppForUser(int userId) throws IOException {
+        Instrumentation instrumentation = getInstrumentation();
         // Copy apk from resources to a local file and install from there.
-        final Resources resources = getContext().getResources();
+        final Resources resources = instrumentation.getContext().getResources();
         final InputStream in = resources.openRawResource(
                 resources.getIdentifier("aardwolf_dummy_app",
-                        "raw", getContext().getPackageName()));
-        final String apkFilename = getInstrumentation().getTargetContext().
-                getFilesDir().getPath() + "/dummy_app.apk";
+                        "raw", instrumentation.getContext().getPackageName()));
+        final String apkFilename = instrumentation.getTargetContext()
+                        .getFilesDir().getPath() + "/dummy_app.apk";
 
         try (PackageInstallCheck pic = new PackageInstallCheck()) {
             final FileOutputStream out = new FileOutputStream(apkFilename);
@@ -85,7 +95,7 @@
             in.close();
             out.close();
 
-            final String result = UiDevice.getInstance(getInstrumentation())
+            final String result = UiDevice.getInstance(instrumentation)
                     .executeShellCommand("pm install --user " + userId + " " + apkFilename);
             Assert.assertTrue(
                     "Failed to install wellbeing test apk; make sure the device is rooted",
@@ -177,13 +187,33 @@
         }
     }
 
+    /**
+     * Runs the callback on the UI thread and returns the result.
+     */
+    public static <T> T getOnUiThread(final Callable<T> callback) {
+        try {
+            FutureTask<T> task = new FutureTask<>(callback);
+            if (Looper.myLooper() == Looper.getMainLooper()) {
+                task.run();
+            } else {
+                new Handler(Looper.getMainLooper()).post(task);
+            }
+            return task.get(DEFAULT_UI_TIMEOUT, TimeUnit.MILLISECONDS);
+        } catch (TimeoutException e) {
+            Log.e(TAG, "Timeout in getOnUiThread, sending SIGABRT", e);
+            Process.sendSignal(Process.myPid(), OsConstants.SIGABRT);
+            throw new RuntimeException(e);
+        } catch (Throwable e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     /** Interface to indicate a runnable which can throw any exception. */
     public interface UncheckedRunnable {
         /** Method to run the task */
         void run() throws Exception;
     }
 
-
     private static class PackageInstallCheck extends LauncherApps.Callback
             implements AutoCloseable {
 
@@ -191,7 +221,8 @@
         final LauncherApps mLauncherApps;
 
         PackageInstallCheck() {
-            mLauncherApps = getTargetContext().getSystemService(LauncherApps.class);
+            mLauncherApps = getInstrumentation().getTargetContext()
+                    .getSystemService(LauncherApps.class);
             mLauncherApps.registerCallback(this, new Handler(Looper.getMainLooper()));
         }
 
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 1111d32..2286d7e 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -138,17 +138,6 @@
         OUTSIDE_WITH_KEYCODE,
     }
 
-    /**
-     * Represents a point in the code at which a callback can run.
-     */
-    public enum CALLBACK_RUN_POINT {
-        CALLBACK_HOLD_BEFORE_DROP,
-        CALLBACK_HOVER_ENTER,
-        CALLBACK_HOVER_EXIT,
-    }
-
-    private Consumer<CALLBACK_RUN_POINT> mCallbackAtRunPoint = null;
-
     // Base class for launcher containers.
     abstract static class VisibleContainer {
         protected final LauncherInstrumentation mLauncher;
@@ -2064,22 +2053,6 @@
     }
 
     /**
-     * Sets the consumer to run callbacks at all run-points.
-     */
-    public void setRunPointCallback(Consumer<CALLBACK_RUN_POINT> callback) {
-        mCallbackAtRunPoint = callback;
-    }
-
-    /**
-     * Runs the callback at the specified point if it exists.
-     */
-    void runCallbackIfActive(CALLBACK_RUN_POINT runPoint) {
-        if (mCallbackAtRunPoint != null) {
-            mCallbackAtRunPoint.accept(runPoint);
-        }
-    }
-
-    /**
      * Waits until a particular condition is true. Based on WaitMixin.
      */
     boolean waitAndGet(BooleanSupplier condition, long timeout, long interval) {
diff --git a/tests/tapl/com/android/launcher3/tapl/OverviewTaskMenuItem.java b/tests/tapl/com/android/launcher3/tapl/OverviewTaskMenuItem.java
index b2cc92d..e3035bf 100644
--- a/tests/tapl/com/android/launcher3/tapl/OverviewTaskMenuItem.java
+++ b/tests/tapl/com/android/launcher3/tapl/OverviewTaskMenuItem.java
@@ -15,13 +15,7 @@
  */
 package com.android.launcher3.tapl;
 
-import static com.android.launcher3.tapl.LauncherInstrumentation.CALLBACK_RUN_POINT.CALLBACK_HOVER_ENTER;
-import static com.android.launcher3.tapl.LauncherInstrumentation.CALLBACK_RUN_POINT.CALLBACK_HOVER_EXIT;
-
-import android.graphics.Point;
 import android.graphics.Rect;
-import android.os.SystemClock;
-import android.view.MotionEvent;
 
 import androidx.test.uiautomator.UiObject2;
 
@@ -42,28 +36,4 @@
     public Rect getVisibleBounds() {
         return mMenuItem.getVisibleBounds();
     }
-
-    /**
-     * Emulate the cursor entering and exiting a hover over this menu item.
-     */
-    public void hoverCursor() {
-        try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck();
-             LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
-                     "cursor hover entering menu item")) {
-            long downTime = SystemClock.uptimeMillis();
-            mLauncher.sendPointer(downTime, downTime, MotionEvent.ACTION_HOVER_ENTER,
-                    new Point(mMenuItem.getVisibleCenter().x, mMenuItem.getVisibleCenter().y),
-                    null);
-            mLauncher.runCallbackIfActive(CALLBACK_HOVER_ENTER);
-
-            try (LauncherInstrumentation.Closable c1 = mLauncher.addContextLayer(
-                    "cursor hover exiting menu item")) {
-                downTime = SystemClock.uptimeMillis();
-                mLauncher.sendPointer(downTime, downTime, MotionEvent.ACTION_HOVER_EXIT,
-                        new Point(mMenuItem.getVisibleCenter().x, mMenuItem.getVisibleCenter().y),
-                        null);
-                mLauncher.runCallbackIfActive(CALLBACK_HOVER_EXIT);
-            }
-        }
-    }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/Workspace.java b/tests/tapl/com/android/launcher3/tapl/Workspace.java
index 7bb02cb..8604988 100644
--- a/tests/tapl/com/android/launcher3/tapl/Workspace.java
+++ b/tests/tapl/com/android/launcher3/tapl/Workspace.java
@@ -18,7 +18,6 @@
 
 import static android.view.accessibility.AccessibilityEvent.TYPE_VIEW_SCROLLED;
 
-import static com.android.launcher3.tapl.LauncherInstrumentation.CALLBACK_RUN_POINT.CALLBACK_HOLD_BEFORE_DROP;
 import static com.android.launcher3.testing.shared.TestProtocol.ALL_APPS_STATE_ORDINAL;
 import static com.android.launcher3.testing.shared.TestProtocol.NORMAL_STATE_ORDINAL;
 
@@ -553,7 +552,6 @@
             launcher.movePointer(dragStart, targetDest, DEFAULT_DRAG_STEPS, isDecelerating,
                     downTime, SystemClock.uptimeMillis(), false,
                     LauncherInstrumentation.GestureScope.INSIDE);
-            launcher.runCallbackIfActive(CALLBACK_HOLD_BEFORE_DROP);
             dropDraggedIcon(launcher, targetDest, downTime, expectDropEvents);
         }
     }
@@ -587,7 +585,6 @@
             launcher.movePointer(dragStart, targetDest, DEFAULT_DRAG_STEPS, isDecelerating,
                     downTime, SystemClock.uptimeMillis(), false,
                     LauncherInstrumentation.GestureScope.INSIDE);
-            launcher.runCallbackIfActive(CALLBACK_HOLD_BEFORE_DROP);
             dropDraggedIcon(launcher, targetDest, downTime, expectDropEvents);
         }
     }