diff --git a/iconloaderlib/src/com/android/launcher3/icons/cache/BaseIconCache.java b/iconloaderlib/src/com/android/launcher3/icons/cache/BaseIconCache.java
index 6f63d88..e807791 100644
--- a/iconloaderlib/src/com/android/launcher3/icons/cache/BaseIconCache.java
+++ b/iconloaderlib/src/com/android/launcher3/icons/cache/BaseIconCache.java
@@ -45,6 +45,7 @@
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.icons.BaseIconFactory;
 import com.android.launcher3.icons.BitmapInfo;
@@ -250,9 +251,9 @@
      * @param replaceExisting if true, it will recreate the bitmap even if it already exists in
      *                        the memory. This is useful then the previous bitmap was created using
      *                        old data.
-     * package private
      */
-    protected synchronized <T> void addIconToDBAndMemCache(T object, CachingLogic<T> cachingLogic,
+    @VisibleForTesting
+    public synchronized <T> void addIconToDBAndMemCache(T object, CachingLogic<T> cachingLogic,
             PackageInfo info, long userSerial, boolean replaceExisting) {
         UserHandle user = cachingLogic.getUser(object);
         ComponentName componentName = cachingLogic.getComponent(object);
diff --git a/robolectric_tests/Android.mk b/robolectric_tests/Android.mk
index 310d43c..86a6e8c 100644
--- a/robolectric_tests/Android.mk
+++ b/robolectric_tests/Android.mk
@@ -19,6 +19,8 @@
 include $(CLEAR_VARS)
 
 LOCAL_MODULE := LauncherRoboTests
+LOCAL_MODULE_CLASS := JAVA_LIBRARIES
+
 LOCAL_SDK_VERSION := current
 LOCAL_SRC_FILES := $(call all-java-files-under, src)
 LOCAL_STATIC_JAVA_LIBRARIES := \
@@ -34,6 +36,9 @@
 LOCAL_INSTRUMENTATION_FOR := Launcher3
 LOCAL_MODULE_TAGS := optional
 
+# Generate test_config.properties
+include external/robolectric-shadows/gen_test_config.mk
+
 include $(BUILD_STATIC_JAVA_LIBRARY)
 
 ############################################
@@ -43,14 +48,11 @@
 
 LOCAL_MODULE := RunLauncherRoboTests
 LOCAL_SDK_VERSION := current
-LOCAL_JAVA_LIBRARIES := \
-    LauncherRoboTests
+LOCAL_JAVA_LIBRARIES := LauncherRoboTests
 
 LOCAL_RESOURCE_DIR := $(LOCAL_PATH)/res
-
 LOCAL_TEST_PACKAGE := Launcher3
-
-LOCAL_INSTRUMENT_SOURCE_DIRS := $(dir $(LOCAL_PATH))../src \
+LOCAL_INSTRUMENT_SOURCE_DIRS := $(dir $(LOCAL_PATH))../src
 
 LOCAL_ROBOTEST_TIMEOUT := 36000
 
diff --git a/robolectric_tests/config/robolectric.properties b/robolectric_tests/config/robolectric.properties
index e0d6e53..932b01b 100644
--- a/robolectric_tests/config/robolectric.properties
+++ b/robolectric_tests/config/robolectric.properties
@@ -1,2 +1 @@
-manifest=packages/apps/Launcher3/AndroidManifest.xml
-sdk=26
+sdk=28
diff --git a/robolectric_tests/src/com/android/launcher3/config/FlagOverrideRule.java b/robolectric_tests/src/com/android/launcher3/config/FlagOverrideRule.java
index 4bb9a53..d33fecd 100644
--- a/robolectric_tests/src/com/android/launcher3/config/FlagOverrideRule.java
+++ b/robolectric_tests/src/com/android/launcher3/config/FlagOverrideRule.java
@@ -14,6 +14,7 @@
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.function.Function;
@@ -35,6 +36,8 @@
  */
 public final class FlagOverrideRule implements TestRule {
 
+    private final HashMap<String, Boolean> mDefaultOverrides = new HashMap<>();
+
     /**
      * Container annotation for handling multiple {@link FlagOverride} annotations.
      * <p>
@@ -60,6 +63,14 @@
         return new MyStatement(base, description);
     }
 
+    /**
+     * Sets a default override to apply on all tests
+     */
+    public FlagOverrideRule setOverride(BaseTogglableFlag flag, boolean value) {
+        mDefaultOverrides.put(flag.getKey(), value);
+        return this;
+    }
+
     private class MyStatement extends Statement {
 
         private final Statement mBase;
@@ -87,11 +98,15 @@
                         overrides = ((FlagOverrides) annotation).value();
                     }
                 }
-                for (FlagOverride override : overrides) {
-                    BaseTogglableFlag flag = allFlags.get(override.key());
+
+                HashMap<String, Boolean> allOverrides = new HashMap<>(mDefaultOverrides);
+                Arrays.stream(overrides).forEach(o -> allOverrides.put(o.key(), o.value()));
+
+                allOverrides.forEach((key, val) -> {
+                    BaseTogglableFlag flag = allFlags.get(key);
                     changedValues.put(flag, flag.get());
-                    flag.setForTests(override.value());
-                }
+                    flag.setForTests(val);
+                });
                 mBase.evaluate();
             } finally {
                 // Clear the values
diff --git a/robolectric_tests/src/com/android/launcher3/config/FlagOverrideSampleTest.java b/robolectric_tests/src/com/android/launcher3/config/FlagOverrideSampleTest.java
index 31a037b..2a359df 100644
--- a/robolectric_tests/src/com/android/launcher3/config/FlagOverrideSampleTest.java
+++ b/robolectric_tests/src/com/android/launcher3/config/FlagOverrideSampleTest.java
@@ -4,16 +4,16 @@
 import static org.junit.Assert.assertTrue;
 
 import com.android.launcher3.config.FlagOverrideRule.FlagOverride;
+import com.android.launcher3.util.LauncherRoboTestRunner;
 
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.robolectric.RobolectricTestRunner;
 
 /**
  * Sample Robolectric test that demonstrates flag-overriding.
  */
-@RunWith(RobolectricTestRunner.class)
+@RunWith(LauncherRoboTestRunner.class)
 public class FlagOverrideSampleTest {
 
     // Check out https://junit.org/junit4/javadoc/4.12/org/junit/Rule.html for more information
diff --git a/robolectric_tests/src/com/android/launcher3/logging/FileLogTest.java b/robolectric_tests/src/com/android/launcher3/logging/FileLogTest.java
index 410a077..48b5a45 100644
--- a/robolectric_tests/src/com/android/launcher3/logging/FileLogTest.java
+++ b/robolectric_tests/src/com/android/launcher3/logging/FileLogTest.java
@@ -3,11 +3,12 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
+import com.android.launcher3.util.LauncherRoboTestRunner;
+
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.robolectric.RobolectricTestRunner;
 import org.robolectric.RuntimeEnvironment;
 import org.robolectric.Shadows;
 import org.robolectric.util.Scheduler;
@@ -20,7 +21,7 @@
 /**
  * Tests for {@link FileLog}
  */
-@RunWith(RobolectricTestRunner.class)
+@RunWith(LauncherRoboTestRunner.class)
 public class FileLogTest {
 
     private File mTempDir;
diff --git a/robolectric_tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java b/robolectric_tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
index d7a2278..ea7c137 100644
--- a/robolectric_tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
@@ -4,54 +4,70 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 
 import android.content.ComponentName;
+import android.content.Context;
 import android.content.Intent;
 import android.graphics.Rect;
 import android.util.Pair;
 
+import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.ItemInfo;
+import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.WorkspaceItemInfo;
+import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.util.ContentWriter;
 import com.android.launcher3.util.GridOccupancy;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSparseArrayMap;
+import com.android.launcher3.util.LauncherModelHelper;
+import com.android.launcher3.util.LauncherRoboTestRunner;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
-import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.annotation.LooperMode.Mode;
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
 
 /**
  * Tests for {@link AddWorkspaceItemsTask}
  */
-@RunWith(RobolectricTestRunner.class)
-public class AddWorkspaceItemsTaskTest extends BaseModelUpdateTaskTestCase {
+@RunWith(LauncherRoboTestRunner.class)
+@LooperMode(Mode.PAUSED)
+public class AddWorkspaceItemsTaskTest {
 
     private final ComponentName mComponent1 = new ComponentName("a", "b");
     private final ComponentName mComponent2 = new ComponentName("b", "b");
 
-    private IntArray existingScreens;
-    private IntArray newScreens;
-    private IntSparseArrayMap<GridOccupancy> screenOccupancy;
+    private Context mTargetContext;
+    private InvariantDeviceProfile mIdp;
+    private LauncherAppState mAppState;
+    private LauncherModelHelper mModelHelper;
+
+    private IntArray mExistingScreens;
+    private IntArray mNewScreens;
+    private IntSparseArrayMap<GridOccupancy> mScreenOccupancy;
 
     @Before
-    public void initData() throws Exception {
-        existingScreens = new IntArray();
-        screenOccupancy = new IntSparseArrayMap<>();
-        newScreens = new IntArray();
+    public void setup() {
+        mModelHelper = new LauncherModelHelper();
+        mTargetContext = RuntimeEnvironment.application;
+        mIdp = InvariantDeviceProfile.INSTANCE.get(mTargetContext);
+        mIdp.numColumns = mIdp.numRows = 5;
+        mAppState = LauncherAppState.getInstance(mTargetContext);
 
-        idp.numColumns = 5;
-        idp.numRows = 5;
+        mExistingScreens = new IntArray();
+        mScreenOccupancy = new IntSparseArrayMap<>();
+        mNewScreens = new IntArray();
     }
 
     private AddWorkspaceItemsTask newTask(ItemInfo... items) {
@@ -70,17 +86,17 @@
         // Second screen has 2 holes of sizes 3x2 and 2x3
         setupWorkspaceWithHoles(nextId, 2, new Rect(2, 0, 5, 2), new Rect(0, 2, 2, 5));
 
-        int[] spaceFound = newTask()
-                .findSpaceForItem(appState, bgDataModel, existingScreens, newScreens, 1, 1);
+        int[] spaceFound = newTask().findSpaceForItem(
+                mAppState, mModelHelper.getBgDataModel(), mExistingScreens, mNewScreens, 1, 1);
         assertEquals(2, spaceFound[0]);
-        assertTrue(screenOccupancy.get(spaceFound[0])
+        assertTrue(mScreenOccupancy.get(spaceFound[0])
                 .isRegionVacant(spaceFound[1], spaceFound[2], 1, 1));
 
         // Find a larger space
-        spaceFound = newTask()
-                .findSpaceForItem(appState, bgDataModel, existingScreens, newScreens, 2, 3);
+        spaceFound = newTask().findSpaceForItem(
+                mAppState, mModelHelper.getBgDataModel(), mExistingScreens, mNewScreens, 2, 3);
         assertEquals(2, spaceFound[0]);
-        assertTrue(screenOccupancy.get(spaceFound[0])
+        assertTrue(mScreenOccupancy.get(spaceFound[0])
                 .isRegionVacant(spaceFound[1], spaceFound[2], 2, 3));
     }
 
@@ -89,11 +105,11 @@
         // First screen has 2 holes of sizes 3x2 and 2x3
         setupWorkspaceWithHoles(1, 1, new Rect(2, 0, 5, 2), new Rect(0, 2, 2, 5));
 
-        IntArray oldScreens = existingScreens.clone();
-        int[] spaceFound = newTask()
-                .findSpaceForItem(appState, bgDataModel, existingScreens, newScreens, 3, 3);
+        IntArray oldScreens = mExistingScreens.clone();
+        int[] spaceFound = newTask().findSpaceForItem(
+                mAppState, mModelHelper.getBgDataModel(), mExistingScreens, mNewScreens, 3, 3);
         assertFalse(oldScreens.contains(spaceFound[0]));
-        assertTrue(newScreens.contains(spaceFound[0]));
+        assertTrue(mNewScreens.contains(spaceFound[0]));
     }
 
     @Test
@@ -105,11 +121,14 @@
         setupWorkspaceWithHoles(1, 1, new Rect(2, 2, 3, 3));
 
         // Nothing was added
-        assertTrue(executeTaskForTest(newTask(info)).isEmpty());
+        assertTrue(mModelHelper.executeTaskForTest(newTask(info)).isEmpty());
     }
 
     @Test
     public void testAddItem_some_items_added() throws Exception {
+        Callbacks callbacks = mock(Callbacks.class);
+        mModelHelper.getModel().initialize(callbacks);
+
         WorkspaceItemInfo info = new WorkspaceItemInfo();
         info.intent = new Intent().setComponent(mComponent1);
 
@@ -119,7 +138,7 @@
         // Setup a screen with a hole
         setupWorkspaceWithHoles(1, 1, new Rect(2, 2, 3, 3));
 
-        executeTaskForTest(newTask(info, info2)).get(0).run();
+        mModelHelper.executeTaskForTest(newTask(info, info2)).get(0).run();
         ArgumentCaptor<ArrayList> notAnimated = ArgumentCaptor.forClass(ArrayList.class);
         ArgumentCaptor<ArrayList> animated = ArgumentCaptor.forClass(ArrayList.class);
 
@@ -134,18 +153,23 @@
     }
 
     private int setupWorkspaceWithHoles(int startId, int screenId, Rect... holes) throws Exception {
-        GridOccupancy occupancy = new GridOccupancy(idp.numColumns, idp.numRows);
-        occupancy.markCells(0, 0, idp.numColumns, idp.numRows, true);
+        return mModelHelper.executeSimpleTask(
+                model -> writeWorkspaceWithHoles(model, startId, screenId, holes));
+    }
+
+    private int writeWorkspaceWithHoles(
+            BgDataModel bgDataModel, int startId, int screenId, Rect... holes) {
+        GridOccupancy occupancy = new GridOccupancy(mIdp.numColumns, mIdp.numRows);
+        occupancy.markCells(0, 0, mIdp.numColumns, mIdp.numRows, true);
         for (Rect r : holes) {
             occupancy.markCells(r, false);
         }
 
-        existingScreens.add(screenId);
-        screenOccupancy.append(screenId, occupancy);
+        mExistingScreens.add(screenId);
+        mScreenOccupancy.append(screenId, occupancy);
 
-        ExecutorService executor = Executors.newSingleThreadExecutor();
-        for (int x = 0; x < idp.numColumns; x++) {
-            for (int y = 0; y < idp.numRows; y++) {
+        for (int x = 0; x < mIdp.numColumns; x++) {
+            for (int y = 0; y < mIdp.numRows; y++) {
                 if (!occupancy.cells[x][y]) {
                     continue;
                 }
@@ -157,20 +181,15 @@
                 info.cellX = x;
                 info.cellY = y;
                 info.container = LauncherSettings.Favorites.CONTAINER_DESKTOP;
-                bgDataModel.addItem(targetContext, info, false);
+                bgDataModel.addItem(mTargetContext, info, false);
 
-                executor.execute(() -> {
-                    ContentWriter writer = new ContentWriter(targetContext);
-                    info.writeToValues(writer);
-                    writer.put(Favorites._ID, info.id);
-                    targetContext.getContentResolver().insert(Favorites.CONTENT_URI,
-                            writer.getValues(targetContext));
-                });
+                ContentWriter writer = new ContentWriter(mTargetContext);
+                info.writeToValues(writer);
+                writer.put(Favorites._ID, info.id);
+                mTargetContext.getContentResolver().insert(Favorites.CONTENT_URI,
+                        writer.getValues(mTargetContext));
             }
         }
-
-        executor.submit(() -> null).get();
-        executor.shutdown();
         return startId;
     }
 }
diff --git a/robolectric_tests/src/com/android/launcher3/model/BaseGridChangesTestCase.java b/robolectric_tests/src/com/android/launcher3/model/BaseGridChangesTestCase.java
deleted file mode 100644
index 07834fc..0000000
--- a/robolectric_tests/src/com/android/launcher3/model/BaseGridChangesTestCase.java
+++ /dev/null
@@ -1,121 +0,0 @@
-package com.android.launcher3.model;
-
-import android.content.ContentValues;
-import android.content.Context;
-import android.content.Intent;
-import android.database.sqlite.SQLiteDatabase;
-
-import com.android.launcher3.LauncherProvider;
-import com.android.launcher3.LauncherSettings;
-import com.android.launcher3.util.TestLauncherProvider;
-
-import org.junit.Before;
-import org.robolectric.Robolectric;
-import org.robolectric.RuntimeEnvironment;
-import org.robolectric.shadows.ShadowContentResolver;
-import org.robolectric.shadows.ShadowLog;
-
-public abstract class BaseGridChangesTestCase {
-
-
-    public static final int DESKTOP = LauncherSettings.Favorites.CONTAINER_DESKTOP;
-    public static final int HOTSEAT = LauncherSettings.Favorites.CONTAINER_HOTSEAT;
-
-    public static final int APP_ICON = LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
-    public static final int SHORTCUT = LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT;
-    public static final int NO__ICON = -1;
-
-    public static final String TEST_PACKAGE = "com.android.launcher3.validpackage";
-
-    public Context mContext;
-    public TestLauncherProvider mProvider;
-    public SQLiteDatabase mDb;
-
-    @Before
-    public void setUpBaseCase() {
-        ShadowLog.stream = System.out;
-
-        mContext = RuntimeEnvironment.application;
-        mProvider = Robolectric.setupContentProvider(TestLauncherProvider.class);
-        ShadowContentResolver.registerProviderInternal(LauncherProvider.AUTHORITY, mProvider);
-        mDb = mProvider.getDb();
-    }
-
-    /**
-     * Adds a dummy item in the DB.
-     * @param type {@link #APP_ICON} or {@link #SHORTCUT} or >= 2 for
-     *             folder (where the type represents the number of items in the folder).
-     */
-    public int addItem(int type, int screen, int container, int x, int y) {
-        int id = LauncherSettings.Settings.call(mContext.getContentResolver(),
-                LauncherSettings.Settings.METHOD_NEW_ITEM_ID)
-                .getInt(LauncherSettings.Settings.EXTRA_VALUE);
-
-        ContentValues values = new ContentValues();
-        values.put(LauncherSettings.Favorites._ID, id);
-        values.put(LauncherSettings.Favorites.CONTAINER, container);
-        values.put(LauncherSettings.Favorites.SCREEN, screen);
-        values.put(LauncherSettings.Favorites.CELLX, x);
-        values.put(LauncherSettings.Favorites.CELLY, y);
-        values.put(LauncherSettings.Favorites.SPANX, 1);
-        values.put(LauncherSettings.Favorites.SPANY, 1);
-
-        if (type == APP_ICON || type == SHORTCUT) {
-            values.put(LauncherSettings.Favorites.ITEM_TYPE, type);
-            values.put(LauncherSettings.Favorites.INTENT,
-                    new Intent(Intent.ACTION_MAIN).setPackage(TEST_PACKAGE).toUri(0));
-        } else {
-            values.put(LauncherSettings.Favorites.ITEM_TYPE,
-                    LauncherSettings.Favorites.ITEM_TYPE_FOLDER);
-            // Add folder items.
-            for (int i = 0; i < type; i++) {
-                addItem(APP_ICON, 0, id, 0, 0);
-            }
-        }
-
-        mContext.getContentResolver().insert(LauncherSettings.Favorites.CONTENT_URI, values);
-        return id;
-    }
-
-    public int[][][] createGrid(int[][][] typeArray) {
-        return createGrid(typeArray, 1);
-    }
-
-    /**
-     * Initializes the DB with dummy elements to represent the provided grid structure.
-     * @param typeArray A 3d array of item types. {@see #addItem(int, long, long, int, int)} for
-     *                  type definitions. The first dimension represents the screens and the next
-     *                  two represent the workspace grid.
-     * @param startScreen First screen id from where the icons will be added.
-     * @return the same grid representation where each entry is the corresponding item id.
-     */
-    public int[][][] createGrid(int[][][] typeArray, int startScreen) {
-        LauncherSettings.Settings.call(mContext.getContentResolver(),
-                LauncherSettings.Settings.METHOD_CREATE_EMPTY_DB);
-        int[][][] ids = new int[typeArray.length][][];
-
-        for (int i = 0; i < typeArray.length; i++) {
-            // Add screen to DB
-            int screenId = startScreen + i;
-
-            // Keep the screen id counter up to date
-            LauncherSettings.Settings.call(mContext.getContentResolver(),
-                    LauncherSettings.Settings.METHOD_NEW_SCREEN_ID);
-
-            ids[i] = new int[typeArray[i].length][];
-            for (int y = 0; y < typeArray[i].length; y++) {
-                ids[i][y] = new int[typeArray[i][y].length];
-                for (int x = 0; x < typeArray[i][y].length; x++) {
-                    if (typeArray[i][y][x] < 0) {
-                        // Empty cell
-                        ids[i][y][x] = -1;
-                    } else {
-                        ids[i][y][x] = addItem(typeArray[i][y][x], screenId, DESKTOP, x, y);
-                    }
-                }
-            }
-        }
-
-        return ids;
-    }
-}
diff --git a/robolectric_tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java b/robolectric_tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java
deleted file mode 100644
index 012258d..0000000
--- a/robolectric_tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java
+++ /dev/null
@@ -1,231 +0,0 @@
-package com.android.launcher3.model;
-
-import static com.android.launcher3.shadows.ShadowLooperExecutor.reinitializeStaticExecutors;
-
-import static org.mockito.Matchers.anyBoolean;
-import static org.mockito.Mockito.atLeast;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
-import android.content.ComponentName;
-import android.content.Context;
-import android.content.Intent;
-import android.graphics.Bitmap;
-import android.graphics.Bitmap.Config;
-import android.graphics.Color;
-import android.os.Process;
-import android.os.UserHandle;
-
-import androidx.annotation.NonNull;
-
-import com.android.launcher3.AppFilter;
-import com.android.launcher3.AppInfo;
-import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.ItemInfo;
-import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.LauncherModel;
-import com.android.launcher3.LauncherModel.ModelUpdateTask;
-import com.android.launcher3.LauncherProvider;
-import com.android.launcher3.icons.BitmapInfo;
-import com.android.launcher3.icons.IconCache;
-import com.android.launcher3.icons.cache.CachingLogic;
-import com.android.launcher3.model.BgDataModel.Callbacks;
-import com.android.launcher3.pm.InstallSessionHelper;
-import com.android.launcher3.util.ComponentKey;
-import com.android.launcher3.util.TestLauncherProvider;
-
-import org.junit.Before;
-import org.mockito.ArgumentCaptor;
-import org.robolectric.Robolectric;
-import org.robolectric.RuntimeEnvironment;
-import org.robolectric.shadows.ShadowContentResolver;
-import org.robolectric.shadows.ShadowLog;
-
-import java.io.BufferedReader;
-import java.io.InputStreamReader;
-import java.lang.reflect.Field;
-import java.util.HashMap;
-import java.util.List;
-import java.util.concurrent.Executor;
-import java.util.function.Supplier;
-
-/**
- * Base class for writing tests for Model update tasks.
- */
-public class BaseModelUpdateTaskTestCase {
-
-    public final HashMap<Class, HashMap<String, Field>> fieldCache = new HashMap<>();
-    public TestLauncherProvider provider;
-
-    public Context targetContext;
-    public UserHandle myUser;
-
-    public InvariantDeviceProfile idp;
-    public LauncherAppState appState;
-    public LauncherModel model;
-    public ModelWriter modelWriter;
-    public MyIconCache iconCache;
-
-    public BgDataModel bgDataModel;
-    public AllAppsList allAppsList;
-    public Callbacks callbacks;
-
-    @Before
-    public void setUp() throws Exception {
-        ShadowLog.stream = System.out;
-        reinitializeStaticExecutors();
-        InstallSessionHelper.INSTANCE.initializeForTesting(null);
-
-        provider = Robolectric.setupContentProvider(TestLauncherProvider.class);
-        ShadowContentResolver.registerProviderInternal(LauncherProvider.AUTHORITY, provider);
-
-        callbacks = mock(Callbacks.class);
-        appState = mock(LauncherAppState.class);
-        model = mock(LauncherModel.class);
-        modelWriter = mock(ModelWriter.class);
-
-        LauncherAppState.INSTANCE.initializeForTesting(appState);
-        when(appState.getModel()).thenReturn(model);
-        when(model.getWriter(anyBoolean(), anyBoolean())).thenReturn(modelWriter);
-        when(model.getCallback()).thenReturn(callbacks);
-
-        myUser = Process.myUserHandle();
-
-        bgDataModel = new BgDataModel();
-        targetContext = RuntimeEnvironment.application;
-
-        idp = new InvariantDeviceProfile();
-        iconCache = new MyIconCache(targetContext, idp);
-
-        allAppsList = new AllAppsList(iconCache, new AppFilter());
-
-        when(appState.getIconCache()).thenReturn(iconCache);
-        when(appState.getInvariantDeviceProfile()).thenReturn(idp);
-        when(appState.getContext()).thenReturn(targetContext);
-    }
-
-    /**
-     * Synchronously executes the task and returns all the UI callbacks posted.
-     */
-    public List<Runnable> executeTaskForTest(ModelUpdateTask task) throws Exception {
-        when(model.isModelLoaded()).thenReturn(true);
-
-        Executor mockExecutor = mock(Executor.class);
-
-        task.init(appState, model, bgDataModel, allAppsList, mockExecutor);
-        task.run();
-        ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
-        verify(mockExecutor, atLeast(0)).execute(captor.capture());
-
-        return captor.getAllValues();
-    }
-
-    /**
-     * Initializes mock data for the test.
-     */
-    public void initializeData(String resourceName) throws Exception {
-        try (BufferedReader reader = new BufferedReader(new InputStreamReader(
-                this.getClass().getResourceAsStream(resourceName)))) {
-            String line;
-            HashMap<String, Class> classMap = new HashMap<>();
-            while((line = reader.readLine()) != null) {
-                line = line.trim();
-                if (line.startsWith("#") || line.isEmpty()) {
-                    continue;
-                }
-                String[] commands = line.split(" ");
-                switch (commands[0]) {
-                    case "classMap":
-                        classMap.put(commands[1], Class.forName(commands[2]));
-                        break;
-                    case "bgItem":
-                        bgDataModel.addItem(targetContext,
-                                (ItemInfo) initItem(classMap.get(commands[1]), commands, 2), false);
-                        break;
-                    case "allApps":
-                        allAppsList.add((AppInfo) initItem(AppInfo.class, commands, 1), null);
-                        break;
-                }
-            }
-        }
-    }
-
-    private Object initItem(Class clazz, String[] fieldDef, int startIndex) throws Exception {
-        HashMap<String, Field> cache = fieldCache.get(clazz);
-        if (cache == null) {
-            cache = new HashMap<>();
-            Class c = clazz;
-            while (c != null) {
-                for (Field f : c.getDeclaredFields()) {
-                    f.setAccessible(true);
-                    cache.put(f.getName(), f);
-                }
-                c = c.getSuperclass();
-            }
-            fieldCache.put(clazz, cache);
-        }
-
-        Object item = clazz.newInstance();
-        for (int i = startIndex; i < fieldDef.length; i++) {
-            String[] fieldData = fieldDef[i].split("=", 2);
-            Field f = cache.get(fieldData[0]);
-            Class type = f.getType();
-            if (type == int.class || type == long.class) {
-                f.set(item, Integer.parseInt(fieldData[1]));
-            } else if (type == CharSequence.class || type == String.class) {
-                f.set(item, fieldData[1]);
-            } else if (type == Intent.class) {
-                if (!fieldData[1].startsWith("#Intent")) {
-                    fieldData[1] = "#Intent;" + fieldData[1] + ";end";
-                }
-                f.set(item, Intent.parseUri(fieldData[1], 0));
-            } else if (type == ComponentName.class) {
-                f.set(item, ComponentName.unflattenFromString(fieldData[1]));
-            } else {
-                throw new Exception("Added parsing logic for "
-                        + f.getName() + " of type " + f.getType());
-            }
-        }
-        return item;
-    }
-
-    public static class MyIconCache extends IconCache {
-
-        private final HashMap<ComponentKey, CacheEntry> mCache = new HashMap<>();
-
-        public MyIconCache(Context context, InvariantDeviceProfile idp) {
-            super(context, idp);
-        }
-
-        @Override
-        protected <T> CacheEntry cacheLocked(
-                @NonNull ComponentName componentName,
-                UserHandle user, @NonNull Supplier<T> infoProvider,
-                @NonNull CachingLogic<T> cachingLogic,
-                boolean usePackageIcon, boolean useLowResIcon) {
-            CacheEntry entry = mCache.get(new ComponentKey(componentName, user));
-            if (entry == null) {
-                entry = new CacheEntry();
-                entry.bitmap = getDefaultIcon(user);
-            }
-            return entry;
-        }
-
-        public void addCache(ComponentName key, String title) {
-            CacheEntry entry = new CacheEntry();
-            entry.bitmap = BitmapInfo.of(newIcon(), Color.RED);
-            entry.title = title;
-            mCache.put(new ComponentKey(key, Process.myUserHandle()), entry);
-        }
-
-        public Bitmap newIcon() {
-            return Bitmap.createBitmap(1, 1, Config.ARGB_8888);
-        }
-
-        @Override
-        public synchronized BitmapInfo getDefaultIcon(UserHandle user) {
-            return BitmapInfo.fromBitmap(newIcon());
-        }
-    }
-}
diff --git a/robolectric_tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java b/robolectric_tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
index 69c5b00..f128e24 100644
--- a/robolectric_tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
@@ -5,15 +5,34 @@
 import static org.junit.Assert.assertNotSame;
 import static org.junit.Assert.assertTrue;
 
+import android.content.ComponentName;
+import android.content.Context;
+import android.content.pm.PackageInfo;
+import android.graphics.Bitmap;
+import android.graphics.Bitmap.Config;
+import android.graphics.Color;
+import android.os.Process;
+import android.os.UserHandle;
+import android.os.UserManager;
+
+import androidx.annotation.NonNull;
+
 import com.android.launcher3.AppInfo;
 import com.android.launcher3.ItemInfo;
+import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.WorkspaceItemInfo;
 import com.android.launcher3.icons.BitmapInfo;
+import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.CachingLogic;
+import com.android.launcher3.util.LauncherModelHelper;
+import com.android.launcher3.util.LauncherRoboTestRunner;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.annotation.LooperMode.Mode;
 
 import java.util.Arrays;
 import java.util.HashSet;
@@ -21,40 +40,73 @@
 /**
  * Tests for {@link CacheDataUpdatedTask}
  */
-@RunWith(RobolectricTestRunner.class)
-public class CacheDataUpdatedTaskTest extends BaseModelUpdateTaskTestCase {
+@RunWith(LauncherRoboTestRunner.class)
+@LooperMode(Mode.PAUSED)
+public class CacheDataUpdatedTaskTest {
 
     private static final String NEW_LABEL_PREFIX = "new-label-";
 
+    private LauncherModelHelper mModelHelper;
+
     @Before
-    public void initData() throws Exception {
-        initializeData("/cache_data_updated_task_data.txt");
+    public void setup() throws Exception {
+        mModelHelper = new LauncherModelHelper();
+        mModelHelper.initializeData("/cache_data_updated_task_data.txt");
+
         // Add dummy entries in the cache to simulate update
-        for (ItemInfo info : bgDataModel.itemsIdMap) {
-            iconCache.addCache(info.getTargetComponent(), NEW_LABEL_PREFIX + info.id);
+        Context context = RuntimeEnvironment.application;
+        IconCache iconCache = LauncherAppState.getInstance(context).getIconCache();
+        CachingLogic<ItemInfo> dummyLogic = new CachingLogic<ItemInfo>() {
+            @Override
+            public ComponentName getComponent(ItemInfo info) {
+                return info.getTargetComponent();
+            }
+
+            @Override
+            public UserHandle getUser(ItemInfo info) {
+                return info.user;
+            }
+
+            @Override
+            public CharSequence getLabel(ItemInfo info) {
+                return NEW_LABEL_PREFIX + info.id;
+            }
+
+            @NonNull
+            @Override
+            public BitmapInfo loadIcon(Context context, ItemInfo info) {
+                return BitmapInfo.of(Bitmap.createBitmap(1, 1, Config.ARGB_8888), Color.RED);
+            }
+        };
+
+        UserManager um = context.getSystemService(UserManager.class);
+        for (ItemInfo info : mModelHelper.getBgDataModel().itemsIdMap) {
+            iconCache.addIconToDBAndMemCache(info, dummyLogic, new PackageInfo(),
+                    um.getSerialNumberForUser(info.user), true);
         }
     }
 
     private CacheDataUpdatedTask newTask(int op, String... pkg) {
-        return new CacheDataUpdatedTask(op, myUser, new HashSet<>(Arrays.asList(pkg)));
+        return new CacheDataUpdatedTask(op, Process.myUserHandle(),
+                new HashSet<>(Arrays.asList(pkg)));
     }
 
     @Test
     public void testCacheUpdate_update_apps() throws Exception {
         // Clear all icons from apps list so that its easy to check what was updated
-        for (AppInfo info : allAppsList.data) {
+        for (AppInfo info : mModelHelper.getAllAppsList().data) {
             info.bitmap = BitmapInfo.LOW_RES_INFO;
         }
 
-        executeTaskForTest(newTask(CacheDataUpdatedTask.OP_CACHE_UPDATE, "app1"));
+        mModelHelper.executeTaskForTest(newTask(CacheDataUpdatedTask.OP_CACHE_UPDATE, "app1"));
 
         // Verify that only the app icons of app1 (id 1 & 2) are updated. Custom shortcut (id 7)
         // is not updated
         verifyUpdate(1, 2);
 
         // Verify that only app1 var updated in allAppsList
-        assertFalse(allAppsList.data.isEmpty());
-        for (AppInfo info : allAppsList.data) {
+        assertFalse(mModelHelper.getAllAppsList().data.isEmpty());
+        for (AppInfo info : mModelHelper.getAllAppsList().data) {
             if (info.componentName.getPackageName().equals("app1")) {
                 assertFalse(info.bitmap.isNullOrLowRes());
             } else {
@@ -65,7 +117,7 @@
 
     @Test
     public void testSessionUpdate_ignores_normal_apps() throws Exception {
-        executeTaskForTest(newTask(CacheDataUpdatedTask.OP_SESSION_UPDATE, "app1"));
+        mModelHelper.executeTaskForTest(newTask(CacheDataUpdatedTask.OP_SESSION_UPDATE, "app1"));
 
         // app1 has no restored shortcuts. Verify that nothing was updated.
         verifyUpdate();
@@ -73,7 +125,7 @@
 
     @Test
     public void testSessionUpdate_updates_pending_apps() throws Exception {
-        executeTaskForTest(newTask(CacheDataUpdatedTask.OP_SESSION_UPDATE, "app3"));
+        mModelHelper.executeTaskForTest(newTask(CacheDataUpdatedTask.OP_SESSION_UPDATE, "app3"));
 
         // app3 has only restored apps (id 5, 6) and shortcuts (id 9). Verify that only apps were
         // were updated
@@ -82,7 +134,7 @@
 
     private void verifyUpdate(Integer... idsUpdated) {
         HashSet<Integer> updates = new HashSet<>(Arrays.asList(idsUpdated));
-        for (ItemInfo info : bgDataModel.itemsIdMap) {
+        for (ItemInfo info : mModelHelper.getBgDataModel().itemsIdMap) {
             if (updates.contains(info.id)) {
                 assertEquals(NEW_LABEL_PREFIX + info.id, info.title);
                 assertFalse(((WorkspaceItemInfo) info).bitmap.isNullOrLowRes());
diff --git a/robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java b/robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
index b7340cf..1442c55 100644
--- a/robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
@@ -36,11 +36,11 @@
 import com.android.launcher3.LauncherProvider.DatabaseHelper;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.R;
+import com.android.launcher3.util.LauncherRoboTestRunner;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.robolectric.RobolectricTestRunner;
 import org.robolectric.RuntimeEnvironment;
 
 import java.io.File;
@@ -48,7 +48,7 @@
 /**
  * Tests for {@link DbDowngradeHelper}
  */
-@RunWith(RobolectricTestRunner.class)
+@RunWith(LauncherRoboTestRunner.class)
 public class DbDowngradeHelperTest {
 
     private static final String SCHEMA_FILE = "test_schema.json";
diff --git a/robolectric_tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java b/robolectric_tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
index 68713d8..e0ddcb1 100644
--- a/robolectric_tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
@@ -22,6 +22,7 @@
 import static org.robolectric.util.ReflectionHelpers.setField;
 
 import android.content.ComponentName;
+import android.content.Context;
 import android.content.pm.PackageInstaller;
 import android.content.pm.PackageInstaller.SessionInfo;
 import android.content.pm.PackageInstaller.SessionParams;
@@ -31,23 +32,20 @@
 import com.android.launcher3.FolderInfo;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.ItemInfo;
+import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherProvider;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.icons.BitmapInfo;
-import com.android.launcher3.pm.InstallSessionHelper;
-import com.android.launcher3.shadows.LShadowLauncherApps;
-import com.android.launcher3.shadows.LShadowUserManager;
-import com.android.launcher3.shadows.ShadowLooperExecutor;
+import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.util.Executors;
 import com.android.launcher3.util.LauncherLayoutBuilder;
-import com.android.launcher3.widget.custom.CustomWidgetManager;
+import com.android.launcher3.util.LauncherModelHelper;
+import com.android.launcher3.util.LauncherRoboTestRunner;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.annotation.Config;
+import org.robolectric.RuntimeEnvironment;
 import org.robolectric.annotation.LooperMode;
 import org.robolectric.annotation.LooperMode.Mode;
 import org.robolectric.shadows.ShadowPackageManager;
@@ -61,10 +59,9 @@
 /**
  * Tests for layout parser for remote layout
  */
-@RunWith(RobolectricTestRunner.class)
-@Config(shadows = {LShadowUserManager.class, LShadowLauncherApps.class, ShadowLooperExecutor.class})
+@RunWith(LauncherRoboTestRunner.class)
 @LooperMode(Mode.PAUSED)
-public class DefaultLayoutProviderTest extends BaseModelUpdateTaskTestCase {
+public class DefaultLayoutProviderTest {
 
     private static final String SETTINGS_APP = "com.android.settings";
     private static final String TEST_PROVIDER_AUTHORITY =
@@ -73,40 +70,37 @@
     private static final int BITMAP_SIZE = 10;
     private static final int GRID_SIZE = 4;
 
+    private LauncherModelHelper mModelHelper;
+    private Context mTargetContext;
+    private InvariantDeviceProfile mIdp;
+
     @Before
-    public void setUp() throws Exception {
-        super.setUp();
-        InvariantDeviceProfile.INSTANCE.initializeForTesting(idp);
-        CustomWidgetManager.INSTANCE.initializeForTesting(mock(CustomWidgetManager.class));
+    public void setUp() {
+        mModelHelper = new LauncherModelHelper();
+        mTargetContext = RuntimeEnvironment.application;
 
-        idp.numRows = idp.numColumns = idp.numHotseatIcons = GRID_SIZE;
-        idp.iconBitmapSize = BITMAP_SIZE;
+        mIdp = InvariantDeviceProfile.INSTANCE.get(mTargetContext);
+        mIdp.numRows = mIdp.numColumns = mIdp.numHotseatIcons = GRID_SIZE;
+        mIdp.iconBitmapSize = BITMAP_SIZE;
 
-        provider.setAllowLoadDefaultFavorites(true);
-        Settings.Secure.putString(targetContext.getContentResolver(),
+        mModelHelper.provider.setAllowLoadDefaultFavorites(true);
+        Settings.Secure.putString(mTargetContext.getContentResolver(),
                 "launcher3.layout.provider", TEST_PROVIDER_AUTHORITY);
 
-        ShadowPackageManager spm = shadowOf(targetContext.getPackageManager());
+        ShadowPackageManager spm = shadowOf(mTargetContext.getPackageManager());
         spm.addProviderIfNotPresent(new ComponentName("com.test", "Dummy")).authority =
                 TEST_PROVIDER_AUTHORITY;
         spm.addActivityIfNotPresent(new ComponentName(SETTINGS_APP, SETTINGS_APP));
     }
 
-    @After
-    public void cleanup() {
-        InvariantDeviceProfile.INSTANCE.initializeForTesting(null);
-        CustomWidgetManager.INSTANCE.initializeForTesting(null);
-        InstallSessionHelper.INSTANCE.initializeForTesting(null);
-    }
-
     @Test
     public void testCustomProfileLoaded_with_icon_on_hotseat() throws Exception {
         writeLayoutAndLoad(new LauncherLayoutBuilder().atHotseat(0)
                 .putApp(SETTINGS_APP, SETTINGS_APP));
 
         // Verify one item in hotseat
-        assertEquals(1, bgDataModel.workspaceItems.size());
-        ItemInfo info = bgDataModel.workspaceItems.get(0);
+        assertEquals(1, mModelHelper.getBgDataModel().workspaceItems.size());
+        ItemInfo info = mModelHelper.getBgDataModel().workspaceItems.get(0);
         assertEquals(LauncherSettings.Favorites.CONTAINER_HOTSEAT, info.container);
         assertEquals(LauncherSettings.Favorites.ITEM_TYPE_APPLICATION, info.itemType);
     }
@@ -120,8 +114,8 @@
                 .build());
 
         // Verify folder
-        assertEquals(1, bgDataModel.workspaceItems.size());
-        ItemInfo info = bgDataModel.workspaceItems.get(0);
+        assertEquals(1, mModelHelper.getBgDataModel().workspaceItems.size());
+        ItemInfo info = mModelHelper.getBgDataModel().workspaceItems.get(0);
         assertEquals(LauncherSettings.Favorites.ITEM_TYPE_FOLDER, info.itemType);
         assertEquals(3, ((FolderInfo) info).contents.size());
     }
@@ -134,7 +128,7 @@
         SessionParams params = new SessionParams(SessionParams.MODE_FULL_INSTALL);
         params.setAppPackageName(pendingAppPkg);
 
-        PackageInstaller installer = targetContext.getPackageManager().getPackageInstaller();
+        PackageInstaller installer = mTargetContext.getPackageManager().getPackageInstaller();
         int sessionId = installer.createSession(params);
         SessionInfo sessionInfo = installer.getSessionInfo(sessionId);
         setField(sessionInfo, "installerPackageName", "com.test");
@@ -144,8 +138,8 @@
                 .putWidget(pendingAppPkg, "DummyWidget", 2, 2));
 
         // Verify widget
-        assertEquals(1, bgDataModel.appWidgets.size());
-        ItemInfo info = bgDataModel.appWidgets.get(0);
+        assertEquals(1, mModelHelper.getBgDataModel().appWidgets.size());
+        ItemInfo info = mModelHelper.getBgDataModel().appWidgets.get(0);
         assertEquals(LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET, info.itemType);
         assertEquals(2, info.spanX);
         assertEquals(2, info.spanY);
@@ -155,13 +149,21 @@
         ByteArrayOutputStream bos = new ByteArrayOutputStream();
         builder.build(new OutputStreamWriter(bos));
 
-        Uri layoutUri = LauncherProvider.getLayoutUri(TEST_PROVIDER_AUTHORITY, targetContext);
-        shadowOf(targetContext.getContentResolver()).registerInputStream(layoutUri,
+        Uri layoutUri = LauncherProvider.getLayoutUri(TEST_PROVIDER_AUTHORITY, mTargetContext);
+        shadowOf(mTargetContext.getContentResolver()).registerInputStream(layoutUri,
                 new ByteArrayInputStream(bos.toByteArray()));
 
-        LoaderResults results = new LoaderResults(appState, bgDataModel, allAppsList, 0,
-                new WeakReference<>(callbacks));
-        LoaderTask task = new LoaderTask(appState, allAppsList, bgDataModel, results);
+        LoaderResults results = new LoaderResults(
+                LauncherAppState.getInstance(mTargetContext),
+                mModelHelper.getBgDataModel(),
+                mModelHelper.getAllAppsList(),
+                0,
+                new WeakReference<>(mock(Callbacks.class)));
+        LoaderTask task = new LoaderTask(
+                LauncherAppState.getInstance(mTargetContext),
+                mModelHelper.getAllAppsList(),
+                mModelHelper.getBgDataModel(),
+                results);
         Executors.MODEL_EXECUTOR.submit(() -> task.loadWorkspace(new ArrayList<>())).get();
     }
 }
diff --git a/robolectric_tests/src/com/android/launcher3/model/GridBackupTableTest.java b/robolectric_tests/src/com/android/launcher3/model/GridBackupTableTest.java
index 53287a9..f46b849 100644
--- a/robolectric_tests/src/com/android/launcher3/model/GridBackupTableTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/GridBackupTableTest.java
@@ -6,33 +6,53 @@
 import static com.android.launcher3.LauncherSettings.Favorites.BACKUP_TABLE_NAME;
 import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.provider.LauncherDbUtils.tableExists;
+import static com.android.launcher3.util.LauncherModelHelper.APP_ICON;
+import static com.android.launcher3.util.LauncherModelHelper.DESKTOP;
+import static com.android.launcher3.util.LauncherModelHelper.NO__ICON;
+import static com.android.launcher3.util.LauncherModelHelper.SHORTCUT;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
 import android.content.ContentValues;
+import android.content.Context;
+import android.database.sqlite.SQLiteDatabase;
 import android.graphics.Point;
 
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.LauncherSettings.Settings;
+import com.android.launcher3.util.LauncherModelHelper;
+import com.android.launcher3.util.LauncherRoboTestRunner;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
 
 /**
  * Unit tests for {@link GridBackupTable}
  */
-@RunWith(RobolectricTestRunner.class)
-public class GridBackupTableTest extends BaseGridChangesTestCase {
+@RunWith(LauncherRoboTestRunner.class)
+public class GridBackupTableTest {
 
     private static final int BACKUP_ITEM_COUNT = 12;
 
+    private LauncherModelHelper mModelHelper;
+    private Context mContext;
+    private SQLiteDatabase mDb;
+
     @Before
-    public void setupGridData() {
-        createGrid(new int[][][]{{
+    public void setUp() {
+        mModelHelper = new LauncherModelHelper();
+        mContext = RuntimeEnvironment.application;
+        mDb = mModelHelper.provider.getDb();
+
+        setupGridData();
+    }
+
+    private void setupGridData() {
+        mModelHelper.createGrid(new int[][][]{{
                 { APP_ICON, APP_ICON, SHORTCUT, SHORTCUT},
                 { SHORTCUT, SHORTCUT, NO__ICON, NO__ICON},
                 { NO__ICON, NO__ICON, SHORTCUT, SHORTCUT},
@@ -81,7 +101,7 @@
 
         assertTrue(tableExists(mDb, BACKUP_TABLE_NAME));
 
-        addItem(1, 2, DESKTOP, 1, 1);
+        mModelHelper.addItem(1, 2, DESKTOP, 1, 1);
         assertFalse(tableExists(mDb, BACKUP_TABLE_NAME));
     }
 
diff --git a/robolectric_tests/src/com/android/launcher3/model/GridSizeMigrationTaskTest.java b/robolectric_tests/src/com/android/launcher3/model/GridSizeMigrationTaskTest.java
index 53f6a06..8dd7588 100644
--- a/robolectric_tests/src/com/android/launcher3/model/GridSizeMigrationTaskTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/GridSizeMigrationTaskTest.java
@@ -1,25 +1,31 @@
 package com.android.launcher3.model;
 
 import static com.android.launcher3.model.GridSizeMigrationTask.getWorkspaceScreenIds;
+import static com.android.launcher3.util.LauncherModelHelper.APP_ICON;
+import static com.android.launcher3.util.LauncherModelHelper.HOTSEAT;
+import static com.android.launcher3.util.LauncherModelHelper.SHORTCUT;
+import static com.android.launcher3.util.LauncherModelHelper.TEST_PACKAGE;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
+import android.content.Context;
 import android.database.Cursor;
+import android.database.sqlite.SQLiteDatabase;
 import android.graphics.Point;
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.config.FeatureFlags;
-import com.android.launcher3.config.FlagOverrideRule;
 import com.android.launcher3.model.GridSizeMigrationTask.MultiStepMigrationTask;
 import com.android.launcher3.util.IntArray;
+import com.android.launcher3.util.LauncherModelHelper;
+import com.android.launcher3.util.LauncherRoboTestRunner;
 
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
 
 import java.util.HashSet;
 import java.util.LinkedList;
@@ -27,30 +33,35 @@
 /**
  * Unit tests for {@link GridSizeMigrationTask}
  */
-@RunWith(RobolectricTestRunner.class)
-public class GridSizeMigrationTaskTest extends BaseGridChangesTestCase {
+@RunWith(LauncherRoboTestRunner.class)
+public class GridSizeMigrationTaskTest {
 
-    @Rule
-    public final FlagOverrideRule flags = new FlagOverrideRule();
+    private LauncherModelHelper mModelHelper;
+    private Context mContext;
+    private SQLiteDatabase mDb;
 
     private HashSet<String> mValidPackages;
     private InvariantDeviceProfile mIdp;
 
     @Before
     public void setUp() {
+        mModelHelper = new LauncherModelHelper();
+        mContext = RuntimeEnvironment.application;
+        mDb = mModelHelper.provider.getDb();
+
         mValidPackages = new HashSet<>();
         mValidPackages.add(TEST_PACKAGE);
-        mIdp = new InvariantDeviceProfile();
+        mIdp = InvariantDeviceProfile.INSTANCE.get(mContext);
     }
 
     @Test
     public void testHotseatMigration_apps_dropped() throws Exception {
         int[] hotseatItems = {
-                addItem(APP_ICON, 0, HOTSEAT, 0, 0),
-                addItem(SHORTCUT, 1, HOTSEAT, 0, 0),
+                mModelHelper.addItem(APP_ICON, 0, HOTSEAT, 0, 0),
+                mModelHelper.addItem(SHORTCUT, 1, HOTSEAT, 0, 0),
                 -1,
-                addItem(SHORTCUT, 3, HOTSEAT, 0, 0),
-                addItem(APP_ICON, 4, HOTSEAT, 0, 0),
+                mModelHelper.addItem(SHORTCUT, 3, HOTSEAT, 0, 0),
+                mModelHelper.addItem(APP_ICON, 4, HOTSEAT, 0, 0),
         };
 
         mIdp.numHotseatIcons = 3;
@@ -63,11 +74,11 @@
     @Test
     public void testHotseatMigration_shortcuts_dropped() throws Exception {
         int[] hotseatItems = {
-                addItem(APP_ICON, 0, HOTSEAT, 0, 0),
-                addItem(30, 1, HOTSEAT, 0, 0),
+                mModelHelper.addItem(APP_ICON, 0, HOTSEAT, 0, 0),
+                mModelHelper.addItem(30, 1, HOTSEAT, 0, 0),
                 -1,
-                addItem(SHORTCUT, 3, HOTSEAT, 0, 0),
-                addItem(10, 4, HOTSEAT, 0, 0),
+                mModelHelper.addItem(SHORTCUT, 3, HOTSEAT, 0, 0),
+                mModelHelper.addItem(10, 4, HOTSEAT, 0, 0),
         };
 
         mIdp.numHotseatIcons = 3;
@@ -109,7 +120,7 @@
 
     @Test
     public void testWorkspace_empty_row_column_removed() throws Exception {
-        int[][][] ids = createGrid(new int[][][]{{
+        int[][][] ids = mModelHelper.createGrid(new int[][][]{{
                 {  0,  0, -1,  1},
                 {  3,  1, -1,  4},
                 { -1, -1, -1, -1},
@@ -129,7 +140,7 @@
 
     @Test
     public void testWorkspace_new_screen_created() throws Exception {
-        int[][][] ids = createGrid(new int[][][]{{
+        int[][][] ids = mModelHelper.createGrid(new int[][][]{{
                 {  0,  0,  0,  1},
                 {  3,  1,  0,  4},
                 { -1, -1, -1, -1},
@@ -151,7 +162,7 @@
 
     @Test
     public void testWorkspace_items_merged_in_next_screen() throws Exception {
-        int[][][] ids = createGrid(new int[][][]{{
+        int[][][] ids = mModelHelper.createGrid(new int[][][]{{
                 {  0,  0,  0,  1},
                 {  3,  1,  0,  4},
                 { -1, -1, -1, -1},
@@ -181,7 +192,7 @@
     public void testWorkspace_items_not_merged_in_next_screen() throws Exception {
         // First screen has 2 items that need to be moved, but second screen has only one
         // empty space after migration (top-left corner)
-        int[][][] ids = createGrid(new int[][][]{{
+        int[][][] ids = mModelHelper.createGrid(new int[][][]{{
                 {  0,  0,  0,  1},
                 {  3,  1,  0,  4},
                 { -1, -1, -1, -1},
@@ -217,7 +228,7 @@
         }
         // The first screen has one item on the 4th column which needs moving, as the first row
         // will be kept empty.
-        int[][][] ids = createGrid(new int[][][]{{
+        int[][][] ids = mModelHelper.createGrid(new int[][][]{{
                 { -1, -1, -1, -1},
                 {  3,  1,  7,  0},
                 {  8,  7,  7, -1},
@@ -244,7 +255,7 @@
             return;
         }
         // Items will get moved to the next screen to keep the first screen empty.
-        int[][][] ids = createGrid(new int[][][]{{
+        int[][][] ids = mModelHelper.createGrid(new int[][][]{{
                 { -1, -1, -1, -1},
                 {  0,  1,  0,  0},
                 {  8,  7,  7, -1},
diff --git a/tests/src/com/android/launcher3/model/LoaderCursorTest.java b/robolectric_tests/src/com/android/launcher3/model/LoaderCursorTest.java
similarity index 76%
rename from tests/src/com/android/launcher3/model/LoaderCursorTest.java
rename to robolectric_tests/src/com/android/launcher3/model/LoaderCursorTest.java
index 0dcfaa8..4854314 100644
--- a/tests/src/com/android/launcher3/model/LoaderCursorTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/LoaderCursorTest.java
@@ -1,3 +1,19 @@
+/*
+ * Copyright (C) 2019 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.model;
 
 import static com.android.launcher3.LauncherSettings.Favorites.CELLX;
@@ -17,6 +33,7 @@
 import static com.android.launcher3.LauncherSettings.Favorites.SCREEN;
 import static com.android.launcher3.LauncherSettings.Favorites.TITLE;
 import static com.android.launcher3.LauncherSettings.Favorites._ID;
+import static com.android.launcher3.util.LauncherModelHelper.TEST_PACKAGE;
 
 import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertFalse;
@@ -24,43 +41,38 @@
 import static junit.framework.Assert.assertNull;
 import static junit.framework.Assert.assertTrue;
 
-import static org.mockito.Mockito.eq;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
+import static org.robolectric.Shadows.shadowOf;
 
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
 import android.content.pm.LauncherApps;
 import android.database.MatrixCursor;
-import android.graphics.Bitmap;
 import android.os.Process;
 
-import androidx.test.InstrumentationRegistry;
-import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
-
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.ItemInfo;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.WorkspaceItemInfo;
-import com.android.launcher3.icons.BitmapInfo;
-import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.util.Executors;
+import com.android.launcher3.util.LauncherRoboTestRunner;
 import com.android.launcher3.util.PackageManagerHelper;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.annotation.LooperMode.Mode;
 
 /**
  * Tests for {@link LoaderCursor}
  */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(LauncherRoboTestRunner.class)
+@LooperMode(Mode.PAUSED)
 public class LoaderCursorTest {
 
-    private LauncherAppState mMockApp;
-    private IconCache mMockIconCache;
+    private LauncherAppState mApp;
 
     private MatrixCursor mCursor;
     private InvariantDeviceProfile mIDP;
@@ -71,22 +83,18 @@
 
     @Before
     public void setup() {
-        mIDP = new InvariantDeviceProfile();
+        mContext = RuntimeEnvironment.application;
+        mIDP = InvariantDeviceProfile.INSTANCE.get(mContext);
+        mApp = LauncherAppState.getInstance(mContext);
+        mLauncherApps = mContext.getSystemService(LauncherApps.class);
+
         mCursor = new MatrixCursor(new String[] {
                 ICON, ICON_PACKAGE, ICON_RESOURCE, TITLE,
                 _ID, CONTAINER, ITEM_TYPE, PROFILE_ID,
                 SCREEN, CELLX, CELLY, RESTORED, INTENT
         });
-        mContext = InstrumentationRegistry.getTargetContext();
 
-        mMockApp = mock(LauncherAppState.class);
-        mMockIconCache = mock(IconCache.class);
-        when(mMockApp.getIconCache()).thenReturn(mMockIconCache);
-        when(mMockApp.getInvariantDeviceProfile()).thenReturn(mIDP);
-        when(mMockApp.getContext()).thenReturn(mContext);
-        mLauncherApps = mContext.getSystemService(LauncherApps.class);
-
-        mLoaderCursor = new LoaderCursor(mCursor, mMockApp);
+        mLoaderCursor = new LoaderCursor(mCursor, mApp);
         mLoaderCursor.allUsers.put(0, Process.myUserHandle());
     }
 
@@ -109,26 +117,31 @@
     }
 
     @Test
-    public void getAppShortcutInfo_dontAllowMissing_validComponent() {
+    public void getAppShortcutInfo_dontAllowMissing_validComponent() throws Exception {
+        ComponentName cn = new ComponentName(TEST_PACKAGE, TEST_PACKAGE);
+        shadowOf(mContext.getPackageManager()).addActivityIfNotPresent(cn);
+
         initCursor(ITEM_TYPE_APPLICATION, "");
         assertTrue(mLoaderCursor.moveToNext());
 
-        ComponentName cn = mLauncherApps.getActivityList(null, mLoaderCursor.user)
-                .get(0).getComponentName();
-        WorkspaceItemInfo info = mLoaderCursor.getAppShortcutInfo(
-                new Intent().setComponent(cn), false /* allowMissingTarget */, true);
+        WorkspaceItemInfo info = Executors.MODEL_EXECUTOR.submit(() ->
+                mLoaderCursor.getAppShortcutInfo(
+                        new Intent().setComponent(cn), false  /* allowMissingTarget */, true))
+                .get();
         assertNotNull(info);
         assertTrue(PackageManagerHelper.isLauncherAppTarget(info.intent));
     }
 
     @Test
-    public void getAppShortcutInfo_allowMissing_invalidComponent() {
+    public void getAppShortcutInfo_allowMissing_invalidComponent() throws Exception {
         initCursor(ITEM_TYPE_APPLICATION, "");
         assertTrue(mLoaderCursor.moveToNext());
 
         ComponentName cn = new ComponentName(mContext.getPackageName(), "dummy-do");
-        WorkspaceItemInfo info = mLoaderCursor.getAppShortcutInfo(
-                new Intent().setComponent(cn), true  /* allowMissingTarget */, true);
+        WorkspaceItemInfo info = Executors.MODEL_EXECUTOR.submit(() ->
+                mLoaderCursor.getAppShortcutInfo(
+                        new Intent().setComponent(cn), true  /* allowMissingTarget */, true))
+                .get();
         assertNotNull(info);
         assertTrue(PackageManagerHelper.isLauncherAppTarget(info.intent));
     }
@@ -138,11 +151,8 @@
         initCursor(ITEM_TYPE_SHORTCUT, "my-shortcut");
         assertTrue(mLoaderCursor.moveToNext());
 
-        Bitmap icon = Bitmap.createBitmap(1, 1, Bitmap.Config.ALPHA_8);
-        when(mMockIconCache.getDefaultIcon(eq(mLoaderCursor.user)))
-                .thenReturn(BitmapInfo.fromBitmap(icon));
         WorkspaceItemInfo info = mLoaderCursor.loadSimpleWorkspaceItem();
-        assertEquals(icon, info.bitmap.icon);
+        assertTrue(mApp.getIconCache().isDefaultIcon(info.bitmap, info.user));
         assertEquals("my-shortcut", info.title);
         assertEquals(ITEM_TYPE_SHORTCUT, info.itemType);
     }
diff --git a/robolectric_tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java b/robolectric_tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
index a1a4561..bd71f01 100644
--- a/robolectric_tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
@@ -6,11 +6,14 @@
 import com.android.launcher3.LauncherAppWidgetInfo;
 import com.android.launcher3.WorkspaceItemInfo;
 import com.android.launcher3.pm.PackageInstallInfo;
+import com.android.launcher3.util.LauncherModelHelper;
+import com.android.launcher3.util.LauncherRoboTestRunner;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.robolectric.RobolectricTestRunner;
+import org.robolectric.annotation.LooperMode;
+import org.robolectric.annotation.LooperMode.Mode;
 
 import java.util.Arrays;
 import java.util.HashSet;
@@ -18,12 +21,16 @@
 /**
  * Tests for {@link PackageInstallStateChangedTask}
  */
-@RunWith(RobolectricTestRunner.class)
-public class PackageInstallStateChangedTaskTest extends BaseModelUpdateTaskTestCase {
+@RunWith(LauncherRoboTestRunner.class)
+@LooperMode(Mode.PAUSED)
+public class PackageInstallStateChangedTaskTest {
+
+    private LauncherModelHelper mModelHelper;
 
     @Before
-    public void initData() throws Exception {
-        initializeData("/package_install_state_change_task_data.txt");
+    public void setup() throws Exception {
+        mModelHelper = new LauncherModelHelper();
+        mModelHelper.initializeData("/package_install_state_change_task_data.txt");
     }
 
     private PackageInstallStateChangedTask newTask(String pkg, int progress) {
@@ -35,7 +42,7 @@
 
     @Test
     public void testSessionUpdate_ignore_installed() throws Exception {
-        executeTaskForTest(newTask("app1", 30));
+        mModelHelper.executeTaskForTest(newTask("app1", 30));
 
         // No shortcuts were updated
         verifyProgressUpdate(0);
@@ -43,21 +50,21 @@
 
     @Test
     public void testSessionUpdate_shortcuts_updated() throws Exception {
-        executeTaskForTest(newTask("app3", 30));
+        mModelHelper.executeTaskForTest(newTask("app3", 30));
 
         verifyProgressUpdate(30, 5, 6, 7);
     }
 
     @Test
     public void testSessionUpdate_widgets_updated() throws Exception {
-        executeTaskForTest(newTask("app4", 30));
+        mModelHelper.executeTaskForTest(newTask("app4", 30));
 
         verifyProgressUpdate(30, 8, 9);
     }
 
     private void verifyProgressUpdate(int progress, Integer... idsUpdated) {
         HashSet<Integer> updates = new HashSet<>(Arrays.asList(idsUpdated));
-        for (ItemInfo info : bgDataModel.itemsIdMap) {
+        for (ItemInfo info : mModelHelper.getBgDataModel().itemsIdMap) {
             if (info instanceof WorkspaceItemInfo) {
                 assertEquals(updates.contains(info.id) ? progress: 0,
                         ((WorkspaceItemInfo) info).getInstallProgress());
diff --git a/robolectric_tests/src/com/android/launcher3/popup/PopupPopulatorTest.java b/robolectric_tests/src/com/android/launcher3/popup/PopupPopulatorTest.java
index 83bf7da..7612ae1 100644
--- a/robolectric_tests/src/com/android/launcher3/popup/PopupPopulatorTest.java
+++ b/robolectric_tests/src/com/android/launcher3/popup/PopupPopulatorTest.java
@@ -27,9 +27,10 @@
 
 import android.content.pm.ShortcutInfo;
 
+import com.android.launcher3.util.LauncherRoboTestRunner;
+
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.robolectric.RobolectricTestRunner;
 import org.robolectric.RuntimeEnvironment;
 
 import java.util.ArrayList;
@@ -39,7 +40,7 @@
 /**
  * Tests the sorting and filtering of shortcuts in {@link PopupPopulator}.
  */
-@RunWith(RobolectricTestRunner.class)
+@RunWith(LauncherRoboTestRunner.class)
 public class PopupPopulatorTest {
 
     @Test
diff --git a/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java b/robolectric_tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
similarity index 78%
rename from tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
rename to robolectric_tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
index 27990f4..7ef670c 100644
--- a/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
+++ b/robolectric_tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
@@ -1,3 +1,18 @@
+/*
+ * Copyright (C) 2019 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.provider;
 
 import static org.junit.Assert.assertEquals;
@@ -6,21 +21,18 @@
 import android.database.Cursor;
 import android.database.sqlite.SQLiteDatabase;
 
-import androidx.test.InstrumentationRegistry;
-import androidx.test.filters.MediumTest;
-import androidx.test.runner.AndroidJUnit4;
-
 import com.android.launcher3.LauncherProvider.DatabaseHelper;
 import com.android.launcher3.LauncherSettings.Favorites;
+import com.android.launcher3.util.LauncherRoboTestRunner;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.RuntimeEnvironment;
 
 /**
  * Tests for {@link RestoreDbTask}
  */
-@MediumTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(LauncherRoboTestRunner.class)
 public class RestoreDbTaskTest {
 
     @Test
@@ -83,7 +95,7 @@
         private final long mProfileId;
 
         MyDatabaseHelper(long profileId) {
-            super(InstrumentationRegistry.getContext(), null);
+            super(RuntimeEnvironment.application, null);
             mProfileId = profileId;
         }
 
diff --git a/robolectric_tests/src/com/android/launcher3/shadows/LShadowAppWidgetManager.java b/robolectric_tests/src/com/android/launcher3/shadows/LShadowAppWidgetManager.java
new file mode 100644
index 0000000..696ffd0
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/shadows/LShadowAppWidgetManager.java
@@ -0,0 +1,48 @@
+/*
+ * Copyright (C) 2019 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.shadows;
+
+import android.appwidget.AppWidgetManager;
+import android.appwidget.AppWidgetProviderInfo;
+import android.os.Process;
+import android.os.UserHandle;
+
+import org.robolectric.annotation.Implementation;
+import org.robolectric.annotation.Implements;
+import org.robolectric.shadows.ShadowAppWidgetManager;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+/**
+ * Extension of {@link ShadowAppWidgetManager} with missing shadow methods
+ */
+@Implements(value = AppWidgetManager.class)
+public class LShadowAppWidgetManager extends ShadowAppWidgetManager {
+
+    @Override
+    protected List<AppWidgetProviderInfo> getInstalledProviders() {
+        return getInstalledProvidersForProfile(null);
+    }
+
+    @Implementation
+    public List<AppWidgetProviderInfo> getInstalledProvidersForProfile(UserHandle profile) {
+        UserHandle user = profile == null ? Process.myUserHandle() : profile;
+        return super.getInstalledProviders().stream().filter(
+                info -> user.equals(info.getProfile())).collect(Collectors.toList());
+    }
+}
diff --git a/robolectric_tests/src/com/android/launcher3/shadows/LShadowBitmap.java b/robolectric_tests/src/com/android/launcher3/shadows/LShadowBitmap.java
new file mode 100644
index 0000000..abd90bb
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/shadows/LShadowBitmap.java
@@ -0,0 +1,36 @@
+/*
+ * Copyright (C) 2019 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.shadows;
+
+import android.graphics.Bitmap;
+import android.graphics.Paint;
+
+import org.robolectric.annotation.Implementation;
+import org.robolectric.annotation.Implements;
+import org.robolectric.shadows.ShadowBitmap;
+
+/**
+ * Extension of {@link ShadowBitmap} with missing shadow methods
+ */
+@Implements(value = Bitmap.class)
+public class LShadowBitmap extends ShadowBitmap {
+
+    @Implementation
+    protected Bitmap extractAlpha(Paint paint, int[] offsetXY) {
+        return extractAlpha();
+    }
+}
diff --git a/robolectric_tests/src/com/android/launcher3/shadows/LShadowLauncherApps.java b/robolectric_tests/src/com/android/launcher3/shadows/LShadowLauncherApps.java
index 204ec9b..ccbc18a 100644
--- a/robolectric_tests/src/com/android/launcher3/shadows/LShadowLauncherApps.java
+++ b/robolectric_tests/src/com/android/launcher3/shadows/LShadowLauncherApps.java
@@ -77,7 +77,7 @@
     protected LauncherActivityInfo resolveActivity(Intent intent, UserHandle user) {
         ResolveInfo ri = RuntimeEnvironment.application.getPackageManager()
                 .resolveActivity(intent, 0);
-        return getLauncherActivityInfo(ri.activityInfo);
+        return ri == null ? null : getLauncherActivityInfo(ri.activityInfo);
     }
 
     public LauncherActivityInfo getLauncherActivityInfo(ActivityInfo activityInfo) {
diff --git a/robolectric_tests/src/com/android/launcher3/shadows/ShadowLooperExecutor.java b/robolectric_tests/src/com/android/launcher3/shadows/ShadowLooperExecutor.java
index d56de3c..a3b7dc7 100644
--- a/robolectric_tests/src/com/android/launcher3/shadows/ShadowLooperExecutor.java
+++ b/robolectric_tests/src/com/android/launcher3/shadows/ShadowLooperExecutor.java
@@ -18,25 +18,16 @@
 
 import static com.android.launcher3.util.Executors.createAndStartNewLooper;
 
-import static org.robolectric.shadow.api.Shadow.invokeConstructor;
-import static org.robolectric.util.ReflectionHelpers.ClassParameter.from;
+import static org.robolectric.shadow.api.Shadow.directlyOn;
 import static org.robolectric.util.ReflectionHelpers.setField;
 
 import android.os.Handler;
-import android.os.Looper;
 
-import com.android.launcher3.util.Executors;
 import com.android.launcher3.util.LooperExecutor;
 
 import org.robolectric.annotation.Implementation;
 import org.robolectric.annotation.Implements;
 import org.robolectric.annotation.RealObject;
-import org.robolectric.util.ReflectionHelpers;
-
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.Set;
-import java.util.WeakHashMap;
 
 /**
  * Shadow for {@link LooperExecutor} to provide reset functionality for static executors.
@@ -44,25 +35,18 @@
 @Implements(value = LooperExecutor.class, isInAndroidSdk = false)
 public class ShadowLooperExecutor {
 
-    // Keep reference to all created Loopers so they can be torn down after test
-    private static Set<LooperExecutor> executors =
-            Collections.synchronizedSet(Collections.newSetFromMap(new WeakHashMap<>()));
-
-    @RealObject private LooperExecutor realExecutor;
+    @RealObject private LooperExecutor mRealExecutor;
 
     @Implementation
-    protected void __constructor__(Looper looper) {
-        invokeConstructor(LooperExecutor.class, realExecutor, from(Looper.class, looper));
-        executors.add(realExecutor);
-    }
-
-    /**
-     * Re-initializes any executor which may have been reset when a test finished
-     */
-    public static void reinitializeStaticExecutors() {
-        for (LooperExecutor executor : new ArrayList<>(executors)) {
-            setField(executor, "mHandler",
-                    new Handler(createAndStartNewLooper(executor.getThread().getName())));
+    protected Handler getHandler() {
+        Handler handler = directlyOn(mRealExecutor, LooperExecutor.class, "getHandler");
+        Thread thread = handler.getLooper().getThread();
+        if (!thread.isAlive()) {
+            // Robolectric destroys all loopers at the end of every test. Since Launcher maintains
+            // some static threads, they need to be reinitialized in case they were destroyed.
+            setField(mRealExecutor, "mHandler",
+                    new Handler(createAndStartNewLooper(thread.getName())));
         }
+        return directlyOn(mRealExecutor, LooperExecutor.class, "getHandler");
     }
 }
diff --git a/robolectric_tests/src/com/android/launcher3/shadows/ShadowMainThreadInitializedObject.java b/robolectric_tests/src/com/android/launcher3/shadows/ShadowMainThreadInitializedObject.java
new file mode 100644
index 0000000..6e2ccf8
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/shadows/ShadowMainThreadInitializedObject.java
@@ -0,0 +1,61 @@
+/*
+ * Copyright (C) 2019 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.shadows;
+
+import static org.robolectric.shadow.api.Shadow.invokeConstructor;
+import static org.robolectric.util.ReflectionHelpers.ClassParameter.from;
+
+import com.android.launcher3.util.MainThreadInitializedObject;
+import com.android.launcher3.util.MainThreadInitializedObject.ObjectProvider;
+
+import org.robolectric.annotation.Implementation;
+import org.robolectric.annotation.Implements;
+import org.robolectric.annotation.RealObject;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Set;
+import java.util.WeakHashMap;
+
+/**
+ * Shadow for {@link MainThreadInitializedObject} to provide reset functionality for static sObjects
+ */
+@Implements(value = MainThreadInitializedObject.class, isInAndroidSdk = false)
+public class ShadowMainThreadInitializedObject {
+
+    // Keep reference to all created MainThreadInitializedObject so they can be cleared after test
+    private static Set<MainThreadInitializedObject> sObjects =
+            Collections.synchronizedSet(Collections.newSetFromMap(new WeakHashMap<>()));
+
+    @RealObject private MainThreadInitializedObject mRealObject;
+
+    @Implementation
+    protected void __constructor__(ObjectProvider provider) {
+        invokeConstructor(MainThreadInitializedObject.class, mRealObject,
+                from(ObjectProvider.class, provider));
+        sObjects.add(mRealObject);
+    }
+
+    /**
+     * Resets all the initialized sObjects to be null
+     */
+    public static void resetInitializedObjects() {
+        for (MainThreadInitializedObject object : new ArrayList<>(sObjects)) {
+            object.initializeForTesting(null);
+        }
+    }
+}
diff --git a/robolectric_tests/src/com/android/launcher3/shadows/ShadowTogglableFlag.java b/robolectric_tests/src/com/android/launcher3/shadows/ShadowTogglableFlag.java
new file mode 100644
index 0000000..3603dd8
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/shadows/ShadowTogglableFlag.java
@@ -0,0 +1,38 @@
+/*
+ * Copyright (C) 2019 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.shadows;
+
+import android.content.Context;
+
+import com.android.launcher3.uioverrides.TogglableFlag;
+import com.android.launcher3.util.LooperExecutor;
+
+import org.robolectric.annotation.Implementation;
+import org.robolectric.annotation.Implements;
+
+/**
+ * Shadow for {@link LooperExecutor} to provide reset functionality for static executors.
+ */
+@Implements(value = TogglableFlag.class, isInAndroidSdk = false)
+public class ShadowTogglableFlag {
+
+    /**
+     * Mock change listener as it uses internal system classes not available to robolectric
+     */
+    @Implementation
+    protected void addChangeListener(Context context, Runnable r) { }
+}
diff --git a/robolectric_tests/src/com/android/launcher3/util/GridOccupancyTest.java b/robolectric_tests/src/com/android/launcher3/util/GridOccupancyTest.java
index aa51ad2..e453e31 100644
--- a/robolectric_tests/src/com/android/launcher3/util/GridOccupancyTest.java
+++ b/robolectric_tests/src/com/android/launcher3/util/GridOccupancyTest.java
@@ -2,7 +2,6 @@
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.robolectric.RobolectricTestRunner;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -11,7 +10,7 @@
 /**
  * Unit tests for {@link GridOccupancy}
  */
-@RunWith(RobolectricTestRunner.class)
+@RunWith(LauncherRoboTestRunner.class)
 public class GridOccupancyTest {
 
     @Test
diff --git a/robolectric_tests/src/com/android/launcher3/util/IntArrayTest.java b/robolectric_tests/src/com/android/launcher3/util/IntArrayTest.java
index c08e198..5974ea5 100644
--- a/robolectric_tests/src/com/android/launcher3/util/IntArrayTest.java
+++ b/robolectric_tests/src/com/android/launcher3/util/IntArrayTest.java
@@ -19,12 +19,11 @@
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.robolectric.RobolectricTestRunner;
 
 /**
  * Robolectric unit tests for {@link IntArray}
  */
-@RunWith(RobolectricTestRunner.class)
+@RunWith(LauncherRoboTestRunner.class)
 public class IntArrayTest {
 
     @Test
diff --git a/robolectric_tests/src/com/android/launcher3/util/IntSetTest.java b/robolectric_tests/src/com/android/launcher3/util/IntSetTest.java
index 8513353..aedf71e 100644
--- a/robolectric_tests/src/com/android/launcher3/util/IntSetTest.java
+++ b/robolectric_tests/src/com/android/launcher3/util/IntSetTest.java
@@ -20,8 +20,6 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
-import org.robolectric.RobolectricTestRunner;
-
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
@@ -29,7 +27,7 @@
 /**
  * Robolectric unit tests for {@link IntSet}
  */
-@RunWith(RobolectricTestRunner.class)
+@RunWith(LauncherRoboTestRunner.class)
 public class IntSetTest {
 
     @Test
diff --git a/robolectric_tests/src/com/android/launcher3/util/LauncherModelHelper.java b/robolectric_tests/src/com/android/launcher3/util/LauncherModelHelper.java
new file mode 100644
index 0000000..1a03f9f
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/util/LauncherModelHelper.java
@@ -0,0 +1,288 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.util;
+
+import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
+
+import static org.mockito.Mockito.atLeast;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+
+import android.content.ComponentName;
+import android.content.ContentValues;
+import android.content.Context;
+import android.content.Intent;
+
+import com.android.launcher3.AppInfo;
+import com.android.launcher3.ItemInfo;
+import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherModel;
+import com.android.launcher3.LauncherModel.ModelUpdateTask;
+import com.android.launcher3.LauncherProvider;
+import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.model.AllAppsList;
+import com.android.launcher3.model.BgDataModel;
+
+import org.mockito.ArgumentCaptor;
+import org.robolectric.Robolectric;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.shadows.ShadowContentResolver;
+import org.robolectric.util.ReflectionHelpers;
+
+import java.io.BufferedReader;
+import java.io.InputStreamReader;
+import java.lang.reflect.Field;
+import java.util.HashMap;
+import java.util.List;
+import java.util.concurrent.Executor;
+import java.util.function.Function;
+
+/**
+ * Utility class to help manage Launcher Model and related objects for test.
+ */
+public class LauncherModelHelper {
+
+    public static final int DESKTOP = LauncherSettings.Favorites.CONTAINER_DESKTOP;
+    public static final int HOTSEAT = LauncherSettings.Favorites.CONTAINER_HOTSEAT;
+
+    public static final int APP_ICON = LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
+    public static final int SHORTCUT = LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT;
+    public static final int NO__ICON = -1;
+    public static final String TEST_PACKAGE = "com.android.launcher3.validpackage";
+
+    private final HashMap<Class, HashMap<String, Field>> mFieldCache = new HashMap<>();
+    public final TestLauncherProvider provider;
+
+    private BgDataModel mDataModel;
+    private AllAppsList mAllAppsList;
+
+    public LauncherModelHelper() {
+        provider = Robolectric.setupContentProvider(TestLauncherProvider.class);
+        ShadowContentResolver.registerProviderInternal(LauncherProvider.AUTHORITY, provider);
+    }
+
+    public LauncherModel getModel() {
+        return LauncherAppState.getInstance(RuntimeEnvironment.application).getModel();
+    }
+
+    public synchronized BgDataModel getBgDataModel() {
+        if (mDataModel == null) {
+            mDataModel = ReflectionHelpers.getField(getModel(), "mBgDataModel");
+        }
+        return mDataModel;
+    }
+
+    public synchronized AllAppsList getAllAppsList() {
+        if (mAllAppsList == null) {
+            mAllAppsList = ReflectionHelpers.getField(getModel(), "mBgAllAppsList");
+        }
+        return mAllAppsList;
+    }
+
+    /**
+     * Synchronously executes the task and returns all the UI callbacks posted.
+     */
+    public List<Runnable> executeTaskForTest(ModelUpdateTask task) throws Exception {
+        LauncherModel model = getModel();
+        if (!model.isModelLoaded()) {
+            ReflectionHelpers.setField(model, "mModelLoaded", true);
+        }
+        Executor mockExecutor = mock(Executor.class);
+        model.enqueueModelUpdateTask(new ModelUpdateTask() {
+            @Override
+            public void init(LauncherAppState app, LauncherModel model, BgDataModel dataModel,
+                    AllAppsList allAppsList, Executor uiExecutor) {
+                task.init(app, model, dataModel, allAppsList, mockExecutor);
+            }
+
+            @Override
+            public void run() {
+                task.run();
+            }
+        });
+        MODEL_EXECUTOR.submit(() -> null).get();
+
+        ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
+        verify(mockExecutor, atLeast(0)).execute(captor.capture());
+        return captor.getAllValues();
+    }
+
+    /**
+     * Synchronously executes a task on the model
+     */
+    public <T> T executeSimpleTask(Function<BgDataModel, T> task) throws Exception {
+        BgDataModel dataModel = getBgDataModel();
+        return MODEL_EXECUTOR.submit(() -> task.apply(dataModel)).get();
+    }
+
+    /**
+     * Initializes mock data for the test.
+     */
+    public void initializeData(String resourceName) throws Exception {
+        Context targetContext = RuntimeEnvironment.application;
+        BgDataModel bgDataModel = getBgDataModel();
+        AllAppsList allAppsList = getAllAppsList();
+
+        MODEL_EXECUTOR.submit(() -> {
+            try (BufferedReader reader = new BufferedReader(new InputStreamReader(
+                    this.getClass().getResourceAsStream(resourceName)))) {
+                String line;
+                HashMap<String, Class> classMap = new HashMap<>();
+                while ((line = reader.readLine()) != null) {
+                    line = line.trim();
+                    if (line.startsWith("#") || line.isEmpty()) {
+                        continue;
+                    }
+                    String[] commands = line.split(" ");
+                    switch (commands[0]) {
+                        case "classMap":
+                            classMap.put(commands[1], Class.forName(commands[2]));
+                            break;
+                        case "bgItem":
+                            bgDataModel.addItem(targetContext,
+                                    (ItemInfo) initItem(classMap.get(commands[1]), commands, 2),
+                                    false);
+                            break;
+                        case "allApps":
+                            allAppsList.add((AppInfo) initItem(AppInfo.class, commands, 1), null);
+                            break;
+                    }
+                }
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+        }).get();
+    }
+
+    private Object initItem(Class clazz, String[] fieldDef, int startIndex) throws Exception {
+        HashMap<String, Field> cache = mFieldCache.get(clazz);
+        if (cache == null) {
+            cache = new HashMap<>();
+            Class c = clazz;
+            while (c != null) {
+                for (Field f : c.getDeclaredFields()) {
+                    f.setAccessible(true);
+                    cache.put(f.getName(), f);
+                }
+                c = c.getSuperclass();
+            }
+            mFieldCache.put(clazz, cache);
+        }
+
+        Object item = clazz.newInstance();
+        for (int i = startIndex; i < fieldDef.length; i++) {
+            String[] fieldData = fieldDef[i].split("=", 2);
+            Field f = cache.get(fieldData[0]);
+            Class type = f.getType();
+            if (type == int.class || type == long.class) {
+                f.set(item, Integer.parseInt(fieldData[1]));
+            } else if (type == CharSequence.class || type == String.class) {
+                f.set(item, fieldData[1]);
+            } else if (type == Intent.class) {
+                if (!fieldData[1].startsWith("#Intent")) {
+                    fieldData[1] = "#Intent;" + fieldData[1] + ";end";
+                }
+                f.set(item, Intent.parseUri(fieldData[1], 0));
+            } else if (type == ComponentName.class) {
+                f.set(item, ComponentName.unflattenFromString(fieldData[1]));
+            } else {
+                throw new Exception("Added parsing logic for "
+                        + f.getName() + " of type " + f.getType());
+            }
+        }
+        return item;
+    }
+
+    /**
+     * Adds a dummy item in the DB.
+     * @param type {@link #APP_ICON} or {@link #SHORTCUT} or >= 2 for
+     *             folder (where the type represents the number of items in the folder).
+     */
+    public int addItem(int type, int screen, int container, int x, int y) {
+        Context context = RuntimeEnvironment.application;
+        int id = LauncherSettings.Settings.call(context.getContentResolver(),
+                LauncherSettings.Settings.METHOD_NEW_ITEM_ID)
+                .getInt(LauncherSettings.Settings.EXTRA_VALUE);
+
+        ContentValues values = new ContentValues();
+        values.put(LauncherSettings.Favorites._ID, id);
+        values.put(LauncherSettings.Favorites.CONTAINER, container);
+        values.put(LauncherSettings.Favorites.SCREEN, screen);
+        values.put(LauncherSettings.Favorites.CELLX, x);
+        values.put(LauncherSettings.Favorites.CELLY, y);
+        values.put(LauncherSettings.Favorites.SPANX, 1);
+        values.put(LauncherSettings.Favorites.SPANY, 1);
+
+        if (type == APP_ICON || type == SHORTCUT) {
+            values.put(LauncherSettings.Favorites.ITEM_TYPE, type);
+            values.put(LauncherSettings.Favorites.INTENT,
+                    new Intent(Intent.ACTION_MAIN).setPackage(TEST_PACKAGE).toUri(0));
+        } else {
+            values.put(LauncherSettings.Favorites.ITEM_TYPE,
+                    LauncherSettings.Favorites.ITEM_TYPE_FOLDER);
+            // Add folder items.
+            for (int i = 0; i < type; i++) {
+                addItem(APP_ICON, 0, id, 0, 0);
+            }
+        }
+
+        context.getContentResolver().insert(LauncherSettings.Favorites.CONTENT_URI, values);
+        return id;
+    }
+
+    public int[][][] createGrid(int[][][] typeArray) {
+        return createGrid(typeArray, 1);
+    }
+
+    /**
+     * Initializes the DB with dummy elements to represent the provided grid structure.
+     * @param typeArray A 3d array of item types. {@see #addItem(int, long, long, int, int)} for
+     *                  type definitions. The first dimension represents the screens and the next
+     *                  two represent the workspace grid.
+     * @param startScreen First screen id from where the icons will be added.
+     * @return the same grid representation where each entry is the corresponding item id.
+     */
+    public int[][][] createGrid(int[][][] typeArray, int startScreen) {
+        Context context = RuntimeEnvironment.application;
+        LauncherSettings.Settings.call(context.getContentResolver(),
+                LauncherSettings.Settings.METHOD_CREATE_EMPTY_DB);
+        int[][][] ids = new int[typeArray.length][][];
+
+        for (int i = 0; i < typeArray.length; i++) {
+            // Add screen to DB
+            int screenId = startScreen + i;
+
+            // Keep the screen id counter up to date
+            LauncherSettings.Settings.call(context.getContentResolver(),
+                    LauncherSettings.Settings.METHOD_NEW_SCREEN_ID);
+
+            ids[i] = new int[typeArray[i].length][];
+            for (int y = 0; y < typeArray[i].length; y++) {
+                ids[i][y] = new int[typeArray[i][y].length];
+                for (int x = 0; x < typeArray[i][y].length; x++) {
+                    if (typeArray[i][y][x] < 0) {
+                        // Empty cell
+                        ids[i][y][x] = -1;
+                    } else {
+                        ids[i][y][x] = addItem(typeArray[i][y][x], screenId, DESKTOP, x, y);
+                    }
+                }
+            }
+        }
+
+        return ids;
+    }
+}
diff --git a/robolectric_tests/src/com/android/launcher3/util/LauncherRoboTestRunner.java b/robolectric_tests/src/com/android/launcher3/util/LauncherRoboTestRunner.java
new file mode 100644
index 0000000..5c6b486
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/util/LauncherRoboTestRunner.java
@@ -0,0 +1,90 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.util;
+
+import static org.mockito.Mockito.mock;
+
+import com.android.launcher3.shadows.LShadowAppWidgetManager;
+import com.android.launcher3.shadows.LShadowBitmap;
+import com.android.launcher3.shadows.LShadowLauncherApps;
+import com.android.launcher3.shadows.LShadowUserManager;
+import com.android.launcher3.shadows.ShadowLooperExecutor;
+import com.android.launcher3.shadows.ShadowMainThreadInitializedObject;
+import com.android.launcher3.shadows.ShadowTogglableFlag;
+import com.android.launcher3.uioverrides.plugins.PluginManagerWrapper;
+
+import org.junit.runners.model.InitializationError;
+import org.robolectric.DefaultTestLifecycle;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.TestLifecycle;
+import org.robolectric.annotation.Config;
+import org.robolectric.shadows.ShadowLog;
+
+import java.lang.reflect.Method;
+
+import javax.annotation.Nonnull;
+
+/**
+ * Test runner with Launcher specific configurations
+ */
+public class LauncherRoboTestRunner extends RobolectricTestRunner {
+
+    private static final Class<?>[] SHADOWS = new Class<?>[] {
+            LShadowAppWidgetManager.class,
+            LShadowUserManager.class,
+            LShadowLauncherApps.class,
+            LShadowBitmap.class,
+
+            ShadowLooperExecutor.class,
+            ShadowMainThreadInitializedObject.class,
+            ShadowTogglableFlag.class,
+    };
+
+    public LauncherRoboTestRunner(Class<?> testClass) throws InitializationError {
+        super(testClass);
+    }
+
+    @Override
+    protected Config buildGlobalConfig() {
+        return new Config.Builder().setShadows(SHADOWS).build();
+    }
+
+    @Nonnull
+    @Override
+    protected Class<? extends TestLifecycle> getTestLifecycleClass() {
+        return LauncherTestLifecycle.class;
+    }
+
+    public static class LauncherTestLifecycle extends DefaultTestLifecycle {
+
+        @Override
+        public void beforeTest(Method method) {
+            super.beforeTest(method);
+            ShadowLog.stream = System.out;
+
+            // Disable plugins
+            PluginManagerWrapper.INSTANCE.initializeForTesting(mock(PluginManagerWrapper.class));
+        }
+
+        @Override
+        public void afterTest(Method method) {
+            super.afterTest(method);
+
+            ShadowLog.stream = null;
+            ShadowMainThreadInitializedObject.resetInitializedObjects();
+        }
+    }
+}
diff --git a/tests/src/com/android/launcher3/widget/WidgetsListAdapterTest.java b/robolectric_tests/src/com/android/launcher3/widget/WidgetsListAdapterTest.java
similarity index 83%
rename from tests/src/com/android/launcher3/widget/WidgetsListAdapterTest.java
rename to robolectric_tests/src/com/android/launcher3/widget/WidgetsListAdapterTest.java
index 57b0b09..daae818 100644
--- a/tests/src/com/android/launcher3/widget/WidgetsListAdapterTest.java
+++ b/robolectric_tests/src/com/android/launcher3/widget/WidgetsListAdapterTest.java
@@ -19,16 +19,15 @@
 import static org.mockito.Matchers.isNull;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.robolectric.Shadows.shadowOf;
 
 import android.appwidget.AppWidgetProviderInfo;
+import android.content.ComponentName;
 import android.content.Context;
 import android.graphics.Bitmap;
 import android.view.LayoutInflater;
 
 import androidx.recyclerview.widget.RecyclerView;
-import androidx.test.InstrumentationRegistry;
-import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppWidgetProviderInfo;
@@ -37,19 +36,21 @@
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.PackageItemInfo;
 import com.android.launcher3.model.WidgetItem;
-import com.android.launcher3.util.MultiHashMap;
+import com.android.launcher3.util.LauncherRoboTestRunner;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.shadows.ShadowPackageManager;
+import org.robolectric.util.ReflectionHelpers;
 
 import java.util.ArrayList;
-import java.util.Map;
+import java.util.Collections;
 
-@SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(LauncherRoboTestRunner.class)
 public class WidgetsListAdapterTest {
 
     @Mock private LayoutInflater mMockLayoutInflater;
@@ -64,7 +65,7 @@
     @Before
     public void setup() {
         MockitoAnnotations.initMocks(this);
-        mContext = InstrumentationRegistry.getTargetContext();
+        mContext = RuntimeEnvironment.application;
         mTestProfile = new InvariantDeviceProfile();
         mTestProfile.numRows = 5;
         mTestProfile.numColumns = 5;
@@ -121,15 +122,19 @@
     /**
      * Helper method to generate the sample widget model map that can be used for the tests
      * @param num the number of WidgetItem the map should contain
-     * @return
      */
     private ArrayList<WidgetListRowEntry> generateSampleMap(int num) {
         ArrayList<WidgetListRowEntry> result = new ArrayList<>();
         if (num <= 0) return result;
+        ShadowPackageManager spm = shadowOf(mContext.getPackageManager());
 
-        MultiHashMap<PackageItemInfo, WidgetItem> newMap = new MultiHashMap();
-        WidgetManagerHelper widgetManager = new WidgetManagerHelper(mContext);
-        for (AppWidgetProviderInfo widgetInfo : widgetManager.getAllProviders(null)) {
+        for (int i = 0; i < num; i++) {
+            ComponentName cn = new ComponentName("com.dummy.apk" + i, "DummyWidet");
+
+            AppWidgetProviderInfo widgetInfo = new AppWidgetProviderInfo();
+            widgetInfo.provider = cn;
+            ReflectionHelpers.setField(widgetInfo, "providerInfo", spm.addReceiverIfNotPresent(cn));
+
             WidgetItem wi = new WidgetItem(LauncherAppWidgetProviderInfo
                     .fromProviderInfo(mContext, widgetInfo), mTestProfile, mIconCache);
 
@@ -137,13 +142,8 @@
             pInfo.title = pInfo.packageName;
             pInfo.user = wi.user;
             pInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
-            newMap.addToList(pInfo, wi);
-            if (newMap.size() == num) {
-                break;
-            }
-        }
-        for (Map.Entry<PackageItemInfo, ArrayList<WidgetItem>> entry : newMap.entrySet()) {
-            result.add(new WidgetListRowEntry(entry.getKey(), entry.getValue()));
+
+            result.add(new WidgetListRowEntry(pInfo, new ArrayList<>(Collections.singleton(wi))));
         }
 
         return result;
diff --git a/src/com/android/launcher3/LauncherModel.java b/src/com/android/launcher3/LauncherModel.java
index 0bdf8fd..e005320 100644
--- a/src/com/android/launcher3/LauncherModel.java
+++ b/src/com/android/launcher3/LauncherModel.java
@@ -108,14 +108,14 @@
      * All the static data should be accessed on the background thread, A lock should be acquired
      * on this object when accessing any data from this model.
      */
-    static final BgDataModel sBgDataModel = new BgDataModel();
+    private final BgDataModel mBgDataModel = new BgDataModel();
 
     // Runnable to check if the shortcuts permission has changed.
     private final Runnable mShortcutPermissionCheckRunnable = new Runnable() {
         @Override
         public void run() {
             if (mModelLoaded && hasShortcutsPermission(mApp.getContext())
-                    != sBgDataModel.hasShortcutHostPermission) {
+                    != mBgDataModel.hasShortcutHostPermission) {
                 forceReload();
             }
         }
@@ -138,7 +138,7 @@
     }
 
     public ModelWriter getWriter(boolean hasVerticalHotseat, boolean verifyChanges) {
-        return new ModelWriter(mApp.getContext(), this, sBgDataModel,
+        return new ModelWriter(mApp.getContext(), this, mBgDataModel,
                 hasVerticalHotseat, verifyChanges);
     }
 
@@ -303,7 +303,7 @@
 
                 // If there is already one running, tell it to stop.
                 stopLoader();
-                LoaderResults loaderResults = new LoaderResults(mApp, sBgDataModel,
+                LoaderResults loaderResults = new LoaderResults(mApp, mBgDataModel,
                         mBgAllAppsList, synchronousBindPage, mCallbacks);
                 if (mModelLoaded && !mIsLoaderTaskRunning) {
                     // Divide the set of loaded items into those that we are binding synchronously,
@@ -339,7 +339,7 @@
     public void startLoaderForResults(LoaderResults results) {
         synchronized (mLock) {
             stopLoader();
-            mLoaderTask = new LoaderTask(mApp, mBgAllAppsList, sBgDataModel, results);
+            mLoaderTask = new LoaderTask(mApp, mBgAllAppsList, mBgDataModel, results);
 
             // Always post the loader task, instead of running directly (even on same thread) so
             // that we exit any nested synchronized blocks
@@ -487,7 +487,7 @@
     }
 
     public void enqueueModelUpdateTask(ModelUpdateTask task) {
-        task.init(mApp, this, sBgDataModel, mBgAllAppsList, MAIN_EXECUTOR);
+        task.init(mApp, this, mBgDataModel, mBgAllAppsList, MAIN_EXECUTOR);
         MODEL_EXECUTOR.execute(task);
     }
 
@@ -558,7 +558,7 @@
                         + " componentName=" + info.componentName.getPackageName());
             }
         }
-        sBgDataModel.dump(prefix, fd, writer, args);
+        mBgDataModel.dump(prefix, fd, writer, args);
     }
 
     public Callbacks getCallback() {
diff --git a/src/com/android/launcher3/util/Executors.java b/src/com/android/launcher3/util/Executors.java
index 4d5ee49..0a32734 100644
--- a/src/com/android/launcher3/util/Executors.java
+++ b/src/com/android/launcher3/util/Executors.java
@@ -19,7 +19,6 @@
 import android.os.Looper;
 import android.os.Process;
 
-import java.util.concurrent.Executor;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
@@ -36,9 +35,9 @@
     private static final int KEEP_ALIVE = 1;
 
     /**
-     * An {@link Executor} to be used with async task with no limit on the queue size.
+     * An {@link ThreadPoolExecutor} to be used with async task with no limit on the queue size.
      */
-    public static final Executor THREAD_POOL_EXECUTOR = new ThreadPoolExecutor(
+    public static final ThreadPoolExecutor THREAD_POOL_EXECUTOR = new ThreadPoolExecutor(
             CORE_POOL_SIZE, MAXIMUM_POOL_SIZE, KEEP_ALIVE,
             TimeUnit.SECONDS, new LinkedBlockingQueue<>());
 
diff --git a/src/com/android/launcher3/util/LooperExecutor.java b/src/com/android/launcher3/util/LooperExecutor.java
index 8ac600f..3a8a13c 100644
--- a/src/com/android/launcher3/util/LooperExecutor.java
+++ b/src/com/android/launcher3/util/LooperExecutor.java
@@ -41,10 +41,10 @@
 
     @Override
     public void execute(Runnable runnable) {
-        if (mHandler.getLooper() == Looper.myLooper()) {
+        if (getHandler().getLooper() == Looper.myLooper()) {
             runnable.run();
         } else {
-            mHandler.post(runnable);
+            getHandler().post(runnable);
         }
     }
 
@@ -52,7 +52,7 @@
      * Same as execute, but never runs the action inline.
      */
     public void post(Runnable runnable) {
-        mHandler.post(runnable);
+        getHandler().post(runnable);
     }
 
     /**
@@ -96,14 +96,14 @@
      * Returns the thread for this executor
      */
     public Thread getThread() {
-        return mHandler.getLooper().getThread();
+        return getHandler().getLooper().getThread();
     }
 
     /**
      * Returns the looper for this executor
      */
     public Looper getLooper() {
-        return mHandler.getLooper();
+        return getHandler().getLooper();
     }
 
     /**