diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index dce1ab8..d601322 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -83,7 +83,7 @@
     DeviceProfile landscapeProfile;
     DeviceProfile portraitProfile;
 
-    InvariantDeviceProfile() {
+    public InvariantDeviceProfile() {
     }
 
     public InvariantDeviceProfile(InvariantDeviceProfile p) {
diff --git a/src/com/android/launcher3/LauncherBackupAgentHelper.java b/src/com/android/launcher3/LauncherBackupAgentHelper.java
index bf9c668..b192ba3 100644
--- a/src/com/android/launcher3/LauncherBackupAgentHelper.java
+++ b/src/com/android/launcher3/LauncherBackupAgentHelper.java
@@ -101,12 +101,9 @@
                         LauncherSettings.Settings.METHOD_UPDATE_FOLDER_ITEMS_RANK);
             }
 
-            // TODO: Update this logic to handle grid difference of 2. as well as hotseat difference
             if (GridSizeMigrationTask.ENABLED && mHelper.shouldAttemptWorkspaceMigration()) {
                 GridSizeMigrationTask.markForMigration(getApplicationContext(),
-                        (int) mHelper.migrationCompatibleProfileData.desktopCols,
-                        (int) mHelper.migrationCompatibleProfileData.desktopRows,
-                        mHelper.widgetSizes);
+                        mHelper.widgetSizes, mHelper.migrationCompatibleProfileData);
             }
 
             LauncherSettings.Settings.call(getContentResolver(),
diff --git a/src/com/android/launcher3/LauncherBackupHelper.java b/src/com/android/launcher3/LauncherBackupHelper.java
index 5f58e28..05d729e 100644
--- a/src/com/android/launcher3/LauncherBackupHelper.java
+++ b/src/com/android/launcher3/LauncherBackupHelper.java
@@ -315,14 +315,13 @@
             return true;
         }
 
-        if (GridSizeMigrationTask.ENABLED &&
-                (oldProfile.desktopCols - currentProfile.desktopCols <= 1) &&
-                (oldProfile.desktopRows - currentProfile.desktopRows <= 1)) {
-            // Allow desktop migration when row and/or column count contracts by 1.
-
+        if (GridSizeMigrationTask.ENABLED) {
+            // One time migrate the workspace when launcher starts.
             migrationCompatibleProfileData = initDeviceProfileData(mIdp);
             migrationCompatibleProfileData.desktopCols = oldProfile.desktopCols;
             migrationCompatibleProfileData.desktopRows = oldProfile.desktopRows;
+            migrationCompatibleProfileData.hotseatCount = oldProfile.hotseatCount;
+            migrationCompatibleProfileData.allappsRank = oldProfile.allappsRank;
             return true;
         }
         return false;
diff --git a/src/com/android/launcher3/LauncherModel.java b/src/com/android/launcher3/LauncherModel.java
index 0eb1a90..92ef3ea 100644
--- a/src/com/android/launcher3/LauncherModel.java
+++ b/src/com/android/launcher3/LauncherModel.java
@@ -1651,25 +1651,10 @@
             int countX = profile.numColumns;
             int countY = profile.numRows;
 
-            if (GridSizeMigrationTask.ENABLED && GridSizeMigrationTask.shouldRunTask(mContext)) {
-                long migrationStartTime = System.currentTimeMillis();
-                Log.v(TAG, "Starting workspace migration after restore");
-                try {
-                    GridSizeMigrationTask task = new GridSizeMigrationTask(mContext);
-                    // Clear the flags before starting the task, so that we do not run the task
-                    // again, in case there was an uncaught error.
-                    GridSizeMigrationTask.clearFlags(mContext);
-                    task.execute();
-                } catch (Exception e) {
-                    Log.e(TAG, "Error during grid migration", e);
-
-                    // Clear workspace.
-                    mFlags = mFlags | LOADER_FLAG_CLEAR_WORKSPACE;
-                }
-                Log.v(TAG, "Workspace migration completed in "
-                        + (System.currentTimeMillis() - migrationStartTime));
-
-                GridSizeMigrationTask.saveCurrentConfig(mContext);
+            if (GridSizeMigrationTask.ENABLED &&
+                    !GridSizeMigrationTask.migrateGridIfNeeded(mContext)) {
+                // Migration failed. Clear workspace.
+                mFlags = mFlags | LOADER_FLAG_CLEAR_WORKSPACE;
             }
 
             if ((mFlags & LOADER_FLAG_CLEAR_WORKSPACE) != 0) {
diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index 3fc0b94..ac9b321 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -77,7 +77,7 @@
 
     private static final Object LISTENER_LOCK = new Object();
     @Thunk LauncherProviderChangeListener mListener;
-    @Thunk DatabaseHelper mOpenHelper;
+    protected DatabaseHelper mOpenHelper;
 
     @Override
     public boolean onCreate() {
@@ -104,7 +104,10 @@
         }
     }
 
-    private synchronized void createDbIfNotExists() {
+    /**
+     * Overridden in tests
+     */
+    protected synchronized void createDbIfNotExists() {
         if (mOpenHelper == null) {
             mOpenHelper = new DatabaseHelper(getContext(), this);
         }
@@ -364,7 +367,10 @@
         return folderIds;
     }
 
-    private void notifyListeners() {
+    /**
+     * Overridden in tests
+     */
+    protected void notifyListeners() {
         // always notify the backup agent
         LauncherBackupAgentHelper.dataChanged(getContext());
         synchronized (LISTENER_LOCK) {
@@ -501,7 +507,10 @@
         });
     }
 
-    private static class DatabaseHelper extends SQLiteOpenHelper implements LayoutParserCallback {
+    /**
+     * The class is subclassed in tests to create an in-memory db.
+     */
+    protected static class DatabaseHelper extends SQLiteOpenHelper implements LayoutParserCallback {
         private final LauncherProvider mProvider;
         private final Context mContext;
         @Thunk final AppWidgetHost mAppWidgetHost;
@@ -535,6 +544,19 @@
             }
         }
 
+        /**
+         * Constructor used only in tests.
+         */
+        public DatabaseHelper(Context context, LauncherProvider provider, String tableName) {
+            super(context, tableName, null, DATABASE_VERSION);
+            mContext = context;
+            mProvider = provider;
+
+            mAppWidgetHost = null;
+            mMaxItemId = initializeMaxItemId(getWritableDatabase());
+            mMaxScreenId = initializeMaxScreenId(getWritableDatabase());
+        }
+
         private boolean tableExists(String tableName) {
             Cursor c = getReadableDatabase().query(
                     true, "sqlite_master", new String[] {"tbl_name"},
@@ -565,18 +587,28 @@
 
             // Fresh and clean launcher DB.
             mMaxItemId = initializeMaxItemId(db);
-            setFlagEmptyDbCreated();
+            onEmptyDbCreated();
+        }
+
+        /**
+         * Overriden in tests.
+         */
+        protected void onEmptyDbCreated() {
+            // Set the flag for empty DB
+            Utilities.getPrefs(mContext).edit().putBoolean(EMPTY_DATABASE_CREATED, true).commit();
 
             // When a new DB is created, remove all previously stored managed profile information.
-            ManagedProfileHeuristic.processAllUsers(Collections.<UserHandleCompat>emptyList(), mContext);
+            ManagedProfileHeuristic.processAllUsers(Collections.<UserHandleCompat>emptyList(),
+                    mContext);
+        }
+
+        protected long getDefaultUserSerial() {
+            return UserManagerCompat.getInstance(mContext).getSerialNumberForUser(
+                    UserHandleCompat.myUserHandle());
         }
 
         private void addFavoritesTable(SQLiteDatabase db, boolean optional) {
-            UserManagerCompat userManager = UserManagerCompat.getInstance(mContext);
-            long userSerialNumber = userManager.getSerialNumberForUser(
-                    UserHandleCompat.myUserHandle());
             String ifNotExists = optional ? " IF NOT EXISTS " : "";
-
             db.execSQL("CREATE TABLE " + ifNotExists + TABLE_FAVORITES + " (" +
                     "_id INTEGER PRIMARY KEY," +
                     "title TEXT," +
@@ -599,7 +631,7 @@
                     "appWidgetProvider TEXT," +
                     "modified INTEGER NOT NULL DEFAULT 0," +
                     "restored INTEGER NOT NULL DEFAULT 0," +
-                    "profileId INTEGER DEFAULT " + userSerialNumber + "," +
+                    "profileId INTEGER DEFAULT " + getDefaultUserSerial() + "," +
                     "rank INTEGER NOT NULL DEFAULT 0," +
                     "options INTEGER NOT NULL DEFAULT 0" +
                     ");");
@@ -649,10 +681,6 @@
             Utilities.getPrefs(mContext).edit().putBoolean(EMPTY_DATABASE_CREATED, false).commit();
         }
 
-        private void setFlagEmptyDbCreated() {
-            Utilities.getPrefs(mContext).edit().putBoolean(EMPTY_DATABASE_CREATED, true).commit();
-        }
-
         @Override
         public void onUpgrade(SQLiteDatabase db, int oldVersion, int newVersion) {
             if (LOGD) Log.d(TAG, "onUpgrade triggered: " + oldVersion);
diff --git a/src/com/android/launcher3/LauncherSettings.java b/src/com/android/launcher3/LauncherSettings.java
index 9ee6a21..55a5378 100644
--- a/src/com/android/launcher3/LauncherSettings.java
+++ b/src/com/android/launcher3/LauncherSettings.java
@@ -115,7 +115,7 @@
         /**
          * The content:// style URL for this table
          */
-        static final Uri CONTENT_URI = Uri.parse("content://" +
+        public static final Uri CONTENT_URI = Uri.parse("content://" +
                 ProviderConfig.AUTHORITY + "/" + TABLE_NAME);
 
         /**
diff --git a/src/com/android/launcher3/model/GridSizeMigrationTask.java b/src/com/android/launcher3/model/GridSizeMigrationTask.java
index 08c3dc0..19ec3ed 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationTask.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationTask.java
@@ -9,6 +9,7 @@
 import android.content.pm.PackageInfo;
 import android.database.Cursor;
 import android.graphics.Point;
+import android.net.Uri;
 import android.text.TextUtils;
 import android.util.Log;
 
@@ -21,6 +22,7 @@
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.Utilities;
+import com.android.launcher3.backup.nano.BackupProtos;
 import com.android.launcher3.compat.PackageInstallerCompat;
 import com.android.launcher3.compat.UserHandleCompat;
 import com.android.launcher3.util.LongArrayMap;
@@ -58,15 +60,14 @@
     private static final float WT_FOLDER_FACTOR = 0.5f;
 
     private final Context mContext;
-    private final ContentValues mTempValues = new ContentValues();
-    private final HashMap<String, Point> mWidgetMinSize;
     private final InvariantDeviceProfile mIdp;
 
-    private HashSet<String> mValidPackages;
-    public ArrayList<Long> mEntryToRemove;
-    private ArrayList<ContentProviderOperation> mUpdateOperations;
-
-    private ArrayList<DbEntry> mCarryOver;
+    private final HashMap<String, Point> mWidgetMinSize = new HashMap<>();
+    private final ContentValues mTempValues = new ContentValues();
+    private final ArrayList<Long> mEntryToRemove = new ArrayList<>();
+    private final ArrayList<ContentProviderOperation> mUpdateOperations = new ArrayList<>();
+    private final ArrayList<DbEntry> mCarryOver = new ArrayList<>();
+    private final HashSet<String> mValidPackages;
 
     private final int mSrcX, mSrcY;
     private final int mTrgX, mTrgY;
@@ -74,73 +75,54 @@
 
     private final int mSrcHotseatSize;
     private final int mSrcAllAppsRank;
+    private final int mDestHotseatSize;
+    private final int mDestAllAppsRank;
 
-    /**
-     * TODO: Create a generic constructor which can be unit tested.
-     */
-    public GridSizeMigrationTask(Context context) {
+    protected GridSizeMigrationTask(Context context, InvariantDeviceProfile idp,
+            HashSet<String> validPackages, HashMap<String, Point> widgetMinSize,
+            Point sourceSize, Point targetSize) {
         mContext = context;
+        mValidPackages = validPackages;
+        mWidgetMinSize.putAll(widgetMinSize);
+        mIdp = idp;
 
-
-        mIdp = LauncherAppState.getInstance().getInvariantDeviceProfile();
-        mTrgX = mIdp.numColumns;
-        mTrgY = mIdp.numRows;
-
-        SharedPreferences prefs = Utilities.getPrefs(context);
-        Point sourceSize = parsePoint(
-                prefs.getString(KEY_MIGRATION_SRC_WORKSPACE_SIZE, getPointString(mTrgX, mTrgY)));
         mSrcX = sourceSize.x;
         mSrcY = sourceSize.y;
 
-        // Hotseat
-        Point hotseatSize = parsePoint(
-                prefs.getString(KEY_MIGRATION_SRC_HOTSEAT_SIZE,
-                        getPointString(mIdp.numHotseatIcons, mIdp.hotseatAllAppsRank)));
-        mSrcHotseatSize = hotseatSize.x;
-        mSrcAllAppsRank = hotseatSize.y;
-
-        // Widget sizes
-        mWidgetMinSize = new HashMap<String, Point>();
-        for (String s : prefs.getStringSet(KEY_MIGRATION_WIDGET_MINSIZE,
-                Collections.<String>emptySet())) {
-            String[] parts = s.split("#");
-            mWidgetMinSize.put(parts[0], parsePoint(parts[1]));
-        }
+        mTrgX = targetSize.x;
+        mTrgY = targetSize.y;
 
         mShouldRemoveX = mTrgX < mSrcX;
         mShouldRemoveY = mTrgY < mSrcY;
+
+        // Non-used variables
+        mSrcHotseatSize = mSrcAllAppsRank = mDestHotseatSize = mDestAllAppsRank = -1;
     }
 
-    public void execute() throws Exception {
-        mEntryToRemove = new ArrayList<>();
-        mUpdateOperations = new ArrayList<>();
+    protected GridSizeMigrationTask(Context context,
+            InvariantDeviceProfile idp, HashSet<String> validPackages,
+            int srcHotseatSize, int srcAllAppsRank,
+            int destHotseatSize, int destAllAppsRank) {
+        mContext = context;
+        mIdp = idp;
+        mValidPackages = validPackages;
 
-        // Initialize list of valid packages. This contain all the packages which are already on
-        // the device and packages which are being installed. Any item which doesn't belong to
-        // this set is removed.
-        // Since the loader removes such items anyway, removing these items here doesn't cause any
-        // extra data loss and gives us more free space on the grid for better migration.
-        mValidPackages = new HashSet<>();
-        for (PackageInfo info : mContext.getPackageManager().getInstalledPackages(0)) {
-            mValidPackages.add(info.packageName);
-        }
-        mValidPackages.addAll(PackageInstallerCompat.getInstance(mContext)
-                .updateAndGetActiveSessionCache().keySet());
+        mSrcHotseatSize = srcHotseatSize;
+        mSrcAllAppsRank = srcAllAppsRank;
 
-        // Migrate hotseat
-        if (mSrcHotseatSize != mIdp.numHotseatIcons || mSrcAllAppsRank != mIdp.hotseatAllAppsRank) {
-            migrateHotseat();
-        }
+        mDestHotseatSize = destHotseatSize;
+        mDestAllAppsRank = destAllAppsRank;
 
-        if (mShouldRemoveX || mShouldRemoveY) {
-            if ((mSrcY - mTrgX) > 1 || (mSrcY - mSrcY) > 1) {
-                // TODO: support this.
-                throw new Exception("The universe is too large for migration");
-            } else {
-                migrateWorkspace();
-            }
-        }
+        // Non-used variables
+        mSrcX = mSrcY = mTrgX = mTrgY = -1;
+        mShouldRemoveX = mShouldRemoveY = false;
+    }
 
+    /**
+     * Applied all the pending DB operations
+     * @return true if any DB operation was commited.
+     */
+    private boolean applyOperations() throws Exception {
         // Update items
         if (!mUpdateOperations.isEmpty()) {
             mContext.getContentResolver().applyBatch(LauncherProvider.AUTHORITY, mUpdateOperations);
@@ -155,16 +137,7 @@
                             LauncherSettings.Favorites._ID, mEntryToRemove), null);
         }
 
-        if (!mUpdateOperations.isEmpty() || !mEntryToRemove.isEmpty()) {
-            // Make sure we haven't removed everything.
-            final Cursor c = mContext.getContentResolver().query(
-                    LauncherSettings.Favorites.CONTENT_URI, null, null, null, null);
-            boolean hasData = c.moveToNext();
-            c.close();
-            if (!hasData) {
-                throw new Exception("Removed every thing during grid resize");
-            }
-        }
+        return !mUpdateOperations.isEmpty() || !mEntryToRemove.isEmpty();
     }
 
     /**
@@ -173,11 +146,12 @@
      * entries is more than what can fit in the new hotseat, we drop the entries with least weight.
      * For weight calculation {@see #WT_SHORTCUT}, {@see #WT_APPLICATION}
      * & {@see #WT_FOLDER_FACTOR}.
+     * @return true if any DB change was made
      */
-    private void migrateHotseat() {
+    protected boolean migrateHotseat() throws Exception {
         ArrayList<DbEntry> items = loadHotseatEntries();
 
-        int requiredCount = mIdp.numHotseatIcons - 1;
+        int requiredCount = mDestHotseatSize - 1;
 
         while (items.size() > requiredCount) {
             // Pick the center item by default.
@@ -209,15 +183,18 @@
             }
 
             newScreenId++;
-            if (newScreenId == mIdp.hotseatAllAppsRank) {
+            if (newScreenId == mDestAllAppsRank) {
                 newScreenId++;
             }
         }
+
+        return applyOperations();
     }
 
-    private void migrateWorkspace() throws Exception {
-        mCarryOver = new ArrayList<>();
-
+    /**
+     * @return true if any DB change was made
+     */
+    protected boolean migrateWorkspace() throws Exception {
         ArrayList<Long> allScreens = LauncherModel.loadWorkspaceScreensDb(mContext);
         if (allScreens.isEmpty()) {
             throw new Exception("Unable to get workspace screens");
@@ -250,6 +227,7 @@
                             mContext.getContentResolver(),
                             LauncherSettings.Settings.METHOD_NEW_SCREEN_ID)
                             .getLong(LauncherSettings.Settings.EXTRA_VALUE);
+
                     allScreens.add(newScreenId);
                     for (DbEntry item : placement.finalPlacedItems) {
                         if (!mCarryOver.remove(itemMap.get(item.id))) {
@@ -264,10 +242,19 @@
 
             } while (!mCarryOver.isEmpty());
 
-
-            LauncherAppState.getInstance().getModel()
-                .updateWorkspaceScreenOrder(mContext, allScreens);
+            // Update screens
+            final Uri uri = LauncherSettings.WorkspaceScreens.CONTENT_URI;
+            mUpdateOperations.add(ContentProviderOperation.newDelete(uri).build());
+            int count = allScreens.size();
+            for (int i = 0; i < count; i++) {
+                ContentValues v = new ContentValues();
+                long screenId = allScreens.get(i);
+                v.put(LauncherSettings.WorkspaceScreens._ID, screenId);
+                v.put(LauncherSettings.WorkspaceScreens.SCREEN_RANK, i);
+                mUpdateOperations.add(ContentProviderOperation.newInsert(uri).withValues(v).build());
+            }
         }
+        return applyOperations();
     }
 
     /**
@@ -700,96 +687,96 @@
      * Loads entries for a particular screen id.
      */
     private ArrayList<DbEntry> loadWorkspaceEntries(long screen) {
-       Cursor c =  mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
-                new String[] {
-                    Favorites._ID,                  // 0
-                    Favorites.ITEM_TYPE,            // 1
-                    Favorites.CELLX,                // 2
-                    Favorites.CELLY,                // 3
-                    Favorites.SPANX,                // 4
-                    Favorites.SPANY,                // 5
-                    Favorites.INTENT,               // 6
-                    Favorites.APPWIDGET_PROVIDER},  // 7
+        Cursor c =  mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
+                new String[]{
+                        Favorites._ID,                  // 0
+                        Favorites.ITEM_TYPE,            // 1
+                        Favorites.CELLX,                // 2
+                        Favorites.CELLY,                // 3
+                        Favorites.SPANX,                // 4
+                        Favorites.SPANY,                // 5
+                        Favorites.INTENT,               // 6
+                        Favorites.APPWIDGET_PROVIDER},  // 7
                 Favorites.CONTAINER + " = " + Favorites.CONTAINER_DESKTOP
-                    + " AND " + Favorites.SCREEN + " = " + screen, null, null, null);
+                        + " AND " + Favorites.SCREEN + " = " + screen, null, null, null);
 
-       final int indexId = c.getColumnIndexOrThrow(Favorites._ID);
-       final int indexItemType = c.getColumnIndexOrThrow(Favorites.ITEM_TYPE);
-       final int indexCellX = c.getColumnIndexOrThrow(Favorites.CELLX);
-       final int indexCellY = c.getColumnIndexOrThrow(Favorites.CELLY);
-       final int indexSpanX = c.getColumnIndexOrThrow(Favorites.SPANX);
-       final int indexSpanY = c.getColumnIndexOrThrow(Favorites.SPANY);
-       final int indexIntent = c.getColumnIndexOrThrow(Favorites.INTENT);
-       final int indexAppWidgetProvider = c.getColumnIndexOrThrow(Favorites.APPWIDGET_PROVIDER);
+        final int indexId = c.getColumnIndexOrThrow(Favorites._ID);
+        final int indexItemType = c.getColumnIndexOrThrow(Favorites.ITEM_TYPE);
+        final int indexCellX = c.getColumnIndexOrThrow(Favorites.CELLX);
+        final int indexCellY = c.getColumnIndexOrThrow(Favorites.CELLY);
+        final int indexSpanX = c.getColumnIndexOrThrow(Favorites.SPANX);
+        final int indexSpanY = c.getColumnIndexOrThrow(Favorites.SPANY);
+        final int indexIntent = c.getColumnIndexOrThrow(Favorites.INTENT);
+        final int indexAppWidgetProvider = c.getColumnIndexOrThrow(Favorites.APPWIDGET_PROVIDER);
 
-       ArrayList<DbEntry> entries = new ArrayList<>();
-       while (c.moveToNext()) {
-           DbEntry entry = new DbEntry();
-           entry.id = c.getLong(indexId);
-           entry.itemType = c.getInt(indexItemType);
-           entry.cellX = c.getInt(indexCellX);
-           entry.cellY = c.getInt(indexCellY);
-           entry.spanX = c.getInt(indexSpanX);
-           entry.spanY = c.getInt(indexSpanY);
-           entry.screenId = screen;
+        ArrayList<DbEntry> entries = new ArrayList<>();
+        while (c.moveToNext()) {
+            DbEntry entry = new DbEntry();
+            entry.id = c.getLong(indexId);
+            entry.itemType = c.getInt(indexItemType);
+            entry.cellX = c.getInt(indexCellX);
+            entry.cellY = c.getInt(indexCellY);
+            entry.spanX = c.getInt(indexSpanX);
+            entry.spanY = c.getInt(indexSpanY);
+            entry.screenId = screen;
 
-           try {
-               // calculate weight
-               switch (entry.itemType) {
-                   case Favorites.ITEM_TYPE_SHORTCUT:
-                   case Favorites.ITEM_TYPE_APPLICATION: {
-                       verifyIntent(c.getString(indexIntent));
-                       entry.weight = entry.itemType == Favorites.ITEM_TYPE_SHORTCUT
-                           ? WT_SHORTCUT : WT_APPLICATION;
-                       break;
-                   }
-                   case Favorites.ITEM_TYPE_APPWIDGET: {
-                       String provider = c.getString(indexAppWidgetProvider);
-                       ComponentName cn = ComponentName.unflattenFromString(provider);
-                       verifyPackage(cn.getPackageName());
-                       entry.weight = Math.max(WT_WIDGET_MIN, WT_WIDGET_FACTOR
-                               * entry.spanX * entry.spanY);
+            try {
+                // calculate weight
+                switch (entry.itemType) {
+                    case Favorites.ITEM_TYPE_SHORTCUT:
+                    case Favorites.ITEM_TYPE_APPLICATION: {
+                        verifyIntent(c.getString(indexIntent));
+                        entry.weight = entry.itemType == Favorites.ITEM_TYPE_SHORTCUT
+                            ? WT_SHORTCUT : WT_APPLICATION;
+                        break;
+                    }
+                    case Favorites.ITEM_TYPE_APPWIDGET: {
+                        String provider = c.getString(indexAppWidgetProvider);
+                        ComponentName cn = ComponentName.unflattenFromString(provider);
+                        verifyPackage(cn.getPackageName());
+                        entry.weight = Math.max(WT_WIDGET_MIN, WT_WIDGET_FACTOR
+                                * entry.spanX * entry.spanY);
 
-                       // Migration happens for current user only.
-                       LauncherAppWidgetProviderInfo pInfo = LauncherModel.getProviderInfo(
-                               mContext, cn, UserHandleCompat.myUserHandle());
-                       Point spans = pInfo == null ?
-                               mWidgetMinSize.get(provider) : pInfo.getMinSpans(mIdp, mContext);
-                       if (spans != null) {
-                           entry.minSpanX = spans.x > 0 ? spans.x : entry.spanX;
-                           entry.minSpanY = spans.y > 0 ? spans.y : entry.spanY;
-                       } else {
-                           // Assume that the widget be resized down to 2x2
-                           entry.minSpanX = entry.minSpanY = 2;
-                       }
+                        // Migration happens for current user only.
+                        LauncherAppWidgetProviderInfo pInfo = LauncherModel.getProviderInfo(
+                                mContext, cn, UserHandleCompat.myUserHandle());
+                        Point spans = pInfo == null ?
+                                mWidgetMinSize.get(provider) : pInfo.getMinSpans(mIdp, mContext);
+                        if (spans != null) {
+                            entry.minSpanX = spans.x > 0 ? spans.x : entry.spanX;
+                            entry.minSpanY = spans.y > 0 ? spans.y : entry.spanY;
+                        } else {
+                            // Assume that the widget be resized down to 2x2
+                            entry.minSpanX = entry.minSpanY = 2;
+                        }
 
-                       if (entry.minSpanX > mTrgX || entry.minSpanY > mTrgY) {
-                           throw new Exception("Widget can't be resized down to fit the grid");
-                       }
-                       break;
-                   }
-                   case Favorites.ITEM_TYPE_FOLDER: {
-                       int total = getFolderItemsCount(entry.id);
-                       if (total == 0) {
-                           throw new Exception("Folder is empty");
-                       }
-                       entry.weight = WT_FOLDER_FACTOR * total;
-                       break;
-                   }
-                   default:
-                       throw new Exception("Invalid item type");
-               }
-           } catch (Exception e) {
-               if (DEBUG) {
-                   Log.d(TAG, "Removing item " + entry.id, e);
-               }
-               mEntryToRemove.add(entry.id);
-               continue;
-           }
-           entries.add(entry);
-       }
-       c.close();
-       return entries;
+                        if (entry.minSpanX > mTrgX || entry.minSpanY > mTrgY) {
+                            throw new Exception("Widget can't be resized down to fit the grid");
+                        }
+                        break;
+                    }
+                    case Favorites.ITEM_TYPE_FOLDER: {
+                        int total = getFolderItemsCount(entry.id);
+                        if (total == 0) {
+                            throw new Exception("Folder is empty");
+                        }
+                        entry.weight = WT_FOLDER_FACTOR * total;
+                        break;
+                    }
+                    default:
+                        throw new Exception("Invalid item type");
+                }
+            } catch (Exception e) {
+                if (DEBUG) {
+                    Log.d(TAG, "Removing item " + entry.id, e);
+                }
+                mEntryToRemove.add(entry.id);
+                continue;
+            }
+            entries.add(entry);
+        }
+        c.close();
+        return entries;
     }
 
     /**
@@ -797,7 +784,7 @@
      */
     private int getFolderItemsCount(long folderId) {
         Cursor c =  mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
-                new String[] {Favorites._ID, Favorites.INTENT},
+                new String[]{Favorites._ID, Favorites.INTENT},
                 Favorites.CONTAINER + " = " + folderId, null, null, null);
 
         int total = 0;
@@ -897,42 +884,147 @@
         return new Point(Integer.parseInt(split[0]), Integer.parseInt(split[1]));
     }
 
-    public static void markForMigration(Context context, int srcX, int srcY,
-            HashSet<String> widgets) {
-        Utilities.getPrefs(context).edit()
-                .putString(KEY_MIGRATION_SRC_WORKSPACE_SIZE, getPointString(srcX, srcY))
-                .putStringSet(KEY_MIGRATION_WIDGET_MINSIZE, widgets)
-                .apply();
-    }
-
-    public static boolean shouldRunTask(Context context) {
-        SharedPreferences prefs = Utilities.getPrefs(context);
-        InvariantDeviceProfile idp = LauncherAppState.getInstance().getInvariantDeviceProfile();
-
-        // Run task if workspace or hotseat size has changed.
-        return !getPointString(idp.numColumns, idp.numRows).equals(
-                    prefs.getString(KEY_MIGRATION_SRC_WORKSPACE_SIZE, ""))
-                || !getPointString(idp.numHotseatIcons, idp.hotseatAllAppsRank).equals(
-                        prefs.getString(KEY_MIGRATION_SRC_HOTSEAT_SIZE, ""));
-    }
-
-    public static void clearFlags(Context context) {
-        Utilities.getPrefs(context).edit().remove(KEY_MIGRATION_WIDGET_MINSIZE).commit();
-    }
-
-    public static void saveCurrentConfig(Context context) {
-        InvariantDeviceProfile idp = LauncherAppState.getInstance().getInvariantDeviceProfile();
-        Utilities.getPrefs(context).edit()
-                .putString(KEY_MIGRATION_SRC_WORKSPACE_SIZE,
-                        getPointString(idp.numColumns, idp.numRows))
-                .putString(KEY_MIGRATION_SRC_HOTSEAT_SIZE,
-                        getPointString(idp.numHotseatIcons, idp.hotseatAllAppsRank))
-                .remove(KEY_MIGRATION_WIDGET_MINSIZE)
-                .commit();
-    }
-
     private static String getPointString(int x, int y) {
         return String.format(Locale.ENGLISH, "%d,%d", x, y);
     }
 
+    public static void markForMigration(
+            Context context, HashSet<String> widgets, BackupProtos.DeviceProfieData srcProfile) {
+        Utilities.getPrefs(context).edit()
+                .putString(KEY_MIGRATION_SRC_WORKSPACE_SIZE,
+                        getPointString((int) srcProfile.desktopCols, (int) srcProfile.desktopRows))
+                .putString(KEY_MIGRATION_SRC_HOTSEAT_SIZE,
+                        getPointString((int) srcProfile.hotseatCount, srcProfile.allappsRank))
+                .putStringSet(KEY_MIGRATION_WIDGET_MINSIZE, widgets)
+                .apply();
+    }
+
+    /**
+     * Migrates the workspace and hotseat in case their sizes changed.
+     * @return false if the migration failed.
+     */
+    public static boolean migrateGridIfNeeded(Context context) {
+        SharedPreferences prefs = Utilities.getPrefs(context);
+        InvariantDeviceProfile idp = LauncherAppState.getInstance().getInvariantDeviceProfile();
+
+        String gridSizeString = getPointString(idp.numColumns, idp.numRows);
+        String hotseatSizeString = getPointString(idp.numHotseatIcons, idp.hotseatAllAppsRank);
+
+        if (gridSizeString.equals(prefs.getString(KEY_MIGRATION_SRC_WORKSPACE_SIZE, "")) &&
+                hotseatSizeString.equals(prefs.getString(KEY_MIGRATION_SRC_HOTSEAT_SIZE, ""))) {
+            // Skip if workspace and hotseat sizes have not changed.
+            return true;
+        }
+
+        long migrationStartTime = System.currentTimeMillis();
+        try {
+            boolean dbChanged = false;
+
+            // Initialize list of valid packages. This contain all the packages which are already on
+            // the device and packages which are being installed. Any item which doesn't belong to
+            // this set is removed.
+            // Since the loader removes such items anyway, removing these items here doesn't cause
+            // any extra data loss and gives us more free space on the grid for better migration.
+            HashSet validPackages = new HashSet<>();
+            for (PackageInfo info : context.getPackageManager().getInstalledPackages(0)) {
+                validPackages.add(info.packageName);
+            }
+            validPackages.addAll(PackageInstallerCompat.getInstance(context)
+                    .updateAndGetActiveSessionCache().keySet());
+
+            // Hotseat
+            Point srcHotseatSize = parsePoint(prefs.getString(
+                    KEY_MIGRATION_SRC_HOTSEAT_SIZE, hotseatSizeString));
+            if (srcHotseatSize.x != idp.numHotseatIcons ||
+                    srcHotseatSize.y != idp.hotseatAllAppsRank) {
+                // Migrate hotseat.
+
+                dbChanged = new GridSizeMigrationTask(context,
+                        LauncherAppState.getInstance().getInvariantDeviceProfile(),
+                        validPackages,
+                        srcHotseatSize.x, srcHotseatSize.y,
+                        idp.numHotseatIcons, idp.hotseatAllAppsRank).migrateHotseat();
+            }
+
+            // Grid size
+            Point targetSize = new Point(idp.numColumns, idp.numRows);
+            Point sourceSize = parsePoint(prefs.getString(
+                    KEY_MIGRATION_SRC_WORKSPACE_SIZE, gridSizeString));
+
+            if (!targetSize.equals(sourceSize)) {
+
+                // The following list defines all possible grid sizes (and intermediate steps
+                // during migration). Note that at each step, dx <= 1 && dy <= 1. Any grid size
+                // which is not in this list is not migrated.
+                ArrayList<Point> gridSizeSteps = new ArrayList<>();
+                gridSizeSteps.add(new Point(2, 3));
+                gridSizeSteps.add(new Point(3, 3));
+                gridSizeSteps.add(new Point(3, 4));
+                gridSizeSteps.add(new Point(4, 4));
+                gridSizeSteps.add(new Point(5, 5));
+                gridSizeSteps.add(new Point(5, 6));
+                gridSizeSteps.add(new Point(6, 6));
+                gridSizeSteps.add(new Point(7, 7));
+
+                int sourceSizeIndex = gridSizeSteps.indexOf(sourceSize);
+                int targetSizeIndex = gridSizeSteps.indexOf(targetSize);
+
+                if (sourceSizeIndex <= -1 || targetSizeIndex <= -1) {
+                    throw new Exception("Unable to migrate grid size from " + sourceSize
+                            + " to " + targetSize);
+                }
+
+                // Min widget sizes
+                HashMap<String, Point> widgetMinSize = new HashMap<>();
+                for (String s : Utilities.getPrefs(context).getStringSet(KEY_MIGRATION_WIDGET_MINSIZE,
+                        Collections.<String>emptySet())) {
+                    String[] parts = s.split("#");
+                    widgetMinSize.put(parts[0], parsePoint(parts[1]));
+                }
+
+                // Migrate the workspace grid, step by step.
+                while (targetSizeIndex < sourceSizeIndex ) {
+                    // We only need to migrate the grid if source size is greater
+                    // than the target size.
+                    Point stepTargetSize = gridSizeSteps.get(sourceSizeIndex - 1);
+                    Point stepSourceSize = gridSizeSteps.get(sourceSizeIndex);
+
+                    if (new GridSizeMigrationTask(context,
+                            LauncherAppState.getInstance().getInvariantDeviceProfile(),
+                            validPackages, widgetMinSize,
+                            stepSourceSize, stepTargetSize).migrateWorkspace()) {
+                        dbChanged = true;
+                    }
+                    sourceSizeIndex--;
+                }
+            }
+
+            if (dbChanged) {
+                // Make sure we haven't removed everything.
+                final Cursor c = context.getContentResolver().query(
+                        LauncherSettings.Favorites.CONTENT_URI, null, null, null, null);
+                boolean hasData = c.moveToNext();
+                c.close();
+                if (!hasData) {
+                    throw new Exception("Removed every thing during grid resize");
+                }
+            }
+
+            return true;
+        } catch (Exception e) {
+            Log.e(TAG, "Error during grid migration", e);
+
+            return false;
+        } finally {
+            Log.v(TAG, "Workspace migration completed in "
+                    + (System.currentTimeMillis() - migrationStartTime));
+
+            // Save current configuration, so that the migration does not run again.
+            prefs.edit()
+                    .putString(KEY_MIGRATION_SRC_WORKSPACE_SIZE, gridSizeString)
+                    .putString(KEY_MIGRATION_SRC_HOTSEAT_SIZE, hotseatSizeString)
+                    .remove(KEY_MIGRATION_WIDGET_MINSIZE)
+                    .apply();
+        }
+    }
 }
diff --git a/tests/src/com/android/launcher3/model/GridSizeMigrationTaskTest.java b/tests/src/com/android/launcher3/model/GridSizeMigrationTaskTest.java
new file mode 100644
index 0000000..46dac0a
--- /dev/null
+++ b/tests/src/com/android/launcher3/model/GridSizeMigrationTaskTest.java
@@ -0,0 +1,321 @@
+package com.android.launcher3.model;
+
+import android.content.ContentValues;
+import android.content.Intent;
+import android.database.Cursor;
+import android.graphics.Point;
+import android.test.ProviderTestCase2;
+
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherModel;
+import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.config.ProviderConfig;
+import com.android.launcher3.util.TestLauncherProvider;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+
+/**
+ * Unit tests for {@link GridSizeMigrationTask}
+ */
+public class GridSizeMigrationTaskTest extends ProviderTestCase2<TestLauncherProvider> {
+
+    private static final long DESKTOP = LauncherSettings.Favorites.CONTAINER_DESKTOP;
+    private static final long HOTSEAT = LauncherSettings.Favorites.CONTAINER_HOTSEAT;
+
+    private static final int APPLICATION = LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
+    private static final int SHORTCUT = LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT;
+
+    private static final String TEST_PACKAGE = "com.android.launcher3.validpackage";
+    private static final String VALID_INTENT =
+            new Intent(Intent.ACTION_MAIN).setPackage(TEST_PACKAGE).toUri(0);
+
+    private HashSet<String> mValidPackages;
+    private InvariantDeviceProfile mIdp;
+
+    public GridSizeMigrationTaskTest() {
+        super(TestLauncherProvider.class, ProviderConfig.AUTHORITY);
+    }
+
+    @Override
+    protected void setUp() throws Exception {
+        super.setUp();
+        mValidPackages = new HashSet<>();
+        mValidPackages.add(TEST_PACKAGE);
+
+        mIdp = new InvariantDeviceProfile();
+    }
+
+    public void testHotseatMigration_apps_dropped() throws Exception {
+        long[] hotseatItems = {
+                addItem(APPLICATION, 0, HOTSEAT, 0, 0),
+                addItem(SHORTCUT, 1, HOTSEAT, 0, 0),
+                -1,
+                addItem(SHORTCUT, 3, HOTSEAT, 0, 0),
+                addItem(APPLICATION, 4, HOTSEAT, 0, 0),
+        };
+
+        new GridSizeMigrationTask(getMockContext(), mIdp, mValidPackages, 5, 2, 3, 1)
+                .migrateHotseat();
+        // First & last items are dropped as they have the least weight.
+        verifyHotseat(hotseatItems[1], -1, hotseatItems[3]);
+    }
+
+    public void testHotseatMigration_shortcuts_dropped() throws Exception {
+        long[] hotseatItems = {
+                addItem(APPLICATION, 0, HOTSEAT, 0, 0),
+                addItem(30, 1, HOTSEAT, 0, 0),
+                -1,
+                addItem(SHORTCUT, 3, HOTSEAT, 0, 0),
+                addItem(10, 4, HOTSEAT, 0, 0),
+        };
+
+        new GridSizeMigrationTask(getMockContext(), mIdp, mValidPackages, 5, 2, 3, 1)
+                .migrateHotseat();
+        // First & third items are dropped as they have the least weight.
+        verifyHotseat(hotseatItems[1], -1, hotseatItems[4]);
+    }
+
+    private void verifyHotseat(long... sortedIds) {
+        int screenId = 0;
+        int total = 0;
+
+        for (long id : sortedIds) {
+            Cursor c = getMockContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
+                    new String[]{LauncherSettings.Favorites._ID},
+                    "container=-101 and screen=" + screenId, null, null, null);
+
+            if (id == -1) {
+                assertEquals(0, c.getCount());
+            } else {
+                assertEquals(1, c.getCount());
+                c.moveToNext();
+                assertEquals(id, c.getLong(0));
+                total ++;
+            }
+            c.close();
+
+            screenId++;
+        }
+
+        // Verify that not other entry exist in the DB.
+        Cursor c = getMockContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
+                new String[]{LauncherSettings.Favorites._ID},
+                "container=-101", null, null, null);
+        assertEquals(total, c.getCount());
+        c.close();
+    }
+
+    public void testWorkspace_empty_row_column_removed() throws Exception {
+        long[][][] ids = createGrid(new int[][][]{{
+                {  0,  0, -1,  1},
+                {  3,  1, -1,  4},
+                { -1, -1, -1, -1},
+                {  5,  2, -1,  6},
+        }});
+
+        new GridSizeMigrationTask(getMockContext(), mIdp, mValidPackages, new HashMap<String, Point>(),
+                new Point(4, 4), new Point(3, 3)).migrateWorkspace();
+
+        // Column 2 and row 2 got removed.
+        verifyWorkspace(new long[][][] {{
+                {ids[0][0][0], ids[0][0][1], ids[0][0][3]},
+                {ids[0][1][0], ids[0][1][1], ids[0][1][3]},
+                {ids[0][3][0], ids[0][3][1], ids[0][3][3]},
+        }});
+    }
+
+    public void testWorkspace_new_screen_created() throws Exception {
+        long[][][] ids = createGrid(new int[][][]{{
+                {  0,  0,  0,  1},
+                {  3,  1,  0,  4},
+                { -1, -1, -1, -1},
+                {  5,  2, -1,  6},
+        }});
+
+        new GridSizeMigrationTask(getMockContext(), mIdp, mValidPackages, new HashMap<String, Point>(),
+                new Point(4, 4), new Point(3, 3)).migrateWorkspace();
+
+        // Items in the second column get moved to new screen
+        verifyWorkspace(new long[][][] {{
+                {ids[0][0][0], ids[0][0][1], ids[0][0][3]},
+                {ids[0][1][0], ids[0][1][1], ids[0][1][3]},
+                {ids[0][3][0], ids[0][3][1], ids[0][3][3]},
+        }, {
+                {ids[0][0][2], ids[0][1][2], -1},
+        }});
+    }
+
+    public void testWorkspace_items_merged_in_next_screen() throws Exception {
+        long[][][] ids = createGrid(new int[][][]{{
+                {  0,  0,  0,  1},
+                {  3,  1,  0,  4},
+                { -1, -1, -1, -1},
+                {  5,  2, -1,  6},
+        },{
+                {  0,  0, -1,  1},
+                {  3,  1, -1,  4},
+        }});
+
+        new GridSizeMigrationTask(getMockContext(), mIdp, mValidPackages, new HashMap<String, Point>(),
+                new Point(4, 4), new Point(3, 3)).migrateWorkspace();
+
+        // Items in the second column of the first screen should get placed on the 3rd
+        // row of the second screen
+        verifyWorkspace(new long[][][] {{
+                {ids[0][0][0], ids[0][0][1], ids[0][0][3]},
+                {ids[0][1][0], ids[0][1][1], ids[0][1][3]},
+                {ids[0][3][0], ids[0][3][1], ids[0][3][3]},
+        }, {
+                {ids[1][0][0], ids[1][0][1], ids[1][0][3]},
+                {ids[1][1][0], ids[1][1][1], ids[1][1][3]},
+                {ids[0][0][2], ids[0][1][2], -1},
+        }});
+    }
+
+    public void testWorkspace_items_not_merged_in_next_screen() throws Exception {
+        // First screen has 2 items that need to be moved, but second screen has only one
+        // empty space after migration (top-left corner)
+        long[][][] ids = createGrid(new int[][][]{{
+                {  0,  0,  0,  1},
+                {  3,  1,  0,  4},
+                { -1, -1, -1, -1},
+                {  5,  2, -1,  6},
+        },{
+                { -1,  0, -1,  1},
+                {  3,  1, -1,  4},
+                { -1, -1, -1, -1},
+                {  5,  2, -1,  6},
+        }});
+
+        new GridSizeMigrationTask(getMockContext(), mIdp, mValidPackages, new HashMap<String, Point>(),
+                new Point(4, 4), new Point(3, 3)).migrateWorkspace();
+
+        // Items in the second column of the first screen should get placed on a new screen.
+        verifyWorkspace(new long[][][] {{
+                {ids[0][0][0], ids[0][0][1], ids[0][0][3]},
+                {ids[0][1][0], ids[0][1][1], ids[0][1][3]},
+                {ids[0][3][0], ids[0][3][1], ids[0][3][3]},
+        }, {
+                {          -1, ids[1][0][1], ids[1][0][3]},
+                {ids[1][1][0], ids[1][1][1], ids[1][1][3]},
+                {ids[1][3][0], ids[1][3][1], ids[1][3][3]},
+        }, {
+                {ids[0][0][2], ids[0][1][2], -1},
+        }});
+    }
+
+    /**
+     * Initializes the DB with dummy elements to represent the provided grid structure.
+     * @param typeArray A 3d array of item types. {@see #addItem(int, long, long, int, int)} for
+     *                  type definitions. The first dimension represents the screens and the next
+     *                  two represent the workspace grid.
+     * @return the same grid representation where each entry is the corresponding item id.
+     */
+    private long[][][] createGrid(int[][][] typeArray) throws Exception {
+        long[][][] ids = new long[typeArray.length][][];
+
+        for (int i = 0; i < typeArray.length; i++) {
+            // Add screen to DB
+            long screenId = LauncherSettings.Settings.call(getMockContentResolver(),
+                    LauncherSettings.Settings.METHOD_NEW_SCREEN_ID)
+                    .getLong(LauncherSettings.Settings.EXTRA_VALUE);
+
+            ContentValues v = new ContentValues();
+            v.put(LauncherSettings.WorkspaceScreens._ID, screenId);
+            v.put(LauncherSettings.WorkspaceScreens.SCREEN_RANK, i);
+            getMockContentResolver().insert(LauncherSettings.WorkspaceScreens.CONTENT_URI, v);
+
+            ids[i] = new long[typeArray[i].length][];
+            for (int y = 0; y < typeArray[i].length; y++) {
+                ids[i][y] = new long[typeArray[i][y].length];
+                for (int x = 0; x < typeArray[i][y].length; x++) {
+                    if (typeArray[i][y][x] < 0) {
+                        // Empty cell
+                        ids[i][y][x] = -1;
+                    } else {
+                        ids[i][y][x] = addItem(typeArray[i][y][x], screenId, DESKTOP, x, y);
+                    }
+                }
+            }
+        }
+        return ids;
+    }
+
+    /**
+     * Verifies that the workspace items are arranged in the provided order.
+     * @param ids A 3d array where the first dimension represents the screen, and the rest two
+     *            represent the workspace grid.
+     */
+    private void verifyWorkspace(long[][][] ids) {
+        ArrayList<Long> allScreens = LauncherModel.loadWorkspaceScreensDb(getMockContext());
+        assertEquals(ids.length, allScreens.size());
+        int total = 0;
+
+        for (int i = 0; i < ids.length; i++) {
+            long screenId = allScreens.get(i);
+            for (int y = 0; y < ids[i].length; y++) {
+                for (int x = 0; x < ids[i][y].length; x++) {
+                    long id = ids[i][y][x];
+
+                    Cursor c = getMockContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
+                            new String[]{LauncherSettings.Favorites._ID},
+                            "container=-100 and screen=" + screenId +
+                                    " and cellX=" + x + " and cellY=" + y, null, null, null);
+                    if (id == -1) {
+                        assertEquals(0, c.getCount());
+                    } else {
+                        assertEquals(1, c.getCount());
+                        c.moveToNext();
+                        assertEquals(id, c.getLong(0));
+                        total++;
+                    }
+                    c.close();
+                }
+            }
+        }
+
+        // Verify that not other entry exist in the DB.
+        Cursor c = getMockContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
+                new String[]{LauncherSettings.Favorites._ID},
+                "container=-100", null, null, null);
+        assertEquals(total, c.getCount());
+        c.close();
+    }
+
+    /**
+     * Adds a dummy item in the DB.
+     * @param type {@link #APPLICATION} or {@link #SHORTCUT} or >= 2 for
+     *             folder (where the type represents the number of items in the folder).
+     */
+    private long addItem(int type, long screen, long container, int x, int y) throws Exception {
+        long id = LauncherSettings.Settings.call(getMockContentResolver(),
+                LauncherSettings.Settings.METHOD_NEW_ITEM_ID)
+                .getLong(LauncherSettings.Settings.EXTRA_VALUE);
+
+        ContentValues values = new ContentValues();
+        values.put(LauncherSettings.Favorites._ID, id);
+        values.put(LauncherSettings.Favorites.CONTAINER, container);
+        values.put(LauncherSettings.Favorites.SCREEN, screen);
+        values.put(LauncherSettings.Favorites.CELLX, x);
+        values.put(LauncherSettings.Favorites.CELLY, y);
+        values.put(LauncherSettings.Favorites.SPANX, 1);
+        values.put(LauncherSettings.Favorites.SPANY, 1);
+
+        if (type == APPLICATION || type == SHORTCUT) {
+            values.put(LauncherSettings.Favorites.ITEM_TYPE, type);
+            values.put(LauncherSettings.Favorites.INTENT, VALID_INTENT);
+        } else {
+            values.put(LauncherSettings.Favorites.ITEM_TYPE,
+                    LauncherSettings.Favorites.ITEM_TYPE_FOLDER);
+            // Add folder items.
+            for (int i = 0; i < type; i++) {
+                addItem(APPLICATION, 0, id, 0, 0);
+            }
+        }
+
+        getMockContentResolver().insert(LauncherSettings.Favorites.CONTENT_URI, values);
+        return id;
+    }
+}
diff --git a/tests/src/com/android/launcher3/util/TestLauncherProvider.java b/tests/src/com/android/launcher3/util/TestLauncherProvider.java
new file mode 100644
index 0000000..aef3240
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/TestLauncherProvider.java
@@ -0,0 +1,40 @@
+package com.android.launcher3.util;
+
+import android.content.Context;
+
+import com.android.launcher3.LauncherProvider;
+
+/**
+ * An extension of LauncherProvider backed up by in-memory database.
+ */
+public class TestLauncherProvider extends LauncherProvider {
+
+    @Override
+    public boolean onCreate() {
+        return true;
+    }
+
+    @Override
+    protected synchronized void createDbIfNotExists() {
+        if (mOpenHelper == null) {
+            mOpenHelper = new MyDatabaseHelper(getContext(), this);
+        }
+    }
+
+    @Override
+    protected void notifyListeners() { }
+
+    private static class MyDatabaseHelper extends DatabaseHelper {
+        public MyDatabaseHelper(Context context, LauncherProvider provider) {
+            super(context, provider, null);
+        }
+
+        @Override
+        protected long getDefaultUserSerial() {
+            return 0;
+        }
+
+        @Override
+        protected void onEmptyDbCreated() { }
+    }
+}
\ No newline at end of file