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
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