Adding support for migrating the grid between any two valid screens sizes.

The grid is migrated in steps where each step consists of at max one column change and at max one row change.
Adding some unit tests for GridMigrationLogic

Bug: 25958224
Change-Id: Ie54e872ea0925cc4c463edbba0a7201d62b373a0

9 files changed, 701 insertions, 239 deletions

diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index dce1ab887..d60132270 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -83,7 +83,7 @@ public class InvariantDeviceProfile {
     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 bf9c66822..b192ba3cc 100644
--- a/src/com/android/launcher3/LauncherBackupAgentHelper.java
+++ b/src/com/android/launcher3/LauncherBackupAgentHelper.java
@@ -101,12 +101,9 @@ public class LauncherBackupAgentHelper extends BackupAgentHelper {
                         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 5f58e284a..05d729e78 100644
--- a/src/com/android/launcher3/LauncherBackupHelper.java
+++ b/src/com/android/launcher3/LauncherBackupHelper.java
@@ -315,14 +315,13 @@ public class LauncherBackupHelper implements BackupHelper {
             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 0eb1a90b0..92ef3eae7 100644
--- a/src/com/android/launcher3/LauncherModel.java
+++ b/src/com/android/launcher3/LauncherModel.java
@@ -1651,25 +1651,10 @@ public class LauncherModel extends BroadcastReceiver
             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 3fc0b948c..ac9b32168 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -77,7 +77,7 @@ public class LauncherProvider extends ContentProvider {
 
     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 @@ public class LauncherProvider extends ContentProvider {
         }
     }
 
-    private synchronized void createDbIfNotExists() {
+    /**
+     * Overridden in tests
+     */
+    protected synchronized void createDbIfNotExists() {
         if (mOpenHelper == null) {
             mOpenHelper = new DatabaseHelper(getContext(), this);
         }
@@ -364,7 +367,10 @@ public class LauncherProvider extends ContentProvider {
         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 @@ public class LauncherProvider extends ContentProvider {
         });
     }
 
-    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 @@ public class LauncherProvider extends ContentProvider {
             }
         }
 
+        /**
+         * 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 @@ public class LauncherProvider extends ContentProvider {
 
             // 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);
         }
 
-        private void addFavoritesTable(SQLiteDatabase db, boolean optional) {
-            UserManagerCompat userManager = UserManagerCompat.getInstance(mContext);
-            long userSerialNumber = userManager.getSerialNumberForUser(
+        protected long getDefaultUserSerial() {
+            return UserManagerCompat.getInstance(mContext).getSerialNumberForUser(
                     UserHandleCompat.myUserHandle());
-            String ifNotExists = optional ? " IF NOT EXISTS " : "";
+        }
 
+        private void addFavoritesTable(SQLiteDatabase db, boolean optional) {
+            String ifNotExists = optional ? " IF NOT EXISTS " : "";
             db.execSQL("CREATE TABLE " + ifNotExists + TABLE_FAVORITES + " (" +
                     "_id INTEGER PRIMARY KEY," +
                     "title TEXT," +
@@ -599,7 +631,7 @@ public class LauncherProvider extends ContentProvider {
                     "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 @@ public class LauncherProvider extends ContentProvider {
             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 9ee6a2199..55a53785f 100644
--- a/src/com/android/launcher3/LauncherSettings.java
+++ b/src/com/android/launcher3/LauncherSettings.java
@@ -115,7 +115,7 @@ public class LauncherSettings {
         /**
          * 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 08c3dc0bb..19ec3ed64 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationTask.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationTask.java
@@ -9,6 +9,7 @@ import android.content.SharedPreferences;
 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.LauncherProvider;
 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 @@ public class GridSizeMigrationTask {
     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 @@ public class GridSizeMigrationTask {
 
     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 @@ public class GridSizeMigrationTask {
                             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 @@ public class GridSizeMigrationTask {
      * 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 @@ public class GridSizeMigrationTask {
             }
 
             newScreenId++;
-            if (newScreenId == mIdp.hotseatAllAppsRank) {
+            if (newScreenId == mDestAllAppsRank) {
                 newScreenId++;
             }
         }
-    }
 
-    private void migrateWorkspace() throws Exception {
-        mCarryOver = new ArrayList<>();
+        return applyOperations();
+    }
 
+    /**
+     * @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 @@ public class GridSizeMigrationTask {
                             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 @@ public class GridSizeMigrationTask {
 
             } 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 @@ public class GridSizeMigrationTask {
      * 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);
-
-       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;
-
-           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;
-                       }
-
-                       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;
+                        + " 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);
+
+        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);
+
+                        // 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;
     }
 
     /**
@@ -797,7 +784,7 @@ public class GridSizeMigrationTask {
      */
     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 @@ public class GridSizeMigrationTask {
         return new Point(Integer.parseInt(split[0]), Integer.parseInt(split[1]));
     }
 
-    public static void markForMigration(Context context, int srcX, int srcY,
-            HashSet<String> widgets) {
+    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(srcX, srcY))
+                .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();
     }
 
-    public static boolean shouldRunTask(Context context) {
+    /**
+     * 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();
 
-        // 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, ""));
-    }
+        String gridSizeString = getPointString(idp.numColumns, idp.numRows);
+        String hotseatSizeString = getPointString(idp.numHotseatIcons, idp.hotseatAllAppsRank);
 
-    public static void clearFlags(Context context) {
-        Utilities.getPrefs(context).edit().remove(KEY_MIGRATION_WIDGET_MINSIZE).commit();
-    }
+        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;
+        }
 
-    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();
-    }
+        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();
+            }
 
-    private static String getPointString(int x, int y) {
-        return String.format(Locale.ENGLISH, "%d,%d", x, y);
-    }
+            // 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 000000000..46dac0aab
--- /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 000000000..aef3240ca
--- /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