diff --git a/src/com/android/launcher3/model/GridSizeMigrationTask.java b/src/com/android/launcher3/model/GridSizeMigrationTask.java
index 8e8e551..9db79f0 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationTask.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationTask.java
@@ -150,7 +150,7 @@
     protected boolean migrateHotseat() throws Exception {
         ArrayList<DbEntry> items = loadHotseatEntries();
 
-        int requiredCount = mDestHotseatSize - 1;
+        int requiredCount = FeatureFlags.NO_ALL_APPS_ICON ? mDestHotseatSize : mDestHotseatSize - 1;
 
         while (items.size() > requiredCount) {
             // Pick the center item by default.
diff --git a/src/com/android/launcher3/util/CachedPackageTracker.java b/src/com/android/launcher3/util/CachedPackageTracker.java
index d55d573..293714e 100644
--- a/src/com/android/launcher3/util/CachedPackageTracker.java
+++ b/src/com/android/launcher3/util/CachedPackageTracker.java
@@ -170,7 +170,7 @@
      */
     protected abstract void onLauncherPackageRemoved(String packageName, UserHandleCompat user);
 
-    protected static class LauncherActivityInstallInfo
+    public static class LauncherActivityInstallInfo
             implements Comparable<LauncherActivityInstallInfo> {
         public final LauncherActivityInfoCompat info;
         public final long installTime;
diff --git a/tests/src/com/android/launcher3/model/GridSizeMigrationTaskTest.java b/tests/src/com/android/launcher3/model/GridSizeMigrationTaskTest.java
index 08a81f0..c250cb2 100644
--- a/tests/src/com/android/launcher3/model/GridSizeMigrationTaskTest.java
+++ b/tests/src/com/android/launcher3/model/GridSizeMigrationTaskTest.java
@@ -10,6 +10,7 @@
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherModel;
 import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.config.ProviderConfig;
 import com.android.launcher3.util.TestLauncherProvider;
 
@@ -61,8 +62,13 @@
         mIdp.numHotseatIcons = 3;
         new GridSizeMigrationTask(getMockContext(), mIdp, mValidPackages, 5, 3)
                 .migrateHotseat();
-        // First & last items are dropped as they have the least weight.
-        verifyHotseat(hotseatItems[1], -1, hotseatItems[3]);
+        if (FeatureFlags.NO_ALL_APPS_ICON) {
+            // First item is dropped as it has the least weight.
+            verifyHotseat(hotseatItems[1], hotseatItems[3], hotseatItems[4]);
+        } else {
+            // First & last items are dropped as they have the least weight.
+            verifyHotseat(hotseatItems[1], -1, hotseatItems[3]);
+        }
     }
 
     public void testHotseatMigration_shortcuts_dropped() throws Exception {
@@ -77,8 +83,13 @@
         mIdp.numHotseatIcons = 3;
         new GridSizeMigrationTask(getMockContext(), mIdp, mValidPackages, 5, 3)
                 .migrateHotseat();
-        // First & third items are dropped as they have the least weight.
-        verifyHotseat(hotseatItems[1], -1, hotseatItems[4]);
+        if (FeatureFlags.NO_ALL_APPS_ICON) {
+            // First item is dropped as it has the least weight.
+            verifyHotseat(hotseatItems[1], hotseatItems[3], hotseatItems[4]);
+        } else {
+            // First & third items are dropped as they have the least weight.
+            verifyHotseat(hotseatItems[1], -1, hotseatItems[4]);
+        }
     }
 
     private void verifyHotseat(long... sortedIds) {