Revert "Revert "Migrating all model tests to Instrumentation tests""

This reverts commit 7a4a30d86d471e6c45adc2a9907efb27e9b1799b.

Test: Presubmit
Reason for revert: Fixing original bug

Bug: 196825541
Change-Id: Id4b1eb24a89564d264266d305aebea52917dfcd9

14 files changed, 2732 insertions, 0 deletions

diff --git a/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java b/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
new file mode 100644
index 0000000000..16f024e421
--- /dev/null
+++ b/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
@@ -0,0 +1,221 @@
+package com.android.launcher3.model;
+
+import static org.junit.Assert.assertEquals;
+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 androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.LauncherSettings.Favorites;
+import com.android.launcher3.model.BgDataModel.Callbacks;
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.model.data.WorkspaceItemInfo;
+import com.android.launcher3.util.ContentWriter;
+import com.android.launcher3.util.Executors;
+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 org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Tests for {@link AddWorkspaceItemsTask}
+ */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class AddWorkspaceItemsTaskTest {
+
+    private final ComponentName mComponent1 = new ComponentName("a", "b");
+    private final ComponentName mComponent2 = new ComponentName("b", "b");
+
+    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 setup() {
+        mModelHelper = new LauncherModelHelper();
+        mTargetContext = mModelHelper.sandboxContext;
+        mIdp = InvariantDeviceProfile.INSTANCE.get(mTargetContext);
+        mIdp.numColumns = mIdp.numRows = 5;
+        mAppState = LauncherAppState.getInstance(mTargetContext);
+
+        mExistingScreens = new IntArray();
+        mScreenOccupancy = new IntSparseArrayMap<>();
+        mNewScreens = new IntArray();
+    }
+
+    @After
+    public void tearDown() {
+        mModelHelper.destroy();
+    }
+
+    private AddWorkspaceItemsTask newTask(ItemInfo... items) {
+        List<Pair<ItemInfo, Object>> list = new ArrayList<>();
+        for (ItemInfo item : items) {
+            list.add(Pair.create(item, null));
+        }
+        return new AddWorkspaceItemsTask(list);
+    }
+
+    @Test
+    public void testFindSpaceForItem_prefers_second() throws Exception {
+        mIdp.isSplitDisplay = false;
+
+        // First screen has only one hole of size 1
+        int nextId = setupWorkspaceWithHoles(1, 1, new Rect(2, 2, 3, 3));
+
+        // 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(
+                mAppState, mModelHelper.getBgDataModel(), mExistingScreens, mNewScreens, 1, 1);
+        assertEquals(1, spaceFound[0]);
+        assertTrue(mScreenOccupancy.get(spaceFound[0])
+                .isRegionVacant(spaceFound[1], spaceFound[2], 1, 1));
+
+        // Find a larger space
+        spaceFound = newTask().findSpaceForItem(
+                mAppState, mModelHelper.getBgDataModel(), mExistingScreens, mNewScreens, 2, 3);
+        assertEquals(2, spaceFound[0]);
+        assertTrue(mScreenOccupancy.get(spaceFound[0])
+                .isRegionVacant(spaceFound[1], spaceFound[2], 2, 3));
+    }
+
+    @Test
+    public void testFindSpaceForItem_prefers_third_on_split_display() throws Exception {
+        mIdp.isSplitDisplay = true;
+        // First screen has only one hole of size 1
+        int nextId = setupWorkspaceWithHoles(1, 1, new Rect(2, 2, 3, 3));
+
+        // 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(
+                mAppState, mModelHelper.getBgDataModel(), mExistingScreens, mNewScreens, 1, 1);
+        // For split display, it picks the next screen, even if there is enough space
+        // on previous screen
+        assertEquals(2, spaceFound[0]);
+        assertTrue(mScreenOccupancy.get(spaceFound[0])
+                .isRegionVacant(spaceFound[1], spaceFound[2], 1, 1));
+    }
+
+    @Test
+    public void testFindSpaceForItem_adds_new_screen() throws Exception {
+        // 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 = mExistingScreens.clone();
+        int[] spaceFound = newTask().findSpaceForItem(
+                mAppState, mModelHelper.getBgDataModel(), mExistingScreens, mNewScreens, 3, 3);
+        assertFalse(oldScreens.contains(spaceFound[0]));
+        assertTrue(mNewScreens.contains(spaceFound[0]));
+    }
+
+    @Test
+    public void testAddItem_existing_item_ignored() throws Exception {
+        WorkspaceItemInfo info = new WorkspaceItemInfo();
+        info.intent = new Intent().setComponent(mComponent1);
+
+        // Setup a screen with a hole
+        setupWorkspaceWithHoles(1, 1, new Rect(2, 2, 3, 3));
+
+        // Nothing was added
+        assertTrue(mModelHelper.executeTaskForTest(newTask(info)).isEmpty());
+    }
+
+    @Test
+    public void testAddItem_some_items_added() throws Exception {
+        Callbacks callbacks = mock(Callbacks.class);
+        Executors.MAIN_EXECUTOR.submit(() -> mModelHelper.getModel().addCallbacks(callbacks)).get();
+
+        WorkspaceItemInfo info = new WorkspaceItemInfo();
+        info.intent = new Intent().setComponent(mComponent1);
+
+        WorkspaceItemInfo info2 = new WorkspaceItemInfo();
+        info2.intent = new Intent().setComponent(mComponent2);
+
+        // Setup a screen with a hole
+        setupWorkspaceWithHoles(1, 1, new Rect(2, 2, 3, 3));
+
+        mModelHelper.executeTaskForTest(newTask(info, info2)).get(0).run();
+        ArgumentCaptor<ArrayList> notAnimated = ArgumentCaptor.forClass(ArrayList.class);
+        ArgumentCaptor<ArrayList> animated = ArgumentCaptor.forClass(ArrayList.class);
+
+        // only info2 should be added because info was already added to the workspace
+        // in setupWorkspaceWithHoles()
+        verify(callbacks).bindAppsAdded(any(IntArray.class), notAnimated.capture(),
+                animated.capture());
+        assertTrue(notAnimated.getValue().isEmpty());
+
+        assertEquals(1, animated.getValue().size());
+        assertTrue(animated.getValue().contains(info2));
+    }
+
+    private int setupWorkspaceWithHoles(int startId, int screenId, Rect... holes) throws Exception {
+        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);
+        }
+
+        mExistingScreens.add(screenId);
+        mScreenOccupancy.append(screenId, occupancy);
+
+        for (int x = 0; x < mIdp.numColumns; x++) {
+            for (int y = 0; y < mIdp.numRows; y++) {
+                if (!occupancy.cells[x][y]) {
+                    continue;
+                }
+
+                WorkspaceItemInfo info = new WorkspaceItemInfo();
+                info.intent = new Intent().setComponent(mComponent1);
+                info.id = startId++;
+                info.screenId = screenId;
+                info.cellX = x;
+                info.cellY = y;
+                info.container = LauncherSettings.Favorites.CONTAINER_DESKTOP;
+                bgDataModel.addItem(mTargetContext, info, false);
+
+                ContentWriter writer = new ContentWriter(mTargetContext);
+                info.writeToValues(writer);
+                writer.put(Favorites._ID, info.id);
+                mTargetContext.getContentResolver().insert(Favorites.CONTENT_URI,
+                        writer.getValues(mTargetContext));
+            }
+        }
+        return startId;
+    }
+}
diff --git a/tests/src/com/android/launcher3/model/BackupRestoreTest.java b/tests/src/com/android/launcher3/model/BackupRestoreTest.java
new file mode 100644
index 0000000000..41914de1ab
--- /dev/null
+++ b/tests/src/com/android/launcher3/model/BackupRestoreTest.java
@@ -0,0 +1,211 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.model;
+
+import static android.content.pm.PackageManager.INSTALL_REASON_DEVICE_RESTORE;
+import static android.os.Process.myUserHandle;
+
+import static com.android.launcher3.LauncherSettings.Favorites.BACKUP_TABLE_NAME;
+import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
+import static com.android.launcher3.LauncherSettings.Favorites.addTableToDb;
+import static com.android.launcher3.provider.LauncherDbUtils.dropTable;
+import static com.android.launcher3.provider.LauncherDbUtils.tableExists;
+import static com.android.launcher3.util.LauncherModelHelper.APP_ICON;
+import static com.android.launcher3.util.LauncherModelHelper.NO__ICON;
+import static com.android.launcher3.util.LauncherModelHelper.SHORTCUT;
+import static com.android.launcher3.util.ReflectionHelpers.getField;
+import static com.android.launcher3.util.ReflectionHelpers.setField;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
+
+import android.app.backup.BackupManager;
+import android.content.pm.PackageInstaller;
+import android.database.Cursor;
+import android.database.sqlite.SQLiteDatabase;
+import android.os.UserHandle;
+import android.util.ArrayMap;
+import android.util.LongSparseArray;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.pm.UserCache;
+import com.android.launcher3.provider.RestoreDbTask;
+import com.android.launcher3.util.LauncherModelHelper;
+import com.android.launcher3.util.SafeCloseable;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/**
+ * Tests to verify backup and restore flow.
+ */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class BackupRestoreTest {
+
+    private static final int PER_USER_RANGE = 200000;
+
+
+    private long mCurrentMyProfileId;
+    private long mOldMyProfileId;
+
+    private long mCurrentWorkProfileId;
+    private long mOldWorkProfileId;
+
+    private BackupManager mBackupManager;
+    private LauncherModelHelper mModelHelper;
+    private SQLiteDatabase mDb;
+    private InvariantDeviceProfile mIdp;
+
+    private UserHandle mWorkUserHandle;
+
+    private SafeCloseable mUserChangeListener;
+
+    @Before
+    public void setUp() {
+        mModelHelper = new LauncherModelHelper();
+
+        mCurrentMyProfileId = mModelHelper.defaultProfileId;
+        mOldMyProfileId = mCurrentMyProfileId + 1;
+        mCurrentWorkProfileId = mOldMyProfileId + 1;
+        mOldWorkProfileId = mCurrentWorkProfileId + 1;
+
+        mWorkUserHandle = UserHandle.getUserHandleForUid(PER_USER_RANGE);
+        mUserChangeListener = UserCache.INSTANCE.get(mModelHelper.sandboxContext)
+                .addUserChangeListener(() -> { });
+
+        setupUserManager();
+        setupBackupManager();
+        RestoreDbTask.setPending(mModelHelper.sandboxContext);
+        mDb = mModelHelper.provider.getDb();
+        mIdp = InvariantDeviceProfile.INSTANCE.get(mModelHelper.sandboxContext);
+
+    }
+
+    @After
+    public void tearDown() {
+        mUserChangeListener.close();
+        mModelHelper.destroy();
+    }
+
+    private void setupUserManager() {
+        UserCache cache = UserCache.INSTANCE.get(mModelHelper.sandboxContext);
+        synchronized (cache) {
+            LongSparseArray<UserHandle> users = getField(cache, "mUsers");
+            users.clear();
+            users.put(mCurrentMyProfileId, myUserHandle());
+            users.put(mCurrentWorkProfileId, mWorkUserHandle);
+
+            ArrayMap<UserHandle, Long> userMap = getField(cache, "mUserToSerialMap");
+            userMap.clear();
+            userMap.put(myUserHandle(), mCurrentMyProfileId);
+            userMap.put(mWorkUserHandle, mCurrentWorkProfileId);
+        }
+    }
+
+    private void setupBackupManager() {
+        mBackupManager = spy(new BackupManager(mModelHelper.sandboxContext));
+        doReturn(myUserHandle()).when(mBackupManager)
+                .getUserForAncestralSerialNumber(eq(mOldMyProfileId));
+        doReturn(mWorkUserHandle).when(mBackupManager)
+                .getUserForAncestralSerialNumber(eq(mOldWorkProfileId));
+    }
+
+    @Test
+    public void testOnCreateDbIfNotExists_CreatesBackup() {
+        assertTrue(tableExists(mDb, BACKUP_TABLE_NAME));
+    }
+
+    @Test
+    public void testOnRestoreSessionWithValidCondition_PerformsRestore() throws Exception {
+        setupBackup();
+        verifyTableIsFilled(BACKUP_TABLE_NAME, false);
+        verifyTableIsEmpty(TABLE_NAME);
+        createRestoreSession();
+        verifyTableIsFilled(TABLE_NAME, true);
+    }
+
+    private void setupBackup() {
+        createTableUsingOldProfileId();
+        // setup grid for main user on first screen
+        mModelHelper.createGrid(new int[][][]{{
+                { APP_ICON, APP_ICON, SHORTCUT, SHORTCUT},
+                { SHORTCUT, SHORTCUT, NO__ICON, NO__ICON},
+                { NO__ICON, NO__ICON, SHORTCUT, SHORTCUT},
+                { APP_ICON, SHORTCUT, SHORTCUT, APP_ICON},
+            }}, 1, mOldMyProfileId);
+        // setup grid for work profile on second screen
+        mModelHelper.createGrid(new int[][][]{{
+                { NO__ICON, APP_ICON, SHORTCUT, SHORTCUT},
+                { SHORTCUT, SHORTCUT, NO__ICON, NO__ICON},
+                { NO__ICON, NO__ICON, SHORTCUT, SHORTCUT},
+                { APP_ICON, SHORTCUT, SHORTCUT, NO__ICON},
+            }}, 2, mOldWorkProfileId);
+        // simulates the creation of backup upon restore
+        new GridBackupTable(mModelHelper.sandboxContext, mDb, mIdp.numDatabaseHotseatIcons,
+                mIdp.numColumns, mIdp.numRows).doBackup(
+                mOldMyProfileId, GridBackupTable.OPTION_REQUIRES_SANITIZATION);
+        // reset favorites table
+        createTableUsingOldProfileId();
+    }
+
+    private void verifyTableIsEmpty(String tableName) {
+        assertEquals(0, getCount(mDb, "SELECT * FROM " + tableName));
+    }
+
+    private void verifyTableIsFilled(String tableName, boolean sanitized) {
+        assertEquals(sanitized ? 12 : 13, getCount(mDb,
+                "SELECT * FROM " + tableName + " WHERE profileId = "
+                        + (sanitized ? mCurrentMyProfileId : mOldMyProfileId)));
+        assertEquals(10, getCount(mDb, "SELECT * FROM " + tableName + " WHERE profileId = "
+                + (sanitized ? mCurrentWorkProfileId : mOldWorkProfileId)));
+    }
+
+    private void createTableUsingOldProfileId() {
+        // simulates the creation of favorites table on old device
+        dropTable(mDb, TABLE_NAME);
+        addTableToDb(mDb, mOldMyProfileId, false);
+    }
+
+    private void createRestoreSession() throws Exception {
+        final PackageInstaller.SessionParams params = new PackageInstaller.SessionParams(
+                PackageInstaller.SessionParams.MODE_FULL_INSTALL);
+        final PackageInstaller installer = mModelHelper.sandboxContext.getPackageManager()
+                .getPackageInstaller();
+        final int sessionId = installer.createSession(params);
+        final PackageInstaller.SessionInfo info = installer.getSessionInfo(sessionId);
+        setField(info, "installReason", INSTALL_REASON_DEVICE_RESTORE);
+        // TODO: (b/148410677) we should verify the following call instead
+        //  InstallSessionHelper.INSTANCE.get(getContext()).restoreDbIfApplicable(info);
+        RestoreDbTask.restoreIfPossible(mModelHelper.sandboxContext,
+                mModelHelper.provider.getHelper(), mBackupManager);
+    }
+
+    private static int getCount(SQLiteDatabase db, String sql) {
+        try (Cursor c = db.rawQuery(sql, null)) {
+            return c.getCount();
+        }
+    }
+}
diff --git a/tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java b/tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
new file mode 100644
index 0000000000..dba0a4063f
--- /dev/null
+++ b/tests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
@@ -0,0 +1,151 @@
+package com.android.launcher3.model;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+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 androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.icons.BitmapInfo;
+import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.CachingLogic;
+import com.android.launcher3.model.data.AppInfo;
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.model.data.WorkspaceItemInfo;
+import com.android.launcher3.util.LauncherModelHelper;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.Arrays;
+import java.util.HashSet;
+
+/**
+ * Tests for {@link CacheDataUpdatedTask}
+ */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class CacheDataUpdatedTaskTest {
+
+    private static final String NEW_LABEL_PREFIX = "new-label-";
+
+    private LauncherModelHelper mModelHelper;
+
+    @Before
+    public void setup() throws Exception {
+        mModelHelper = new LauncherModelHelper();
+        mModelHelper.initializeData("cache_data_updated_task_data");
+
+        // Add placeholder entries in the cache to simulate update
+        Context context = mModelHelper.sandboxContext;
+        IconCache iconCache = LauncherAppState.getInstance(context).getIconCache();
+        CachingLogic<ItemInfo> placeholderLogic = 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, placeholderLogic, new PackageInfo(),
+                    um.getSerialNumberForUser(info.user), true);
+        }
+    }
+
+    @After
+    public void tearDown() {
+        mModelHelper.destroy();
+    }
+
+    private CacheDataUpdatedTask newTask(int op, String... 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 : mModelHelper.getAllAppsList().data) {
+            info.bitmap = BitmapInfo.LOW_RES_INFO;
+        }
+
+        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(mModelHelper.getAllAppsList().data.isEmpty());
+        for (AppInfo info : mModelHelper.getAllAppsList().data) {
+            if (info.componentName.getPackageName().equals("app1")) {
+                assertFalse(info.bitmap.isNullOrLowRes());
+            } else {
+                assertTrue(info.bitmap.isNullOrLowRes());
+            }
+        }
+    }
+
+    @Test
+    public void testSessionUpdate_ignores_normal_apps() throws Exception {
+        mModelHelper.executeTaskForTest(newTask(CacheDataUpdatedTask.OP_SESSION_UPDATE, "app1"));
+
+        // app1 has no restored shortcuts. Verify that nothing was updated.
+        verifyUpdate();
+    }
+
+    @Test
+    public void testSessionUpdate_updates_pending_apps() throws Exception {
+        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
+        verifyUpdate(5, 6);
+    }
+
+    private void verifyUpdate(Integer... idsUpdated) {
+        HashSet<Integer> updates = new HashSet<>(Arrays.asList(idsUpdated));
+        for (ItemInfo info : mModelHelper.getBgDataModel().itemsIdMap) {
+            if (updates.contains(info.id)) {
+                assertEquals(NEW_LABEL_PREFIX + info.id, info.title);
+                assertFalse(((WorkspaceItemInfo) info).bitmap.isNullOrLowRes());
+            } else {
+                assertNotSame(NEW_LABEL_PREFIX + info.id, info.title);
+                assertTrue(((WorkspaceItemInfo) info).bitmap.isNullOrLowRes());
+            }
+        }
+    }
+}
diff --git a/tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java b/tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
new file mode 100644
index 0000000000..d849c8fff9
--- /dev/null
+++ b/tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
@@ -0,0 +1,215 @@
+/*
+ * Copyright (C) 2017 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 androidx.test.InstrumentationRegistry.getContext;
+
+import static junit.framework.Assert.assertEquals;
+import static junit.framework.Assert.assertFalse;
+import static junit.framework.Assert.assertNotSame;
+import static junit.framework.Assert.assertTrue;
+
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.when;
+
+import android.content.ContentValues;
+import android.content.Context;
+import android.content.res.Resources;
+import android.database.Cursor;
+import android.database.sqlite.SQLiteDatabase;
+import android.database.sqlite.SQLiteOpenHelper;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+import androidx.test.platform.app.InstrumentationRegistry;
+
+import com.android.launcher3.LauncherProvider;
+import com.android.launcher3.LauncherProvider.DatabaseHelper;
+import com.android.launcher3.LauncherSettings.Favorites;
+import com.android.launcher3.R;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.io.File;
+
+/**
+ * Tests for {@link DbDowngradeHelper}
+ */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class DbDowngradeHelperTest {
+
+    private static final String SCHEMA_FILE = "test_schema.json";
+    private static final String DB_FILE = "test.db";
+
+    private Context mContext;
+    private File mSchemaFile;
+    private File mDbFile;
+
+    @Before
+    public void setup() {
+        mContext = InstrumentationRegistry.getInstrumentation().getTargetContext();
+        mSchemaFile = mContext.getFileStreamPath(SCHEMA_FILE);
+        mDbFile = mContext.getDatabasePath(DB_FILE);
+    }
+
+    @Test
+    public void testDowngradeSchemaMatchesVersion() throws Exception {
+        mSchemaFile.delete();
+        assertFalse(mSchemaFile.exists());
+        DbDowngradeHelper.updateSchemaFile(mSchemaFile, 0, mContext);
+        assertEquals(LauncherProvider.SCHEMA_VERSION, DbDowngradeHelper.parse(mSchemaFile).version);
+    }
+
+    @Test
+    public void testUpdateSchemaFile() throws Exception {
+        // Setup mock resources
+        Resources res = spy(mContext.getResources());
+        Resources myRes = getContext().getResources();
+        doAnswer(i -> myRes.openRawResource(
+                myRes.getIdentifier("db_schema_v10", "raw", getContext().getPackageName())))
+                .when(res).openRawResource(R.raw.downgrade_schema);
+        Context context = spy(mContext);
+        when(context.getResources()).thenReturn(res);
+
+        mSchemaFile.delete();
+        assertFalse(mSchemaFile.exists());
+
+        DbDowngradeHelper.updateSchemaFile(mSchemaFile, 10, context);
+        assertTrue(mSchemaFile.exists());
+        assertEquals(10, DbDowngradeHelper.parse(mSchemaFile).version);
+
+        // Schema is updated on version upgrade
+        assertTrue(mSchemaFile.setLastModified(0));
+        DbDowngradeHelper.updateSchemaFile(mSchemaFile, 11, context);
+        assertNotSame(0, mSchemaFile.lastModified());
+
+        // Schema is not updated when version is same
+        assertTrue(mSchemaFile.setLastModified(0));
+        DbDowngradeHelper.updateSchemaFile(mSchemaFile, 10, context);
+        assertEquals(0, mSchemaFile.lastModified());
+
+        // Schema is not updated on version downgrade
+        DbDowngradeHelper.updateSchemaFile(mSchemaFile, 3, context);
+        assertEquals(0, mSchemaFile.lastModified());
+    }
+
+    @Test
+    public void testDowngrade_success_v24() throws Exception {
+        setupTestDb();
+
+        TestOpenHelper helper = new TestOpenHelper(24);
+        assertEquals(24, helper.getReadableDatabase().getVersion());
+        helper.close();
+    }
+
+    @Test
+    public void testDowngrade_success_v22() throws Exception {
+        setupTestDb();
+
+        SQLiteOpenHelper helper = new TestOpenHelper(22);
+        assertEquals(22, helper.getWritableDatabase().getVersion());
+
+        // Check column does not exist
+        try (Cursor c = helper.getWritableDatabase().query(Favorites.TABLE_NAME,
+                null, null, null, null, null, null)) {
+            assertEquals(-1, c.getColumnIndex(Favorites.OPTIONS));
+
+            // Check data is present
+            assertEquals(10, c.getCount());
+        }
+        helper.close();
+
+        helper = new DatabaseHelper(mContext, DB_FILE, false) {
+            @Override
+            public void onOpen(SQLiteDatabase db) { }
+        };
+        assertEquals(LauncherProvider.SCHEMA_VERSION, helper.getWritableDatabase().getVersion());
+
+        try (Cursor c = helper.getWritableDatabase().query(Favorites.TABLE_NAME,
+                null, null, null, null, null, null)) {
+            // Check column exists
+            assertNotSame(-1, c.getColumnIndex(Favorites.OPTIONS));
+
+            // Check data is present
+            assertEquals(10, c.getCount());
+        }
+        helper.close();
+    }
+
+    @Test(expected = DowngradeFailException.class)
+    public void testDowngrade_fail_v20() throws Exception {
+        setupTestDb();
+
+        TestOpenHelper helper = new TestOpenHelper(20);
+        helper.getReadableDatabase().getVersion();
+    }
+
+    private void setupTestDb() throws Exception {
+        mSchemaFile.delete();
+        mDbFile.delete();
+
+        DbDowngradeHelper.updateSchemaFile(mSchemaFile, LauncherProvider.SCHEMA_VERSION, mContext);
+
+        DatabaseHelper dbHelper = new DatabaseHelper(mContext, DB_FILE, false) {
+            @Override
+            public void onOpen(SQLiteDatabase db) { }
+        };
+        // Insert mock data
+        for (int i = 0; i < 10; i++) {
+            ContentValues values = new ContentValues();
+            values.put(Favorites._ID, i);
+            values.put(Favorites.TITLE, "title " + i);
+            dbHelper.getWritableDatabase().insert(Favorites.TABLE_NAME, null, values);
+        }
+        dbHelper.close();
+    }
+
+    private class TestOpenHelper extends SQLiteOpenHelper {
+
+        public TestOpenHelper(int version) {
+            super(mContext, DB_FILE, null, version);
+        }
+
+        @Override
+        public void onCreate(SQLiteDatabase sqLiteDatabase) {
+            throw new RuntimeException("DB should already be created");
+        }
+
+        @Override
+        public void onUpgrade(SQLiteDatabase db, int oldVersion, int newVersion) {
+            throw new RuntimeException("Only downgrade supported");
+        }
+
+        @Override
+        public void onDowngrade(SQLiteDatabase db, int oldVersion, int newVersion) {
+            try {
+                DbDowngradeHelper.parse(mSchemaFile).onDowngrade(db, oldVersion, newVersion);
+            } catch (Exception e) {
+                throw new DowngradeFailException(e);
+            }
+        }
+    }
+
+    private static class DowngradeFailException extends RuntimeException {
+        public DowngradeFailException(Exception e) {
+            super(e);
+        }
+    }
+}
diff --git a/tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java b/tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
new file mode 100644
index 0000000000..004ed06b32
--- /dev/null
+++ b/tests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
@@ -0,0 +1,133 @@
+/*
+ * 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.util.LauncherModelHelper.TEST_ACTIVITY;
+import static com.android.launcher3.util.LauncherModelHelper.TEST_PACKAGE;
+
+import static org.junit.Assert.assertEquals;
+
+import android.content.Context;
+import android.content.pm.PackageInstaller;
+import android.content.pm.PackageInstaller.SessionParams;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.icons.BitmapInfo;
+import com.android.launcher3.model.data.FolderInfo;
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.util.LauncherLayoutBuilder;
+import com.android.launcher3.util.LauncherModelHelper;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/**
+ * Tests for layout parser for remote layout
+ */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class DefaultLayoutProviderTest {
+
+    private LauncherModelHelper mModelHelper;
+    private Context mTargetContext;
+
+    @Before
+    public void setUp() {
+        mModelHelper = new LauncherModelHelper();
+        mTargetContext = mModelHelper.sandboxContext;
+    }
+
+    @After
+    public void tearDown() {
+        mModelHelper.destroy();
+    }
+
+    @Test
+    public void testCustomProfileLoaded_with_icon_on_hotseat() throws Exception {
+        writeLayoutAndLoad(new LauncherLayoutBuilder().atHotseat(0)
+                .putApp(TEST_PACKAGE, TEST_ACTIVITY));
+
+        // Verify one item in hotseat
+        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);
+    }
+
+    @Test
+    public void testCustomProfileLoaded_with_folder() throws Exception {
+        writeLayoutAndLoad(new LauncherLayoutBuilder().atHotseat(0).putFolder(android.R.string.copy)
+                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
+                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
+                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
+                .build());
+
+        // Verify folder
+        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());
+    }
+
+    @Test
+    public void testCustomProfileLoaded_with_folder_custom_title() throws Exception {
+        writeLayoutAndLoad(new LauncherLayoutBuilder().atHotseat(0).putFolder("CustomFolder")
+                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
+                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
+                .addApp(TEST_PACKAGE, TEST_ACTIVITY)
+                .build());
+
+        // Verify folder
+        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());
+        assertEquals("CustomFolder", info.title.toString());
+    }
+
+    @Test
+    public void testCustomProfileLoaded_with_widget() throws Exception {
+        String pendingAppPkg = "com.test.pending";
+
+        // Add a placeholder session info so that the widget exists
+        SessionParams params = new SessionParams(SessionParams.MODE_FULL_INSTALL);
+        params.setAppPackageName(pendingAppPkg);
+        params.setAppIcon(BitmapInfo.LOW_RES_ICON);
+
+        PackageInstaller installer = mTargetContext.getPackageManager().getPackageInstaller();
+        installer.createSession(params);
+
+        writeLayoutAndLoad(new LauncherLayoutBuilder().atWorkspace(0, 1, 0)
+                .putWidget(pendingAppPkg, "PlaceholderWidget", 2, 2));
+
+        // Verify widget
+        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);
+    }
+
+    private void writeLayoutAndLoad(LauncherLayoutBuilder builder) throws Exception {
+        mModelHelper.setupDefaultLayoutProvider(builder).loadModelSync();
+    }
+}
diff --git a/tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java b/tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
new file mode 100644
index 0000000000..005389e5bb
--- /dev/null
+++ b/tests/src/com/android/launcher3/model/GridSizeMigrationTaskV2Test.java
@@ -0,0 +1,278 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.model;
+
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT;
+import static com.android.launcher3.LauncherSettings.Favorites.TMP_CONTENT_URI;
+import static com.android.launcher3.provider.LauncherDbUtils.dropTable;
+import static com.android.launcher3.util.LauncherModelHelper.APP_ICON;
+import static com.android.launcher3.util.LauncherModelHelper.DESKTOP;
+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.content.Intent;
+import android.database.Cursor;
+import android.database.sqlite.SQLiteDatabase;
+import android.graphics.Point;
+import android.os.Process;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.pm.UserCache;
+import com.android.launcher3.util.LauncherModelHelper;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.HashMap;
+import java.util.HashSet;
+
+/** Unit tests for {@link GridSizeMigrationTaskV2} */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class GridSizeMigrationTaskV2Test {
+
+    private LauncherModelHelper mModelHelper;
+    private Context mContext;
+    private SQLiteDatabase mDb;
+
+    private HashSet<String> mValidPackages;
+    private InvariantDeviceProfile mIdp;
+
+    private final String testPackage1 = "com.android.launcher3.validpackage1";
+    private final String testPackage2 = "com.android.launcher3.validpackage2";
+    private final String testPackage3 = "com.android.launcher3.validpackage3";
+    private final String testPackage4 = "com.android.launcher3.validpackage4";
+    private final String testPackage5 = "com.android.launcher3.validpackage5";
+    private final String testPackage6 = "com.android.launcher3.validpackage6";
+    private final String testPackage7 = "com.android.launcher3.validpackage7";
+    private final String testPackage8 = "com.android.launcher3.validpackage8";
+    private final String testPackage9 = "com.android.launcher3.validpackage9";
+    private final String testPackage10 = "com.android.launcher3.validpackage10";
+
+    @Before
+    public void setUp() {
+        mModelHelper = new LauncherModelHelper();
+        mContext = mModelHelper.sandboxContext;
+        mDb = mModelHelper.provider.getDb();
+
+        mValidPackages = new HashSet<>();
+        mValidPackages.add(TEST_PACKAGE);
+        mValidPackages.add(testPackage1);
+        mValidPackages.add(testPackage2);
+        mValidPackages.add(testPackage3);
+        mValidPackages.add(testPackage4);
+        mValidPackages.add(testPackage5);
+        mValidPackages.add(testPackage6);
+        mValidPackages.add(testPackage7);
+        mValidPackages.add(testPackage8);
+        mValidPackages.add(testPackage9);
+        mValidPackages.add(testPackage10);
+
+        mIdp = InvariantDeviceProfile.INSTANCE.get(mContext);
+
+        long userSerial = UserCache.INSTANCE.get(mContext).getSerialNumberForUser(
+                Process.myUserHandle());
+        dropTable(mDb, LauncherSettings.Favorites.TMP_TABLE);
+        LauncherSettings.Favorites.addTableToDb(mDb, userSerial, false,
+                LauncherSettings.Favorites.TMP_TABLE);
+    }
+
+    @After
+    public void tearDown() {
+        mModelHelper.destroy();
+    }
+
+    @Test
+    public void testMigration() throws Exception {
+        int[] srcHotseatItems = {
+                mModelHelper.addItem(APP_ICON, 0, HOTSEAT, 0, 0, testPackage1, 1, TMP_CONTENT_URI),
+                mModelHelper.addItem(SHORTCUT, 1, HOTSEAT, 0, 0, testPackage2, 2, TMP_CONTENT_URI),
+                -1,
+                mModelHelper.addItem(SHORTCUT, 3, HOTSEAT, 0, 0, testPackage3, 3, TMP_CONTENT_URI),
+                mModelHelper.addItem(APP_ICON, 4, HOTSEAT, 0, 0, testPackage4, 4, TMP_CONTENT_URI),
+        };
+        mModelHelper.addItem(APP_ICON, 0, DESKTOP, 2, 2, testPackage5, 5, TMP_CONTENT_URI);
+        mModelHelper.addItem(APP_ICON, 0, DESKTOP, 2, 3, testPackage6, 6, TMP_CONTENT_URI);
+        mModelHelper.addItem(APP_ICON, 0, DESKTOP, 4, 1, testPackage8, 8, TMP_CONTENT_URI);
+        mModelHelper.addItem(APP_ICON, 0, DESKTOP, 4, 2, testPackage9, 9, TMP_CONTENT_URI);
+        mModelHelper.addItem(APP_ICON, 0, DESKTOP, 4, 3, testPackage10, 10, TMP_CONTENT_URI);
+
+        int[] destHotseatItems = {
+                -1,
+                mModelHelper.addItem(SHORTCUT, 1, HOTSEAT, 0, 0, testPackage2),
+                -1,
+        };
+        mModelHelper.addItem(APP_ICON, 0, DESKTOP, 2, 2, testPackage7);
+
+        mIdp.numDatabaseHotseatIcons = 4;
+        mIdp.numColumns = 4;
+        mIdp.numRows = 4;
+        GridSizeMigrationTaskV2.DbReader srcReader = new GridSizeMigrationTaskV2.DbReader(mDb,
+                LauncherSettings.Favorites.TMP_TABLE, mContext, mValidPackages);
+        GridSizeMigrationTaskV2.DbReader destReader = new GridSizeMigrationTaskV2.DbReader(mDb,
+                LauncherSettings.Favorites.TABLE_NAME, mContext, mValidPackages);
+        GridSizeMigrationTaskV2 task = new GridSizeMigrationTaskV2(mContext, mDb, srcReader,
+                destReader, mIdp.numDatabaseHotseatIcons, new Point(mIdp.numColumns, mIdp.numRows));
+        task.migrate(mIdp);
+
+        // Check hotseat items
+        Cursor c = mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
+                new String[]{LauncherSettings.Favorites.SCREEN, LauncherSettings.Favorites.INTENT},
+                "container=" + CONTAINER_HOTSEAT, null, LauncherSettings.Favorites.SCREEN, null);
+        assertEquals(c.getCount(), mIdp.numDatabaseHotseatIcons);
+        int screenIndex = c.getColumnIndex(LauncherSettings.Favorites.SCREEN);
+        int intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
+        c.moveToNext();
+        assertEquals(c.getInt(screenIndex), 0);
+        assertTrue(c.getString(intentIndex).contains(testPackage1));
+        c.moveToNext();
+        assertEquals(c.getInt(screenIndex), 1);
+        assertTrue(c.getString(intentIndex).contains(testPackage2));
+        c.moveToNext();
+        assertEquals(c.getInt(screenIndex), 2);
+        assertTrue(c.getString(intentIndex).contains(testPackage3));
+        c.moveToNext();
+        assertEquals(c.getInt(screenIndex), 3);
+        assertTrue(c.getString(intentIndex).contains(testPackage4));
+        c.close();
+
+        // Check workspace items
+        c = mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
+                new String[]{LauncherSettings.Favorites.CELLX, LauncherSettings.Favorites.CELLY,
+                        LauncherSettings.Favorites.INTENT},
+                "container=" + CONTAINER_DESKTOP, null, null, null);
+        intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
+        int cellXIndex = c.getColumnIndex(LauncherSettings.Favorites.CELLX);
+        int cellYIndex = c.getColumnIndex(LauncherSettings.Favorites.CELLY);
+
+        HashMap<String, Point> locMap = new HashMap<>();
+        while (c.moveToNext()) {
+            locMap.put(
+                    Intent.parseUri(c.getString(intentIndex), 0).getPackage(),
+                    new Point(c.getInt(cellXIndex), c.getInt(cellYIndex)));
+        }
+        c.close();
+
+        assertEquals(locMap.size(), 6);
+        assertEquals(new Point(0, 2), locMap.get(testPackage8));
+        assertEquals(new Point(0, 3), locMap.get(testPackage6));
+        assertEquals(new Point(1, 3), locMap.get(testPackage10));
+        assertEquals(new Point(2, 3), locMap.get(testPackage5));
+        assertEquals(new Point(3, 3), locMap.get(testPackage9));
+    }
+
+    @Test
+    public void migrateToLargerHotseat() {
+        int[] srcHotseatItems = {
+                mModelHelper.addItem(APP_ICON, 0, HOTSEAT, 0, 0, testPackage1, 1, TMP_CONTENT_URI),
+                mModelHelper.addItem(SHORTCUT, 1, HOTSEAT, 0, 0, testPackage2, 2, TMP_CONTENT_URI),
+                mModelHelper.addItem(APP_ICON, 2, HOTSEAT, 0, 0, testPackage3, 3, TMP_CONTENT_URI),
+                mModelHelper.addItem(SHORTCUT, 3, HOTSEAT, 0, 0, testPackage4, 4, TMP_CONTENT_URI),
+        };
+
+        int numSrcDatabaseHotseatIcons = srcHotseatItems.length;
+        mIdp.numDatabaseHotseatIcons = 6;
+        mIdp.numColumns = 4;
+        mIdp.numRows = 4;
+        GridSizeMigrationTaskV2.DbReader srcReader = new GridSizeMigrationTaskV2.DbReader(mDb,
+                LauncherSettings.Favorites.TMP_TABLE, mContext, mValidPackages);
+        GridSizeMigrationTaskV2.DbReader destReader = new GridSizeMigrationTaskV2.DbReader(mDb,
+                LauncherSettings.Favorites.TABLE_NAME, mContext, mValidPackages);
+        GridSizeMigrationTaskV2 task = new GridSizeMigrationTaskV2(mContext, mDb, srcReader,
+                destReader, mIdp.numDatabaseHotseatIcons, new Point(mIdp.numColumns, mIdp.numRows));
+        task.migrate(mIdp);
+
+        // Check hotseat items
+        Cursor c = mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
+                new String[]{LauncherSettings.Favorites.SCREEN, LauncherSettings.Favorites.INTENT},
+                "container=" + CONTAINER_HOTSEAT, null, LauncherSettings.Favorites.SCREEN, null);
+        assertEquals(c.getCount(), numSrcDatabaseHotseatIcons);
+        int screenIndex = c.getColumnIndex(LauncherSettings.Favorites.SCREEN);
+        int intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
+        c.moveToNext();
+        assertEquals(c.getInt(screenIndex), 0);
+        assertTrue(c.getString(intentIndex).contains(testPackage1));
+        c.moveToNext();
+        assertEquals(c.getInt(screenIndex), 1);
+        assertTrue(c.getString(intentIndex).contains(testPackage2));
+        c.moveToNext();
+        assertEquals(c.getInt(screenIndex), 2);
+        assertTrue(c.getString(intentIndex).contains(testPackage3));
+        c.moveToNext();
+        assertEquals(c.getInt(screenIndex), 3);
+        assertTrue(c.getString(intentIndex).contains(testPackage4));
+
+        c.close();
+    }
+
+    @Test
+    public void migrateFromLargerHotseat() {
+        int[] srcHotseatItems = {
+                mModelHelper.addItem(APP_ICON, 0, HOTSEAT, 0, 0, testPackage1, 1, TMP_CONTENT_URI),
+                -1,
+                mModelHelper.addItem(SHORTCUT, 2, HOTSEAT, 0, 0, testPackage2, 2, TMP_CONTENT_URI),
+                mModelHelper.addItem(APP_ICON, 3, HOTSEAT, 0, 0, testPackage3, 3, TMP_CONTENT_URI),
+                mModelHelper.addItem(SHORTCUT, 4, HOTSEAT, 0, 0, testPackage4, 4, TMP_CONTENT_URI),
+                mModelHelper.addItem(APP_ICON, 5, HOTSEAT, 0, 0, testPackage5, 5, TMP_CONTENT_URI),
+        };
+
+        mIdp.numDatabaseHotseatIcons = 4;
+        mIdp.numColumns = 4;
+        mIdp.numRows = 4;
+        GridSizeMigrationTaskV2.DbReader srcReader = new GridSizeMigrationTaskV2.DbReader(mDb,
+                LauncherSettings.Favorites.TMP_TABLE, mContext, mValidPackages);
+        GridSizeMigrationTaskV2.DbReader destReader = new GridSizeMigrationTaskV2.DbReader(mDb,
+                LauncherSettings.Favorites.TABLE_NAME, mContext, mValidPackages);
+        GridSizeMigrationTaskV2 task = new GridSizeMigrationTaskV2(mContext, mDb, srcReader,
+                destReader, mIdp.numDatabaseHotseatIcons, new Point(mIdp.numColumns, mIdp.numRows));
+        task.migrate(mIdp);
+
+        // Check hotseat items
+        Cursor c = mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
+                new String[]{LauncherSettings.Favorites.SCREEN, LauncherSettings.Favorites.INTENT},
+                "container=" + CONTAINER_HOTSEAT, null, LauncherSettings.Favorites.SCREEN, null);
+        assertEquals(c.getCount(), mIdp.numDatabaseHotseatIcons);
+        int screenIndex = c.getColumnIndex(LauncherSettings.Favorites.SCREEN);
+        int intentIndex = c.getColumnIndex(LauncherSettings.Favorites.INTENT);
+        c.moveToNext();
+        assertEquals(c.getInt(screenIndex), 0);
+        assertTrue(c.getString(intentIndex).contains(testPackage1));
+        c.moveToNext();
+        assertEquals(c.getInt(screenIndex), 1);
+        assertTrue(c.getString(intentIndex).contains(testPackage2));
+        c.moveToNext();
+        assertEquals(c.getInt(screenIndex), 2);
+        assertTrue(c.getString(intentIndex).contains(testPackage3));
+        c.moveToNext();
+        assertEquals(c.getInt(screenIndex), 3);
+        assertTrue(c.getString(intentIndex).contains(testPackage4));
+
+        c.close();
+    }
+}
diff --git a/tests/src/com/android/launcher3/model/LoaderCursorTest.java b/tests/src/com/android/launcher3/model/LoaderCursorTest.java
new file mode 100644
index 0000000000..6444ef6927
--- /dev/null
+++ b/tests/src/com/android/launcher3/model/LoaderCursorTest.java
@@ -0,0 +1,229 @@
+/*
+ * 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 androidx.test.InstrumentationRegistry.getContext;
+
+import static com.android.launcher3.LauncherSettings.Favorites.CELLX;
+import static com.android.launcher3.LauncherSettings.Favorites.CELLY;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT;
+import static com.android.launcher3.LauncherSettings.Favorites.ICON;
+import static com.android.launcher3.LauncherSettings.Favorites.ICON_PACKAGE;
+import static com.android.launcher3.LauncherSettings.Favorites.ICON_RESOURCE;
+import static com.android.launcher3.LauncherSettings.Favorites.INTENT;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT;
+import static com.android.launcher3.LauncherSettings.Favorites.PROFILE_ID;
+import static com.android.launcher3.LauncherSettings.Favorites.RESTORED;
+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_ACTIVITY;
+
+import static junit.framework.Assert.assertEquals;
+import static junit.framework.Assert.assertFalse;
+import static junit.framework.Assert.assertNotNull;
+import static junit.framework.Assert.assertNull;
+import static junit.framework.Assert.assertTrue;
+
+import android.content.ComponentName;
+import android.content.Context;
+import android.content.Intent;
+import android.database.MatrixCursor;
+import android.os.Process;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherSettings.Favorites;
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.model.data.WorkspaceItemInfo;
+import com.android.launcher3.util.Executors;
+import com.android.launcher3.util.LauncherModelHelper;
+import com.android.launcher3.util.PackageManagerHelper;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/**
+ * Tests for {@link LoaderCursor}
+ */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class LoaderCursorTest {
+
+    private LauncherModelHelper mModelHelper;
+    private LauncherAppState mApp;
+
+    private MatrixCursor mCursor;
+    private InvariantDeviceProfile mIDP;
+    private Context mContext;
+
+    private LoaderCursor mLoaderCursor;
+
+    @Before
+    public void setup() {
+        mModelHelper = new LauncherModelHelper();
+        mContext = mModelHelper.sandboxContext;
+        mIDP = InvariantDeviceProfile.INSTANCE.get(mContext);
+        mApp = LauncherAppState.getInstance(mContext);
+
+        mCursor = new MatrixCursor(new String[] {
+                ICON, ICON_PACKAGE, ICON_RESOURCE, TITLE,
+                _ID, CONTAINER, ITEM_TYPE, PROFILE_ID,
+                SCREEN, CELLX, CELLY, RESTORED, INTENT
+        });
+
+        UserManagerState ums = new UserManagerState();
+        mLoaderCursor = new LoaderCursor(mCursor, Favorites.CONTENT_URI, mApp, ums);
+        ums.allUsers.put(0, Process.myUserHandle());
+    }
+
+    @After
+    public void tearDown() {
+        mModelHelper.destroy();
+    }
+
+    private void initCursor(int itemType, String title) {
+        mCursor.newRow()
+                .add(_ID, 1)
+                .add(PROFILE_ID, 0)
+                .add(ITEM_TYPE, itemType)
+                .add(TITLE, title)
+                .add(CONTAINER, CONTAINER_DESKTOP);
+    }
+
+    @Test
+    public void getAppShortcutInfo_dontAllowMissing_invalidComponent() {
+        initCursor(ITEM_TYPE_APPLICATION, "");
+        assertTrue(mLoaderCursor.moveToNext());
+        ComponentName cn = new ComponentName(mContext.getPackageName(), "placeholder-do");
+        assertNull(mLoaderCursor.getAppShortcutInfo(
+                new Intent().setComponent(cn), false /* allowMissingTarget */, true));
+    }
+
+    @Test
+    public void getAppShortcutInfo_dontAllowMissing_validComponent() throws Exception {
+        ComponentName cn = new ComponentName(getContext(), TEST_ACTIVITY);
+        initCursor(ITEM_TYPE_APPLICATION, "");
+        assertTrue(mLoaderCursor.moveToNext());
+
+        WorkspaceItemInfo info = Executors.MODEL_EXECUTOR.submit(() ->
+                mLoaderCursor.getAppShortcutInfo(
+                        new Intent().setComponent(cn), false  /* allowMissingTarget */, true))
+                .get();
+        assertNotNull(info);
+        assertTrue(PackageManagerHelper.isLauncherAppTarget(info.getIntent()));
+    }
+
+    @Test
+    public void getAppShortcutInfo_allowMissing_invalidComponent() throws Exception {
+        initCursor(ITEM_TYPE_APPLICATION, "");
+        assertTrue(mLoaderCursor.moveToNext());
+
+        ComponentName cn = new ComponentName(mContext.getPackageName(), "placeholder-do");
+        WorkspaceItemInfo info = Executors.MODEL_EXECUTOR.submit(() ->
+                mLoaderCursor.getAppShortcutInfo(
+                        new Intent().setComponent(cn), true  /* allowMissingTarget */, true))
+                .get();
+        assertNotNull(info);
+        assertTrue(PackageManagerHelper.isLauncherAppTarget(info.getIntent()));
+    }
+
+    @Test
+    public void loadSimpleShortcut() {
+        initCursor(ITEM_TYPE_SHORTCUT, "my-shortcut");
+        assertTrue(mLoaderCursor.moveToNext());
+
+        WorkspaceItemInfo info = mLoaderCursor.loadSimpleWorkspaceItem();
+        assertTrue(mApp.getIconCache().isDefaultIcon(info.bitmap, info.user));
+        assertEquals("my-shortcut", info.title);
+        assertEquals(ITEM_TYPE_SHORTCUT, info.itemType);
+    }
+
+    @Test
+    public void checkItemPlacement_outsideBounds() {
+        mIDP.numRows = 4;
+        mIDP.numColumns = 4;
+        mIDP.numDatabaseHotseatIcons = 3;
+
+        // Item outside screen bounds are not placed
+        assertFalse(mLoaderCursor.checkItemPlacement(
+                newItemInfo(4, 4, 1, 1, CONTAINER_DESKTOP, 1)));
+    }
+
+    @Test
+    public void checkItemPlacement_overlappingItems() {
+        mIDP.numRows = 4;
+        mIDP.numColumns = 4;
+        mIDP.numDatabaseHotseatIcons = 3;
+
+        // Overlapping mItems are not placed
+        assertTrue(mLoaderCursor.checkItemPlacement(
+                newItemInfo(0, 0, 1, 1, CONTAINER_DESKTOP, 1)));
+        assertFalse(mLoaderCursor.checkItemPlacement(
+                newItemInfo(0, 0, 1, 1, CONTAINER_DESKTOP, 1)));
+
+        assertTrue(mLoaderCursor.checkItemPlacement(
+                newItemInfo(0, 0, 1, 1, CONTAINER_DESKTOP, 2)));
+        assertFalse(mLoaderCursor.checkItemPlacement(
+                newItemInfo(0, 0, 1, 1, CONTAINER_DESKTOP, 2)));
+
+        assertTrue(mLoaderCursor.checkItemPlacement(
+                newItemInfo(1, 1, 1, 1, CONTAINER_DESKTOP, 1)));
+        assertTrue(mLoaderCursor.checkItemPlacement(
+                newItemInfo(2, 2, 2, 2, CONTAINER_DESKTOP, 1)));
+
+        assertFalse(mLoaderCursor.checkItemPlacement(
+                newItemInfo(3, 2, 1, 2, CONTAINER_DESKTOP, 1)));
+    }
+
+    @Test
+    public void checkItemPlacement_hotseat() {
+        mIDP.numRows = 4;
+        mIDP.numColumns = 4;
+        mIDP.numDatabaseHotseatIcons = 3;
+
+        // Hotseat mItems are only placed based on screenId
+        assertTrue(mLoaderCursor.checkItemPlacement(
+                newItemInfo(3, 3, 1, 1, CONTAINER_HOTSEAT, 1)));
+        assertTrue(mLoaderCursor.checkItemPlacement(
+                newItemInfo(3, 3, 1, 1, CONTAINER_HOTSEAT, 2)));
+
+        assertFalse(mLoaderCursor.checkItemPlacement(
+                newItemInfo(3, 3, 1, 1, CONTAINER_HOTSEAT, 3)));
+    }
+
+    private ItemInfo newItemInfo(int cellX, int cellY, int spanX, int spanY,
+            int container, int screenId) {
+        ItemInfo info = new ItemInfo();
+        info.cellX = cellX;
+        info.cellY = cellY;
+        info.spanX = spanX;
+        info.spanY = spanY;
+        info.container = container;
+        info.screenId = screenId;
+        return info;
+    }
+}
diff --git a/tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java b/tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
new file mode 100644
index 0000000000..42c9f11a52
--- /dev/null
+++ b/tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
@@ -0,0 +1,234 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.model;
+
+import static com.android.launcher3.util.LauncherModelHelper.TEST_PACKAGE;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.spy;
+
+import android.os.Process;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.launcher3.model.BgDataModel.Callbacks;
+import com.android.launcher3.model.data.AppInfo;
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.util.Executors;
+import com.android.launcher3.util.IntArray;
+import com.android.launcher3.util.IntSet;
+import com.android.launcher3.util.LauncherLayoutBuilder;
+import com.android.launcher3.util.LauncherModelHelper;
+import com.android.launcher3.util.RunnableList;
+import com.android.launcher3.util.TestUtil;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * Tests to verify multiple callbacks in Loader
+ */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class ModelMultiCallbacksTest {
+
+    private LauncherModelHelper mModelHelper;
+
+    @Before
+    public void setUp() {
+        mModelHelper = new LauncherModelHelper();
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        mModelHelper.destroy();
+        TestUtil.uninstallDummyApp();
+    }
+
+    @Test
+    public void testTwoCallbacks_loadedTogether() throws Exception {
+        setupWorkspacePages(3);
+
+        MyCallbacks cb1 = spy(MyCallbacks.class);
+        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().addCallbacksAndLoad(cb1));
+
+        waitForLoaderAndTempMainThread();
+        cb1.verifySynchronouslyBound(3);
+
+        // Add a new callback
+        cb1.reset();
+        MyCallbacks cb2 = spy(MyCallbacks.class);
+        cb2.mPageToBindSync = IntSet.wrap(2);
+        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().addCallbacksAndLoad(cb2));
+
+        waitForLoaderAndTempMainThread();
+        assertFalse(cb1.bindStarted);
+        cb2.verifySynchronouslyBound(3);
+
+        // Remove callbacks
+        cb1.reset();
+        cb2.reset();
+
+        // No effect on callbacks when removing an callback
+        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().removeCallbacks(cb2));
+        waitForLoaderAndTempMainThread();
+        assertNull(cb1.mPendingTasks);
+        assertNull(cb2.mPendingTasks);
+
+        // Reloading only loads registered callbacks
+        mModelHelper.getModel().startLoader();
+        waitForLoaderAndTempMainThread();
+        cb1.verifySynchronouslyBound(3);
+        assertNull(cb2.mPendingTasks);
+    }
+
+    @Test
+    public void testTwoCallbacks_receiveUpdates() throws Exception {
+        TestUtil.uninstallDummyApp();
+
+        setupWorkspacePages(1);
+
+        MyCallbacks cb1 = spy(MyCallbacks.class);
+        MyCallbacks cb2 = spy(MyCallbacks.class);
+        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().addCallbacksAndLoad(cb1));
+        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().addCallbacksAndLoad(cb2));
+        waitForLoaderAndTempMainThread();
+
+        assertTrue(cb1.allApps().contains(TEST_PACKAGE));
+        assertTrue(cb2.allApps().contains(TEST_PACKAGE));
+
+        // Install package 1
+        TestUtil.installDummyApp();
+        mModelHelper.getModel().onPackageAdded(TestUtil.DUMMY_PACKAGE, Process.myUserHandle());
+        waitForLoaderAndTempMainThread();
+        assertTrue(cb1.allApps().contains(TestUtil.DUMMY_PACKAGE));
+        assertTrue(cb2.allApps().contains(TestUtil.DUMMY_PACKAGE));
+
+        // Uninstall package 2
+        TestUtil.uninstallDummyApp();
+        mModelHelper.getModel().onPackageRemoved(TestUtil.DUMMY_PACKAGE, Process.myUserHandle());
+        waitForLoaderAndTempMainThread();
+        assertFalse(cb1.allApps().contains(TestUtil.DUMMY_PACKAGE));
+        assertFalse(cb2.allApps().contains(TestUtil.DUMMY_PACKAGE));
+
+        // Unregister a callback and verify updates no longer received
+        Executors.MAIN_EXECUTOR.execute(() -> mModelHelper.getModel().removeCallbacks(cb2));
+        TestUtil.installDummyApp();
+        mModelHelper.getModel().onPackageAdded(TestUtil.DUMMY_PACKAGE, Process.myUserHandle());
+        waitForLoaderAndTempMainThread();
+
+        // cb2 didn't get the update
+        assertTrue(cb1.allApps().contains(TestUtil.DUMMY_PACKAGE));
+        assertFalse(cb2.allApps().contains(TestUtil.DUMMY_PACKAGE));
+    }
+
+    private void waitForLoaderAndTempMainThread() throws Exception {
+        Executors.MAIN_EXECUTOR.submit(() -> { }).get();
+        Executors.MODEL_EXECUTOR.submit(() -> { }).get();
+        Executors.MAIN_EXECUTOR.submit(() -> { }).get();
+    }
+
+    private void setupWorkspacePages(int pageCount) throws Exception {
+        // Create a layout with 3 pages
+        LauncherLayoutBuilder builder = new LauncherLayoutBuilder();
+        for (int i = 0; i < pageCount; i++) {
+            builder.atWorkspace(1, 1, i).putApp(TEST_PACKAGE, TEST_PACKAGE);
+        }
+        mModelHelper.setupDefaultLayoutProvider(builder);
+    }
+
+    private abstract static class MyCallbacks implements Callbacks {
+
+        final List<ItemInfo> mItems = new ArrayList<>();
+        IntSet mPageToBindSync = IntSet.wrap(0);
+        IntSet mPageBoundSync = new IntSet();
+        RunnableList mPendingTasks;
+        AppInfo[] mAppInfos;
+        boolean bindStarted;
+
+        MyCallbacks() { }
+
+        @Override
+        public void startBinding() {
+            bindStarted = true;
+        }
+
+        @Override
+        public void onInitialBindComplete(IntSet boundPages, RunnableList pendingTasks) {
+            mPageBoundSync = boundPages;
+            mPendingTasks = pendingTasks;
+        }
+
+        @Override
+        public void bindItems(List<ItemInfo> shortcuts, boolean forceAnimateIcons) {
+            mItems.addAll(shortcuts);
+        }
+
+        @Override
+        public void bindAllApplications(AppInfo[] apps, int flags) {
+            mAppInfos = apps;
+        }
+
+        @Override
+        public IntSet getPagesToBindSynchronously(IntArray orderedScreenIds) {
+            return mPageToBindSync;
+        }
+
+        public void reset() {
+            mItems.clear();
+            mPageBoundSync = new IntSet();
+            mPendingTasks = null;
+            mAppInfos = null;
+            bindStarted = false;
+        }
+
+        public void verifySynchronouslyBound(int totalItems) {
+            // Verify that the requested page is bound synchronously
+            assertTrue(bindStarted);
+            assertEquals(mPageToBindSync, mPageBoundSync);
+            assertEquals(mItems.size(), 1);
+            assertEquals(IntSet.wrap(mItems.get(0).screenId), mPageBoundSync);
+            assertNotNull(mPendingTasks);
+
+            // Verify that all other pages are bound properly
+            mPendingTasks.executeAllAndDestroy();
+            assertEquals(mItems.size(), totalItems);
+        }
+
+        public Set<String> allApps() {
+            return Arrays.stream(mAppInfos)
+                    .map(ai -> ai.getTargetComponent().getPackageName())
+                    .collect(Collectors.toSet());
+        }
+
+        public void verifyApps(String... apps) {
+            assertTrue(allApps().containsAll(Arrays.asList(apps)));
+        }
+    }
+}
diff --git a/tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java b/tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
new file mode 100644
index 0000000000..519191e251
--- /dev/null
+++ b/tests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
@@ -0,0 +1,83 @@
+package com.android.launcher3.model;
+
+import static org.junit.Assert.assertEquals;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.model.data.LauncherAppWidgetInfo;
+import com.android.launcher3.model.data.WorkspaceItemInfo;
+import com.android.launcher3.pm.PackageInstallInfo;
+import com.android.launcher3.util.LauncherModelHelper;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.Arrays;
+import java.util.HashSet;
+
+/**
+ * Tests for {@link PackageInstallStateChangedTask}
+ */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class PackageInstallStateChangedTaskTest {
+
+    private LauncherModelHelper mModelHelper;
+
+    @Before
+    public void setup() throws Exception {
+        mModelHelper = new LauncherModelHelper();
+        mModelHelper.initializeData("package_install_state_change_task_data");
+    }
+
+    @After
+    public void tearDown() {
+        mModelHelper.destroy();
+    }
+
+    private PackageInstallStateChangedTask newTask(String pkg, int progress) {
+        int state = PackageInstallInfo.STATUS_INSTALLING;
+        PackageInstallInfo installInfo = new PackageInstallInfo(pkg, state, progress,
+                android.os.Process.myUserHandle());
+        return new PackageInstallStateChangedTask(installInfo);
+    }
+
+    @Test
+    public void testSessionUpdate_ignore_installed() throws Exception {
+        mModelHelper.executeTaskForTest(newTask("app1", 30));
+
+        // No shortcuts were updated
+        verifyProgressUpdate(0);
+    }
+
+    @Test
+    public void testSessionUpdate_shortcuts_updated() throws Exception {
+        mModelHelper.executeTaskForTest(newTask("app3", 30));
+
+        verifyProgressUpdate(30, 5, 6, 7);
+    }
+
+    @Test
+    public void testSessionUpdate_widgets_updated() throws Exception {
+        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 : mModelHelper.getBgDataModel().itemsIdMap) {
+            if (info instanceof WorkspaceItemInfo) {
+                assertEquals(updates.contains(info.id) ? progress: 100,
+                        ((WorkspaceItemInfo) info).getProgressLevel());
+            } else {
+                assertEquals(updates.contains(info.id) ? progress: -1,
+                        ((LauncherAppWidgetInfo) info).installProgress);
+            }
+        }
+    }
+}
diff --git a/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java b/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
new file mode 100644
index 0000000000..48305eebdf
--- /dev/null
+++ b/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
@@ -0,0 +1,115 @@
+/*
+ * 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;
+
+import android.content.ContentValues;
+import android.database.Cursor;
+import android.database.sqlite.SQLiteDatabase;
+
+import androidx.test.InstrumentationRegistry;
+import androidx.test.filters.SmallTest;
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.launcher3.LauncherProvider.DatabaseHelper;
+import com.android.launcher3.LauncherSettings.Favorites;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/**
+ * Tests for {@link RestoreDbTask}
+ */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class RestoreDbTaskTest {
+
+    @Test
+    public void testGetProfileId() throws Exception {
+        SQLiteDatabase db = new MyDatabaseHelper(23).getWritableDatabase();
+        assertEquals(23, new RestoreDbTask().getDefaultProfileId(db));
+    }
+
+    @Test
+    public void testMigrateProfileId() throws Exception {
+        SQLiteDatabase db = new MyDatabaseHelper(42).getWritableDatabase();
+        // Add some mock data
+        for (int i = 0; i < 5; i++) {
+            ContentValues values = new ContentValues();
+            values.put(Favorites._ID, i);
+            values.put(Favorites.TITLE, "item " + i);
+            db.insert(Favorites.TABLE_NAME, null, values);
+        }
+        // Verify item add
+        assertEquals(5, getCount(db, "select * from favorites where profileId = 42"));
+
+        new RestoreDbTask().migrateProfileId(db, 42, 33);
+
+        // verify data migrated
+        assertEquals(0, getCount(db, "select * from favorites where profileId = 42"));
+        assertEquals(5, getCount(db, "select * from favorites where profileId = 33"));
+    }
+
+    @Test
+    public void testChangeDefaultColumn() throws Exception {
+        SQLiteDatabase db = new MyDatabaseHelper(42).getWritableDatabase();
+        // Add some mock data
+        for (int i = 0; i < 5; i++) {
+            ContentValues values = new ContentValues();
+            values.put(Favorites._ID, i);
+            values.put(Favorites.TITLE, "item " + i);
+            db.insert(Favorites.TABLE_NAME, null, values);
+        }
+        // Verify default column is 42
+        assertEquals(5, getCount(db, "select * from favorites where profileId = 42"));
+
+        new RestoreDbTask().changeDefaultColumn(db, 33);
+
+        // Verify default value changed
+        ContentValues values = new ContentValues();
+        values.put(Favorites._ID, 100);
+        values.put(Favorites.TITLE, "item 100");
+        db.insert(Favorites.TABLE_NAME, null, values);
+        assertEquals(1, getCount(db, "select * from favorites where profileId = 33"));
+    }
+
+    private int getCount(SQLiteDatabase db, String sql) {
+        try (Cursor c = db.rawQuery(sql, null)) {
+            return c.getCount();
+        }
+    }
+
+    private class MyDatabaseHelper extends DatabaseHelper {
+
+        private final long mProfileId;
+
+        MyDatabaseHelper(long profileId) {
+            super(InstrumentationRegistry.getInstrumentation().getTargetContext(), null, false);
+            mProfileId = profileId;
+        }
+
+        @Override
+        public long getDefaultUserSerial() {
+            return mProfileId;
+        }
+
+        @Override
+        protected void handleOneTimeDataUpgrade(SQLiteDatabase db) { }
+
+        protected void onEmptyDbCreated() { }
+    }
+}
diff --git a/tests/src/com/android/launcher3/secondarydisplay/SDLauncherTest.java b/tests/src/com/android/launcher3/secondarydisplay/SDLauncherTest.java
new file mode 100644
index 0000000000..fd86cf1144
--- /dev/null
+++ b/tests/src/com/android/launcher3/secondarydisplay/SDLauncherTest.java
@@ -0,0 +1,53 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.secondarydisplay;
+
+import static androidx.test.core.app.ActivityScenario.launch;
+
+import androidx.test.core.app.ActivityScenario;
+import androidx.test.espresso.intent.Intents;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.MediumTest;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/**
+ * Tests for {@link SecondaryDisplayLauncher}
+ */
+@MediumTest
+@RunWith(AndroidJUnit4.class)
+public class SDLauncherTest {
+
+    @Before
+    public void setUp() {
+        Intents.init();
+    }
+
+    @After
+    public void tearDown() {
+        Intents.release();
+    }
+
+    @Test
+    public void testAllAppsListOpens() {
+        ActivityScenario<SecondaryDisplayLauncher> launcher =
+                launch(SecondaryDisplayLauncher.class);
+        launcher.onActivity(l -> l.showAppDrawer(true));
+    }
+}
diff --git a/tests/src/com/android/launcher3/util/LauncherLayoutBuilder.java b/tests/src/com/android/launcher3/util/LauncherLayoutBuilder.java
new file mode 100644
index 0000000000..4e21dce021
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/LauncherLayoutBuilder.java
@@ -0,0 +1,182 @@
+/**
+ * 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 android.text.TextUtils;
+import android.util.Pair;
+import android.util.Xml;
+
+import org.xmlpull.v1.XmlSerializer;
+
+import java.io.IOException;
+import java.io.StringWriter;
+import java.io.Writer;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Helper class to build xml for Launcher Layout
+ */
+public class LauncherLayoutBuilder {
+
+    // Object Tags
+    private static final String TAG_WORKSPACE = "workspace";
+    private static final String TAG_AUTO_INSTALL = "autoinstall";
+    private static final String TAG_FOLDER = "folder";
+    private static final String TAG_APPWIDGET = "appwidget";
+    private static final String TAG_EXTRA = "extra";
+
+    private static final String ATTR_CONTAINER = "container";
+    private static final String ATTR_RANK = "rank";
+
+    private static final String ATTR_PACKAGE_NAME = "packageName";
+    private static final String ATTR_CLASS_NAME = "className";
+    private static final String ATTR_TITLE = "title";
+    private static final String ATTR_TITLE_TEXT = "titleText";
+    private static final String ATTR_SCREEN = "screen";
+
+    // x and y can be specified as negative integers, in which case -1 represents the
+    // last row / column, -2 represents the second last, and so on.
+    private static final String ATTR_X = "x";
+    private static final String ATTR_Y = "y";
+    private static final String ATTR_SPAN_X = "spanX";
+    private static final String ATTR_SPAN_Y = "spanY";
+
+    private static final String ATTR_CHILDREN = "children";
+
+
+    // Style attrs -- "Extra"
+    private static final String ATTR_KEY = "key";
+    private static final String ATTR_VALUE = "value";
+
+    private static final String CONTAINER_DESKTOP = "desktop";
+    private static final String CONTAINER_HOTSEAT = "hotseat";
+
+    private final ArrayList<Pair<String, HashMap<String, Object>>> mNodes = new ArrayList<>();
+
+    public Location atHotseat(int rank) {
+        Location l = new Location();
+        l.items.put(ATTR_CONTAINER, CONTAINER_HOTSEAT);
+        l.items.put(ATTR_RANK, Integer.toString(rank));
+        return l;
+    }
+
+    public Location atWorkspace(int x, int y, int screen) {
+        Location l = new Location();
+        l.items.put(ATTR_CONTAINER, CONTAINER_DESKTOP);
+        l.items.put(ATTR_X, Integer.toString(x));
+        l.items.put(ATTR_Y, Integer.toString(y));
+        l.items.put(ATTR_SCREEN, Integer.toString(screen));
+        return l;
+    }
+
+    public String build() throws IOException {
+        StringWriter writer = new StringWriter();
+        build(writer);
+        return writer.toString();
+    }
+
+    public void build(Writer writer) throws IOException {
+        XmlSerializer serializer = Xml.newSerializer();
+        serializer.setOutput(writer);
+
+        serializer.startDocument("UTF-8", true);
+        serializer.startTag(null, TAG_WORKSPACE);
+        writeNodes(serializer, mNodes);
+        serializer.endTag(null, TAG_WORKSPACE);
+        serializer.endDocument();
+        serializer.flush();
+    }
+
+    private static void writeNodes(XmlSerializer serializer,
+            ArrayList<Pair<String, HashMap<String, Object>>> nodes) throws IOException {
+        for (Pair<String, HashMap<String, Object>> node : nodes) {
+            ArrayList<Pair<String, HashMap<String, Object>>> children = null;
+
+            serializer.startTag(null, node.first);
+            for (Map.Entry<String, Object> attr : node.second.entrySet()) {
+                if (ATTR_CHILDREN.equals(attr.getKey())) {
+                    children = (ArrayList<Pair<String, HashMap<String, Object>>>) attr.getValue();
+                } else {
+                    serializer.attribute(null, attr.getKey(), (String) attr.getValue());
+                }
+            }
+
+            if (children != null) {
+                writeNodes(serializer, children);
+            }
+            serializer.endTag(null, node.first);
+        }
+    }
+
+    public class Location {
+
+        final HashMap<String, Object> items = new HashMap<>();
+
+        public LauncherLayoutBuilder putApp(String packageName, String className) {
+            items.put(ATTR_PACKAGE_NAME, packageName);
+            items.put(ATTR_CLASS_NAME, TextUtils.isEmpty(className) ? packageName : className);
+            mNodes.add(Pair.create(TAG_AUTO_INSTALL, items));
+            return LauncherLayoutBuilder.this;
+        }
+
+        public LauncherLayoutBuilder putWidget(String packageName, String className,
+                int spanX, int spanY) {
+            items.put(ATTR_PACKAGE_NAME, packageName);
+            items.put(ATTR_CLASS_NAME, className);
+            items.put(ATTR_SPAN_X, Integer.toString(spanX));
+            items.put(ATTR_SPAN_Y, Integer.toString(spanY));
+            mNodes.add(Pair.create(TAG_APPWIDGET, items));
+            return LauncherLayoutBuilder.this;
+        }
+
+        public FolderBuilder putFolder(int titleResId) {
+            items.put(ATTR_TITLE, Integer.toString(titleResId));
+            return putFolder();
+        }
+
+        public FolderBuilder putFolder(String title) {
+            items.put(ATTR_TITLE_TEXT, title);
+            return putFolder();
+        }
+
+        private FolderBuilder putFolder() {
+            FolderBuilder folderBuilder = new FolderBuilder();
+            items.put(ATTR_CHILDREN, folderBuilder.mChildren);
+            mNodes.add(Pair.create(TAG_FOLDER, items));
+            return folderBuilder;
+        }
+    }
+
+    public class FolderBuilder {
+
+        final ArrayList<Pair<String, HashMap<String, Object>>> mChildren = new ArrayList<>();
+
+        public FolderBuilder addApp(String packageName, String className) {
+            HashMap<String, Object> items = new HashMap<>();
+            items.put(ATTR_PACKAGE_NAME, packageName);
+            items.put(ATTR_CLASS_NAME, TextUtils.isEmpty(className) ? packageName : className);
+            mChildren.add(Pair.create(TAG_AUTO_INSTALL, items));
+            return this;
+        }
+
+        public LauncherLayoutBuilder build() {
+            return LauncherLayoutBuilder.this;
+        }
+    }
+}
diff --git a/tests/src/com/android/launcher3/util/LauncherModelHelper.java b/tests/src/com/android/launcher3/util/LauncherModelHelper.java
new file mode 100644
index 0000000000..c9b63aeb11
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/LauncherModelHelper.java
@@ -0,0 +1,569 @@
+/*
+ * 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 androidx.test.core.app.ApplicationProvider.getApplicationContext;
+import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
+
+import static com.android.launcher3.LauncherSettings.Favorites.CONTENT_URI;
+import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
+
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.atLeast;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+
+import android.content.ComponentName;
+import android.content.ContentProvider;
+import android.content.ContentResolver;
+import android.content.ContentValues;
+import android.content.Context;
+import android.content.Intent;
+import android.content.pm.PackageManager;
+import android.content.pm.ProviderInfo;
+import android.content.res.Resources;
+import android.database.sqlite.SQLiteDatabase;
+import android.net.Uri;
+import android.os.ParcelFileDescriptor;
+import android.os.ParcelFileDescriptor.AutoCloseOutputStream;
+import android.os.Process;
+import android.provider.Settings;
+import android.test.mock.MockContentResolver;
+import android.util.ArrayMap;
+
+import androidx.test.core.app.ApplicationProvider;
+import androidx.test.uiautomator.UiDevice;
+
+import com.android.launcher3.InvariantDeviceProfile;
+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 com.android.launcher3.model.BgDataModel.Callbacks;
+import com.android.launcher3.model.ItemInstallQueue;
+import com.android.launcher3.model.data.AppInfo;
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.pm.InstallSessionHelper;
+import com.android.launcher3.pm.UserCache;
+import com.android.launcher3.testing.TestInformationProvider;
+import com.android.launcher3.uioverrides.plugins.PluginManagerWrapper;
+import com.android.launcher3.util.MainThreadInitializedObject.ObjectProvider;
+import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext;
+import com.android.launcher3.widget.custom.CustomWidgetManager;
+
+import org.mockito.ArgumentCaptor;
+
+import java.io.BufferedReader;
+import java.io.ByteArrayOutputStream;
+import java.io.File;
+import java.io.FileNotFoundException;
+import java.io.InputStreamReader;
+import java.io.OutputStreamWriter;
+import java.lang.reflect.Field;
+import java.util.HashMap;
+import java.util.List;
+import java.util.UUID;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+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 = testContext().getPackageName();
+    public static final String TEST_ACTIVITY = "com.android.launcher3.tests.Activity2";
+
+    // Authority for providing a test default-workspace-layout data.
+    private static final String TEST_PROVIDER_AUTHORITY =
+            LauncherModelHelper.class.getName().toLowerCase();
+    private static final int DEFAULT_BITMAP_SIZE = 10;
+    private static final int DEFAULT_GRID_SIZE = 4;
+
+    private final HashMap<Class, HashMap<String, Field>> mFieldCache = new HashMap<>();
+    private final MockContentResolver mMockResolver = new MockContentResolver();
+    public final TestLauncherProvider provider;
+    public final SanboxModelContext sandboxContext;
+
+    public final long defaultProfileId;
+
+    private BgDataModel mDataModel;
+    private AllAppsList mAllAppsList;
+
+    public LauncherModelHelper() {
+        Context context = getApplicationContext();
+        // System settings cache content provider. Ensure that they are statically initialized
+        Settings.Secure.getString(context.getContentResolver(), "test");
+        Settings.System.getString(context.getContentResolver(), "test");
+        Settings.Global.getString(context.getContentResolver(), "test");
+
+        provider = new TestLauncherProvider();
+        sandboxContext = new SanboxModelContext();
+        defaultProfileId = UserCache.INSTANCE.get(sandboxContext)
+                .getSerialNumberForUser(Process.myUserHandle());
+        setupProvider(LauncherProvider.AUTHORITY, provider);
+    }
+
+    protected void setupProvider(String authority, ContentProvider provider) {
+        ProviderInfo providerInfo = new ProviderInfo();
+        providerInfo.authority = authority;
+        providerInfo.applicationInfo = sandboxContext.getApplicationInfo();
+        provider.attachInfo(sandboxContext, providerInfo);
+        mMockResolver.addProvider(providerInfo.authority, provider);
+        doReturn(providerInfo)
+                .when(sandboxContext.mPm)
+                .resolveContentProvider(eq(authority), anyInt());
+    }
+
+    public LauncherModel getModel() {
+        return LauncherAppState.getInstance(sandboxContext).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;
+    }
+
+    public void destroy() {
+        // When destroying the context, make sure that the model thread is blocked, so that no
+        // new jobs get posted while we are cleaning up
+        CountDownLatch l1 = new CountDownLatch(1);
+        CountDownLatch l2 = new CountDownLatch(1);
+        MODEL_EXECUTOR.execute(() -> {
+            l1.countDown();
+            waitOrThrow(l2);
+        });
+        waitOrThrow(l1);
+        sandboxContext.onDestroy();
+        l2.countDown();
+    }
+
+    private void waitOrThrow(CountDownLatch latch) {
+        try {
+            latch.await();
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    /**
+     * 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 {
+        BgDataModel bgDataModel = getBgDataModel();
+        AllAppsList allAppsList = getAllAppsList();
+
+        MODEL_EXECUTOR.submit(() -> {
+            // Copy apk from resources to a local file and install from there.
+            Resources resources = testContext().getResources();
+            int resId = resources.getIdentifier(
+                    resourceName, "raw", testContext().getPackageName());
+            try (BufferedReader reader = new BufferedReader(new InputStreamReader(
+                    resources.openRawResource(resId)))) {
+                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(sandboxContext,
+                                    (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;
+    }
+
+    public int addItem(int type, int screen, int container, int x, int y) {
+        return addItem(type, screen, container, x, y, defaultProfileId, TEST_PACKAGE);
+    }
+
+    public int addItem(int type, int screen, int container, int x, int y, long profileId) {
+        return addItem(type, screen, container, x, y, profileId, TEST_PACKAGE);
+    }
+
+    public int addItem(int type, int screen, int container, int x, int y, String packageName) {
+        return addItem(type, screen, container, x, y, defaultProfileId, packageName);
+    }
+
+    public int addItem(int type, int screen, int container, int x, int y, String packageName,
+            int id, Uri contentUri) {
+        addItem(type, screen, container, x, y, defaultProfileId, packageName, id, contentUri);
+        return id;
+    }
+
+    /**
+     * Adds a mock 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, long profileId,
+            String packageName) {
+        int id = LauncherSettings.Settings.call(sandboxContext.getContentResolver(),
+                LauncherSettings.Settings.METHOD_NEW_ITEM_ID)
+                .getInt(LauncherSettings.Settings.EXTRA_VALUE);
+        addItem(type, screen, container, x, y, profileId, packageName, id, CONTENT_URI);
+        return id;
+    }
+
+    public void addItem(int type, int screen, int container, int x, int y, long profileId,
+            String packageName, int id, Uri contentUri) {
+        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);
+        values.put(LauncherSettings.Favorites.PROFILE_ID, profileId);
+
+        if (type == APP_ICON || type == SHORTCUT) {
+            values.put(LauncherSettings.Favorites.ITEM_TYPE, type);
+            values.put(LauncherSettings.Favorites.INTENT,
+                    new Intent(Intent.ACTION_MAIN).setPackage(packageName).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, profileId);
+            }
+        }
+
+        sandboxContext.getContentResolver().insert(contentUri, values);
+    }
+
+    public int[][][] createGrid(int[][][] typeArray) {
+        return createGrid(typeArray, 1);
+    }
+
+    public int[][][] createGrid(int[][][] typeArray, int startScreen) {
+        LauncherSettings.Settings.call(sandboxContext.getContentResolver(),
+                LauncherSettings.Settings.METHOD_CREATE_EMPTY_DB);
+        LauncherSettings.Settings.call(sandboxContext.getContentResolver(),
+                LauncherSettings.Settings.METHOD_CLEAR_EMPTY_DB_FLAG);
+        return createGrid(typeArray, startScreen, defaultProfileId);
+    }
+
+    /**
+     * Initializes the DB with mock 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, long profileId) {
+        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(sandboxContext.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, profileId);
+                    }
+                }
+            }
+        }
+
+        return ids;
+    }
+
+    /**
+     * Sets up a mock provider to load the provided layout by default, next time the layout loads
+     */
+    public LauncherModelHelper setupDefaultLayoutProvider(LauncherLayoutBuilder builder)
+            throws Exception {
+        InvariantDeviceProfile idp = InvariantDeviceProfile.INSTANCE.get(sandboxContext);
+        idp.numRows = idp.numColumns = idp.numDatabaseHotseatIcons = DEFAULT_GRID_SIZE;
+        idp.iconBitmapSize = DEFAULT_BITMAP_SIZE;
+
+        UiDevice.getInstance(getInstrumentation()).executeShellCommand(
+                "settings put secure launcher3.layout.provider " + TEST_PROVIDER_AUTHORITY);
+        ContentProvider cp = new TestInformationProvider() {
+
+            @Override
+            public ParcelFileDescriptor openFile(Uri uri, String mode)
+                    throws FileNotFoundException {
+                try {
+                    ParcelFileDescriptor[] pipe = ParcelFileDescriptor.createPipe();
+                    AutoCloseOutputStream outputStream = new AutoCloseOutputStream(pipe[1]);
+                    ByteArrayOutputStream bos = new ByteArrayOutputStream();
+                    builder.build(new OutputStreamWriter(bos));
+                    outputStream.write(bos.toByteArray());
+                    outputStream.flush();
+                    outputStream.close();
+                    return pipe[0];
+                } catch (Exception e) {
+                    throw new FileNotFoundException(e.getMessage());
+                }
+            }
+        };
+        setupProvider(TEST_PROVIDER_AUTHORITY, cp);
+        return this;
+    }
+
+    /**
+     * Loads the model in memory synchronously
+     */
+    public void loadModelSync() throws ExecutionException, InterruptedException {
+        Callbacks mockCb = new Callbacks() { };
+        Executors.MAIN_EXECUTOR.submit(() -> getModel().addCallbacksAndLoad(mockCb)).get();
+
+        Executors.MODEL_EXECUTOR.submit(() -> { }).get();
+        Executors.MAIN_EXECUTOR.submit(() -> { }).get();
+        Executors.MAIN_EXECUTOR.submit(() -> getModel().removeCallbacks(mockCb)).get();
+    }
+
+    /**
+     * An extension of LauncherProvider backed up by in-memory database.
+     */
+    public static class TestLauncherProvider extends LauncherProvider {
+
+        @Override
+        public boolean onCreate() {
+            return true;
+        }
+
+        public SQLiteDatabase getDb() {
+            createDbIfNotExists();
+            return mOpenHelper.getWritableDatabase();
+        }
+
+        public DatabaseHelper getHelper() {
+            return mOpenHelper;
+        }
+    }
+
+    public static boolean deleteContents(File dir) {
+        File[] files = dir.listFiles();
+        boolean success = true;
+        if (files != null) {
+            for (File file : files) {
+                if (file.isDirectory()) {
+                    success &= deleteContents(file);
+                }
+                if (!file.delete()) {
+                    success = false;
+                }
+            }
+        }
+        return success;
+    }
+
+    public class SanboxModelContext extends SandboxContext {
+
+        private final ArrayMap<String, Object> mSpiedServices = new ArrayMap<>();
+        private final PackageManager mPm;
+        private final File mDbDir;
+
+        SanboxModelContext() {
+            super(ApplicationProvider.getApplicationContext(),
+                    UserCache.INSTANCE, InstallSessionHelper.INSTANCE,
+                    LauncherAppState.INSTANCE, InvariantDeviceProfile.INSTANCE,
+                    DisplayController.INSTANCE, CustomWidgetManager.INSTANCE,
+                    SettingsCache.INSTANCE, PluginManagerWrapper.INSTANCE,
+                    ItemInstallQueue.INSTANCE);
+            mPm = spy(getBaseContext().getPackageManager());
+            mDbDir = new File(getCacheDir(), UUID.randomUUID().toString());
+        }
+
+        public SanboxModelContext allow(MainThreadInitializedObject object) {
+            mAllowedObjects.add(object);
+            return this;
+        }
+
+        @Override
+        public File getDatabasePath(String name) {
+            if (!mDbDir.exists()) {
+                mDbDir.mkdirs();
+            }
+            return new File(mDbDir, name);
+        }
+
+        @Override
+        public ContentResolver getContentResolver() {
+            return mMockResolver;
+        }
+
+        @Override
+        public void onDestroy() {
+            if (deleteContents(mDbDir)) {
+                mDbDir.delete();
+            }
+            super.onDestroy();
+        }
+
+
+        @Override
+        protected <T> T createObject(ObjectProvider<T> provider) {
+            return spy(provider.get(this));
+        }
+
+        @Override
+        public PackageManager getPackageManager() {
+            return mPm;
+        }
+
+        @Override
+        public Object getSystemService(String name) {
+            Object service = mSpiedServices.get(name);
+            return service != null ? service : super.getSystemService(name);
+        }
+
+        public <T> T spyService(Class<T> tClass) {
+            String name = getSystemServiceName(tClass);
+            Object service = mSpiedServices.get(name);
+            if (service != null) {
+                return (T) service;
+            }
+
+            T result = spy(getSystemService(tClass));
+            mSpiedServices.put(name, result);
+            return result;
+        }
+    }
+
+    private static Context testContext() {
+        return getInstrumentation().getContext();
+    }
+}
diff --git a/tests/src/com/android/launcher3/util/ReflectionHelpers.java b/tests/src/com/android/launcher3/util/ReflectionHelpers.java
new file mode 100644
index 0000000000..d89975de8a
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/ReflectionHelpers.java
@@ -0,0 +1,58 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.util;
+
+import java.lang.reflect.Field;
+
+public class ReflectionHelpers {
+
+    /**
+     * Reflectively get the value of a field.
+     *
+     * @param object Target object.
+     * @param fieldName The field name.
+     * @param <R> The return type.
+     * @return Value of the field on the object.
+     */
+    public static <R> R getField(Object object, String fieldName) {
+        try {
+            Field field = object.getClass().getDeclaredField(fieldName);
+            field.setAccessible(true);
+            return (R) field.get(object);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    /**
+     * Reflectively set the value of a field.
+     *
+     * @param object Target object.
+     * @param fieldName The field name.
+     * @param fieldNewValue New value.
+     */
+    public static void setField(Object object, String fieldName, Object fieldNewValue) {
+        try {
+            Field field = object.getClass().getDeclaredField(fieldName);
+            field.setAccessible(true);
+            field.set(object, fieldNewValue);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+}