Merge "fix app crash on settings for launcher" into ub-launcher3-master
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/RecentsActivity.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/RecentsActivity.java
index e0e20ee..3f5179f 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/RecentsActivity.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/RecentsActivity.java
@@ -18,6 +18,7 @@
 import static com.android.launcher3.QuickstepAppTransitionManagerImpl.RECENTS_LAUNCH_DURATION;
 import static com.android.launcher3.QuickstepAppTransitionManagerImpl.STATUS_BAR_TRANSITION_DURATION;
 import static com.android.launcher3.QuickstepAppTransitionManagerImpl.STATUS_BAR_TRANSITION_PRE_DELAY;
+import static com.android.launcher3.testing.TestProtocol.OVERVIEW_STATE_ORDINAL;
 import static com.android.quickstep.TaskUtils.taskIsATargetWithMode;
 import static com.android.quickstep.TaskViewUtils.getRecentsWindowAnimator;
 import static com.android.systemui.shared.system.RemoteAnimationTargetCompat.MODE_CLOSING;
@@ -38,6 +39,7 @@
 import com.android.launcher3.LauncherAnimationRunner;
 import com.android.launcher3.R;
 import com.android.launcher3.anim.Interpolators;
+import com.android.launcher3.compat.AccessibilityManagerCompat;
 import com.android.launcher3.util.ObjectWrapper;
 import com.android.launcher3.views.BaseDragLayer;
 import com.android.quickstep.fallback.FallbackRecentsView;
@@ -217,6 +219,12 @@
         mFallbackRecentsView.reset();
     }
 
+    @Override
+    protected void onResume() {
+        super.onResume();
+        AccessibilityManagerCompat.sendStateEventToTest(getBaseContext(), OVERVIEW_STATE_ORDINAL);
+    }
+
     public void onTaskLaunched() {
         mFallbackRecentsView.resetTaskVisuals();
     }
diff --git a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
index 85cffaa..abd2180 100644
--- a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
+++ b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
@@ -31,6 +31,8 @@
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.rule.ShellCommandRule.disableHeadsUpNotification;
 import static com.android.launcher3.util.rule.ShellCommandRule.getLauncherCommand;
+import static com.android.launcher3.util.rule.TestStabilityRule.LOCAL;
+import static com.android.launcher3.util.rule.TestStabilityRule.UNBUNDLED_POSTSUBMIT;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
@@ -57,10 +59,10 @@
 import com.android.launcher3.testcomponent.TestCommandReceiver;
 import com.android.launcher3.util.Wait;
 import com.android.launcher3.util.rule.FailureWatcher;
+import com.android.launcher3.util.rule.TestStabilityRule;
 import com.android.quickstep.NavigationModeSwitchRule.NavigationModeSwitch;
 import com.android.quickstep.views.RecentsView;
 
-import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.RuleChain;
@@ -130,7 +132,7 @@
 
     @NavigationModeSwitch
     @Test
-    @Ignore // b/143488140
+    @TestStabilityRule.Stability(flavors = LOCAL | UNBUNDLED_POSTSUBMIT) // b/143488140
     public void goToOverviewFromHome() {
         mDevice.pressHome();
         assertTrue("Fallback Launcher not visible", mDevice.wait(Until.hasObject(By.pkg(
@@ -141,7 +143,7 @@
 
     @NavigationModeSwitch
     @Test
-    @Ignore // b/143488140
+    @TestStabilityRule.Stability(flavors = LOCAL | UNBUNDLED_POSTSUBMIT) // b/143488140
     public void goToOverviewFromApp() {
         startAppFastAndWaitForRecentTask(resolveSystemApp(Intent.CATEGORY_APP_CALCULATOR));
 
@@ -176,7 +178,7 @@
 
     @NavigationModeSwitch
     @Test
-    @Ignore // b/143488140
+    @TestStabilityRule.Stability(flavors = LOCAL | UNBUNDLED_POSTSUBMIT) // b/143488140
     public void testOverview() {
         startAppFastAndWaitForRecentTask(getAppPackageName());
         startAppFastAndWaitForRecentTask(resolveSystemApp(Intent.CATEGORY_APP_CALCULATOR));
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index ae30380..1ae30b7 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -126,6 +126,10 @@
     public static final TogglableFlag ENABLE_DEEP_SHORTCUT_ICON_CACHE = new TogglableFlag(
             "ENABLE_DEEP_SHORTCUT_ICON_CACHE", true, "R/W deep shortcut in IconCache");
 
+    public static final TogglableFlag ENABLE_LAUNCHER_PREVIEW_IN_GRID_PICKER = new TogglableFlag(
+            "ENABLE_LAUNCHER_PREVIEW_IN_GRID_PICKER", false,
+            "Show launcher preview in grid picker");
+
     public static void initialize(Context context) {
         // Avoid the disk read for user builds
         if (Utilities.IS_DEBUG_DEVICE) {
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index 2badb6e..0c5535f 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -19,6 +19,10 @@
 import static android.view.View.MeasureSpec.makeMeasureSpec;
 import static android.view.View.VISIBLE;
 
+import static com.android.launcher3.config.FeatureFlags.ENABLE_LAUNCHER_PREVIEW_IN_GRID_PICKER;
+import static com.android.launcher3.model.ModelUtils.filterCurrentWorkspaceItems;
+import static com.android.launcher3.model.ModelUtils.sortWorkspaceItemsSpatially;
+
 import android.annotation.TargetApi;
 import android.app.Fragment;
 import android.content.Context;
@@ -48,6 +52,10 @@
 import com.android.launcher3.Hotseat;
 import com.android.launcher3.InsettableFrameLayout;
 import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.ItemInfo;
+import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherModel;
+import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
@@ -58,11 +66,20 @@
 import com.android.launcher3.icons.BaseIconFactory;
 import com.android.launcher3.icons.BitmapInfo;
 import com.android.launcher3.icons.BitmapRenderer;
+import com.android.launcher3.model.AllAppsList;
+import com.android.launcher3.model.BgDataModel;
+import com.android.launcher3.model.LoaderResults;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.BaseDragLayer;
 
+import java.util.ArrayList;
 import java.util.concurrent.Callable;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executor;
+import java.util.concurrent.FutureTask;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 /**
  * Utility class for generating the preview of Launcher for a given InvariantDeviceProfile.
@@ -241,22 +258,59 @@
         }
 
         private void renderScreenShot(Canvas canvas) {
-            // Add hotseat icons
-            for (int i = 0; i < mIdp.numHotseatIcons; i++) {
-                WorkspaceItemInfo info = new WorkspaceItemInfo(mWorkspaceItemInfo);
-                info.container = Favorites.CONTAINER_HOTSEAT;
-                info.screenId = i;
-                inflateAndAddIcon(info);
-            }
+            if (ENABLE_LAUNCHER_PREVIEW_IN_GRID_PICKER.get()) {
+                final LauncherModel launcherModel = LauncherAppState.getInstance(
+                        mContext).getModel();
+                final WorkspaceItemsInfoFetcher fetcher = new WorkspaceItemsInfoFetcher();
+                launcherModel.enqueueModelUpdateTask(fetcher);
+                ArrayList<ItemInfo> workspaceItems;
+                try {
+                    workspaceItems = fetcher.mTask.get(5, TimeUnit.SECONDS);
+                } catch (InterruptedException | ExecutionException | TimeoutException e) {
+                    Log.d(TAG, "Error fetching workspace items info", e);
+                    return;
+                }
 
-            // Add workspace icons
-            for (int i = 0; i < mIdp.numColumns; i++) {
-                WorkspaceItemInfo info = new WorkspaceItemInfo(mWorkspaceItemInfo);
-                info.container = Favorites.CONTAINER_DESKTOP;
-                info.screenId = 0;
-                info.cellX = i;
-                info.cellY = mIdp.numRows - 1;
-                inflateAndAddIcon(info);
+                // Separate the items that are on the current screen, and all the other remaining
+                // items
+                ArrayList<ItemInfo> currentWorkspaceItems = new ArrayList<>();
+                ArrayList<ItemInfo> otherWorkspaceItems = new ArrayList<>();
+
+                filterCurrentWorkspaceItems(0 /* currentScreenId */, workspaceItems,
+                        currentWorkspaceItems, otherWorkspaceItems);
+                sortWorkspaceItemsSpatially(mIdp, currentWorkspaceItems);
+
+                for (ItemInfo itemInfo : currentWorkspaceItems) {
+                    switch (itemInfo.itemType) {
+                        case Favorites.ITEM_TYPE_APPLICATION:
+                        case Favorites.ITEM_TYPE_SHORTCUT:
+                        case Favorites.ITEM_TYPE_DEEP_SHORTCUT:
+                            inflateAndAddIcon((WorkspaceItemInfo) itemInfo);
+                            break;
+                        case LauncherSettings.Favorites.ITEM_TYPE_FOLDER:
+                            // TODO: for folder implementation here.
+                            break;
+                        default:
+                            break;
+                    }
+                }
+            } else {
+                // Add hotseat icons
+                for (int i = 0; i < mIdp.numHotseatIcons; i++) {
+                    WorkspaceItemInfo info = new WorkspaceItemInfo(mWorkspaceItemInfo);
+                    info.container = Favorites.CONTAINER_HOTSEAT;
+                    info.screenId = i;
+                    inflateAndAddIcon(info);
+                }
+                // Add workspace icons
+                for (int i = 0; i < mIdp.numColumns; i++) {
+                    WorkspaceItemInfo info = new WorkspaceItemInfo(mWorkspaceItemInfo);
+                    info.container = Favorites.CONTAINER_DESKTOP;
+                    info.screenId = 0;
+                    info.cellX = i;
+                    info.cellY = mIdp.numRows - 1;
+                    inflateAndAddIcon(info);
+                }
             }
 
             // Add first page QSB
@@ -286,6 +340,42 @@
         }
     }
 
+    private static class WorkspaceItemsInfoFetcher implements Callable<ArrayList<ItemInfo>>,
+            LauncherModel.ModelUpdateTask {
+
+        private final FutureTask<ArrayList<ItemInfo>> mTask = new FutureTask<>(this);
+
+        private LauncherAppState mApp;
+        private LauncherModel mModel;
+        private BgDataModel mBgDataModel;
+        private AllAppsList mAllAppsList;
+
+        @Override
+        public void init(LauncherAppState app, LauncherModel model, BgDataModel dataModel,
+                AllAppsList allAppsList, Executor uiExecutor) {
+            mApp = app;
+            mModel = model;
+            mBgDataModel = dataModel;
+            mAllAppsList = allAppsList;
+        }
+
+        @Override
+        public void run() {
+            mTask.run();
+        }
+
+        @Override
+        public ArrayList<ItemInfo> call() throws Exception {
+            if (!mModel.isModelLoaded()) {
+                Log.d(TAG, "Workspace not loaded, loading now");
+                mModel.startLoaderForResults(
+                        new LoaderResults(mApp, mBgDataModel, mAllAppsList, 0, null));
+                return new ArrayList<>();
+            }
+            return mBgDataModel.workspaceItems;
+        }
+    }
+
     private static void measureView(View view, int width, int height) {
         view.measure(makeMeasureSpec(width, EXACTLY), makeMeasureSpec(height, EXACTLY));
         view.layout(0, 0, width, height);
diff --git a/src/com/android/launcher3/model/BaseLoaderResults.java b/src/com/android/launcher3/model/BaseLoaderResults.java
index a00a6bd..76c2951 100644
--- a/src/com/android/launcher3/model/BaseLoaderResults.java
+++ b/src/com/android/launcher3/model/BaseLoaderResults.java
@@ -16,6 +16,8 @@
 
 package com.android.launcher3.model;
 
+import static com.android.launcher3.model.ModelUtils.filterCurrentWorkspaceItems;
+import static com.android.launcher3.model.ModelUtils.sortWorkspaceItemsSpatially;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
 import android.os.Looper;
@@ -27,20 +29,15 @@
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherAppWidgetInfo;
 import com.android.launcher3.LauncherModel.CallbackTask;
-import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.PagedView;
-import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.util.IntArray;
-import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.LooperIdleLock;
 import com.android.launcher3.util.ViewOnDrawExecutor;
 
 import java.lang.ref.WeakReference;
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.Comparator;
-import java.util.Iterator;
 import java.util.concurrent.Executor;
 
 /**
@@ -123,8 +120,9 @@
                 otherWorkspaceItems);
         filterCurrentWorkspaceItems(currentScreenId, appWidgets, currentAppWidgets,
                 otherAppWidgets);
-        sortWorkspaceItemsSpatially(currentWorkspaceItems);
-        sortWorkspaceItemsSpatially(otherWorkspaceItems);
+        final InvariantDeviceProfile idp = mApp.getInvariantDeviceProfile();
+        sortWorkspaceItemsSpatially(idp, currentWorkspaceItems);
+        sortWorkspaceItemsSpatially(idp, otherWorkspaceItems);
 
         // Tell the workspace that we're about to start binding items
         executeCallbacksTask(c -> {
@@ -169,89 +167,6 @@
         }
     }
 
-
-    /** Filters the set of items who are directly or indirectly (via another container) on the
-     * specified screen. */
-    public static <T extends ItemInfo> void filterCurrentWorkspaceItems(int currentScreenId,
-            ArrayList<T> allWorkspaceItems,
-            ArrayList<T> currentScreenItems,
-            ArrayList<T> otherScreenItems) {
-        // Purge any null ItemInfos
-        Iterator<T> iter = allWorkspaceItems.iterator();
-        while (iter.hasNext()) {
-            ItemInfo i = iter.next();
-            if (i == null) {
-                iter.remove();
-            }
-        }
-
-        // Order the set of items by their containers first, this allows use to walk through the
-        // list sequentially, build up a list of containers that are in the specified screen,
-        // as well as all items in those containers.
-        IntSet itemsOnScreen = new IntSet();
-        Collections.sort(allWorkspaceItems,
-                (lhs, rhs) -> Integer.compare(lhs.container, rhs.container));
-
-        for (T info : allWorkspaceItems) {
-            if (info.container == LauncherSettings.Favorites.CONTAINER_DESKTOP) {
-                if (info.screenId == currentScreenId) {
-                    currentScreenItems.add(info);
-                    itemsOnScreen.add(info.id);
-                } else {
-                    otherScreenItems.add(info);
-                }
-            } else if (info.container == LauncherSettings.Favorites.CONTAINER_HOTSEAT) {
-                currentScreenItems.add(info);
-                itemsOnScreen.add(info.id);
-            } else {
-                if (itemsOnScreen.contains(info.container)) {
-                    currentScreenItems.add(info);
-                    itemsOnScreen.add(info.id);
-                } else {
-                    otherScreenItems.add(info);
-                }
-            }
-        }
-    }
-
-    /** Sorts the set of items by hotseat, workspace (spatially from top to bottom, left to
-     * right) */
-    protected void sortWorkspaceItemsSpatially(ArrayList<ItemInfo> workspaceItems) {
-        final InvariantDeviceProfile profile = mApp.getInvariantDeviceProfile();
-        final int screenCols = profile.numColumns;
-        final int screenCellCount = profile.numColumns * profile.numRows;
-        Collections.sort(workspaceItems, new Comparator<ItemInfo>() {
-            @Override
-            public int compare(ItemInfo lhs, ItemInfo rhs) {
-                if (lhs.container == rhs.container) {
-                    // Within containers, order by their spatial position in that container
-                    switch (lhs.container) {
-                        case LauncherSettings.Favorites.CONTAINER_DESKTOP: {
-                            int lr = (lhs.screenId * screenCellCount +
-                                    lhs.cellY * screenCols + lhs.cellX);
-                            int rr = (rhs.screenId * screenCellCount +
-                                    rhs.cellY * screenCols + rhs.cellX);
-                            return Integer.compare(lr, rr);
-                        }
-                        case LauncherSettings.Favorites.CONTAINER_HOTSEAT: {
-                            // We currently use the screen id as the rank
-                            return Integer.compare(lhs.screenId, rhs.screenId);
-                        }
-                        default:
-                            if (FeatureFlags.IS_DOGFOOD_BUILD) {
-                                throw new RuntimeException("Unexpected container type when " +
-                                        "sorting workspace items.");
-                            }
-                            return 0;
-                    }
-                } else {
-                    // Between containers, order by hotseat, desktop
-                    return Integer.compare(lhs.container, rhs.container);
-                }
-            }
-        });
-    }
-
     protected void bindWorkspaceItems(final ArrayList<ItemInfo> workspaceItems,
             final Executor executor) {
         // Bind the workspace items
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 2754cfc..cc994df 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -19,7 +19,7 @@
 import static com.android.launcher3.ItemInfoWithIcon.FLAG_DISABLED_LOCKED_USER;
 import static com.android.launcher3.ItemInfoWithIcon.FLAG_DISABLED_SAFEMODE;
 import static com.android.launcher3.ItemInfoWithIcon.FLAG_DISABLED_SUSPENDED;
-import static com.android.launcher3.model.LoaderResults.filterCurrentWorkspaceItems;
+import static com.android.launcher3.model.ModelUtils.filterCurrentWorkspaceItems;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 import static com.android.launcher3.util.PackageManagerHelper.isSystemApp;
 
diff --git a/src/com/android/launcher3/model/ModelUtils.java b/src/com/android/launcher3/model/ModelUtils.java
new file mode 100644
index 0000000..628dd95
--- /dev/null
+++ b/src/com/android/launcher3/model/ModelUtils.java
@@ -0,0 +1,112 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.model;
+
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.ItemInfo;
+import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.config.FeatureFlags;
+import com.android.launcher3.util.IntSet;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Iterator;
+
+/**
+ * Utils class for {@link com.android.launcher3.LauncherModel}.
+ */
+public class ModelUtils {
+
+    /**
+     * Filters the set of items who are directly or indirectly (via another container) on the
+     * specified screen.
+     */
+    public static <T extends ItemInfo> void filterCurrentWorkspaceItems(int currentScreenId,
+            ArrayList<T> allWorkspaceItems,
+            ArrayList<T> currentScreenItems,
+            ArrayList<T> otherScreenItems) {
+        // Purge any null ItemInfos
+        Iterator<T> iter = allWorkspaceItems.iterator();
+        while (iter.hasNext()) {
+            ItemInfo i = iter.next();
+            if (i == null) {
+                iter.remove();
+            }
+        }
+        // Order the set of items by their containers first, this allows use to walk through the
+        // list sequentially, build up a list of containers that are in the specified screen,
+        // as well as all items in those containers.
+        IntSet itemsOnScreen = new IntSet();
+        Collections.sort(allWorkspaceItems,
+                (lhs, rhs) -> Integer.compare(lhs.container, rhs.container));
+        for (T info : allWorkspaceItems) {
+            if (info.container == LauncherSettings.Favorites.CONTAINER_DESKTOP) {
+                if (info.screenId == currentScreenId) {
+                    currentScreenItems.add(info);
+                    itemsOnScreen.add(info.id);
+                } else {
+                    otherScreenItems.add(info);
+                }
+            } else if (info.container == LauncherSettings.Favorites.CONTAINER_HOTSEAT) {
+                currentScreenItems.add(info);
+                itemsOnScreen.add(info.id);
+            } else {
+                if (itemsOnScreen.contains(info.container)) {
+                    currentScreenItems.add(info);
+                    itemsOnScreen.add(info.id);
+                } else {
+                    otherScreenItems.add(info);
+                }
+            }
+        }
+    }
+
+    /**
+     * Sorts the set of items by hotseat, workspace (spatially from top to bottom, left to right)
+     */
+    public static void sortWorkspaceItemsSpatially(InvariantDeviceProfile profile,
+            ArrayList<ItemInfo> workspaceItems) {
+        final int screenCols = profile.numColumns;
+        final int screenCellCount = profile.numColumns * profile.numRows;
+        Collections.sort(workspaceItems, (lhs, rhs) -> {
+            if (lhs.container == rhs.container) {
+                // Within containers, order by their spatial position in that container
+                switch (lhs.container) {
+                    case LauncherSettings.Favorites.CONTAINER_DESKTOP: {
+                        int lr = (lhs.screenId * screenCellCount + lhs.cellY * screenCols
+                                + lhs.cellX);
+                        int rr = (rhs.screenId * screenCellCount + +rhs.cellY * screenCols
+                                + rhs.cellX);
+                        return Integer.compare(lr, rr);
+                    }
+                    case LauncherSettings.Favorites.CONTAINER_HOTSEAT: {
+                        // We currently use the screen id as the rank
+                        return Integer.compare(lhs.screenId, rhs.screenId);
+                    }
+                    default:
+                        if (FeatureFlags.IS_DOGFOOD_BUILD) {
+                            throw new RuntimeException(
+                                    "Unexpected container type when sorting workspace items.");
+                        }
+                        return 0;
+                }
+            } else {
+                // Between containers, order by hotseat, desktop
+                return Integer.compare(lhs.container, rhs.container);
+            }
+        });
+    }
+}
diff --git a/src/com/android/launcher3/model/ModelWriter.java b/src/com/android/launcher3/model/ModelWriter.java
index c69ace9..bdf3a69 100644
--- a/src/com/android/launcher3/model/ModelWriter.java
+++ b/src/com/android/launcher3/model/ModelWriter.java
@@ -103,7 +103,7 @@
      */
     public void addOrMoveItemInDatabase(ItemInfo item,
             int container, int screenId, int cellX, int cellY) {
-        if (item.container == ItemInfo.NO_ID) {
+        if (item.id == ItemInfo.NO_ID) {
             // From all apps
             addItemToDatabase(item, container, screenId, cellX, cellY);
         } else {
diff --git a/tests/tapl/com/android/launcher3/tapl/Background.java b/tests/tapl/com/android/launcher3/tapl/Background.java
index 6583d32..d9ae778 100644
--- a/tests/tapl/com/android/launcher3/tapl/Background.java
+++ b/tests/tapl/com/android/launcher3/tapl/Background.java
@@ -17,6 +17,7 @@
 package com.android.launcher3.tapl;
 
 import static com.android.launcher3.testing.TestProtocol.BACKGROUND_APP_STATE_ORDINAL;
+import static com.android.launcher3.testing.TestProtocol.OVERVIEW_STATE_ORDINAL;
 
 import android.graphics.Point;
 import android.os.SystemClock;
@@ -54,13 +55,13 @@
         try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
                 "want to switch from background to overview")) {
             verifyActiveContainer();
-            goToOverviewUnchecked(BACKGROUND_APP_STATE_ORDINAL);
+            goToOverviewUnchecked();
             return mLauncher.isFallbackOverview() ?
                     new BaseOverview(mLauncher) : new Overview(mLauncher);
         }
     }
 
-    protected void goToOverviewUnchecked(int expectedState) {
+    protected void goToOverviewUnchecked() {
         switch (mLauncher.getNavigationModel()) {
             case ZERO_BUTTON: {
                 final int centerX = mLauncher.getDevice().getDisplayWidth() / 2;
@@ -81,9 +82,11 @@
                                 start,
                                 end),
                         event -> TestProtocol.PAUSE_DETECTED_MESSAGE.equals(event.getClassName()),
-                        "Pause wasn't detected");
-                mLauncher.sendPointer(
-                        downTime, SystemClock.uptimeMillis(), MotionEvent.ACTION_UP, end);
+                        () -> "Pause wasn't detected");
+                mLauncher.runToState(
+                        () -> mLauncher.sendPointer(
+                                downTime, SystemClock.uptimeMillis(), MotionEvent.ACTION_UP, end),
+                        OVERVIEW_STATE_ORDINAL);
                 break;
             }
 
@@ -105,17 +108,14 @@
                     startY = endY = mLauncher.getDevice().getDisplayHeight() / 2;
                 }
 
-                if (mLauncher.isFallbackOverview()) {
-                    mLauncher.linearGesture(startX, startY, endX, endY, 10, false);
-                    new BaseOverview(mLauncher);
-                } else {
-                    mLauncher.swipeToState(startX, startY, endX, endY, 10, expectedState);
-                }
+                mLauncher.swipeToState(startX, startY, endX, endY, 10, OVERVIEW_STATE_ORDINAL);
                 break;
             }
 
             case THREE_BUTTON:
-                mLauncher.waitForSystemUiObject("recent_apps").click();
+                mLauncher.runToState(
+                        () -> mLauncher.waitForSystemUiObject("recent_apps").click(),
+                        OVERVIEW_STATE_ORDINAL);
                 break;
         }
     }
@@ -167,7 +167,7 @@
             case THREE_BUTTON:
                 // Double press the recents button.
                 UiObject2 recentsButton = mLauncher.waitForSystemUiObject("recent_apps");
-                recentsButton.click();
+                mLauncher.runToState(() -> recentsButton.click(), OVERVIEW_STATE_ORDINAL);
                 mLauncher.getOverview();
                 recentsButton.click();
                 break;
diff --git a/tests/tapl/com/android/launcher3/tapl/Home.java b/tests/tapl/com/android/launcher3/tapl/Home.java
index e0fe933..1e4d937 100644
--- a/tests/tapl/com/android/launcher3/tapl/Home.java
+++ b/tests/tapl/com/android/launcher3/tapl/Home.java
@@ -16,7 +16,6 @@
 
 package com.android.launcher3.tapl;
 
-import static com.android.launcher3.testing.TestProtocol.OVERVIEW_STATE_ORDINAL;
 import static com.android.launcher3.testing.TestProtocol.QUICK_SWITCH_STATE_ORDINAL;
 
 import androidx.annotation.NonNull;
@@ -52,7 +51,7 @@
         try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
                 "want to switch from home to overview")) {
             verifyActiveContainer();
-            goToOverviewUnchecked(OVERVIEW_STATE_ORDINAL);
+            goToOverviewUnchecked();
             try (LauncherInstrumentation.Closable c1 = mLauncher.addContextLayer(
                     "performed the switch action")) {
                 return new Overview(mLauncher);
diff --git a/tests/tapl/com/android/launcher3/tapl/Launchable.java b/tests/tapl/com/android/launcher3/tapl/Launchable.java
index df80a51..6881197 100644
--- a/tests/tapl/com/android/launcher3/tapl/Launchable.java
+++ b/tests/tapl/com/android/launcher3/tapl/Launchable.java
@@ -56,7 +56,7 @@
         mLauncher.executeAndWaitForEvent(
                 () -> mObject.click(),
                 event -> event.getEventType() == TYPE_WINDOW_STATE_CHANGED,
-                "Launching an app didn't open a new window: " + mObject.getText());
+                () -> "Launching an app didn't open a new window: " + mObject.getText());
 
         mLauncher.assertTrue(
                 "App didn't start: " + selector,
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 727d757..2fea1b7 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -22,7 +22,6 @@
 import static android.content.pm.PackageManager.MATCH_DISABLED_COMPONENTS;
 
 import static com.android.launcher3.tapl.TestHelpers.getOverviewPackageName;
-import static com.android.launcher3.testing.TestProtocol.BACKGROUND_APP_STATE_ORDINAL;
 import static com.android.launcher3.testing.TestProtocol.NORMAL_STATE_ORDINAL;
 
 import android.app.ActivityManager;
@@ -78,6 +77,8 @@
 import java.util.concurrent.TimeoutException;
 import java.util.function.Consumer;
 import java.util.function.Function;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
 
 /**
  * The main tapl object. The only object that can be explicitly constructed by the using code. It
@@ -510,7 +511,7 @@
     }
 
     Parcelable executeAndWaitForEvent(Runnable command,
-            UiAutomation.AccessibilityEventFilter eventFilter, String message) {
+            UiAutomation.AccessibilityEventFilter eventFilter, Supplier<String> message) {
         try {
             final AccessibilityEvent event =
                     mInstrumentation.getUiAutomation().executeAndWaitForEvent(
@@ -518,7 +519,7 @@
             assertNotNull("executeAndWaitForEvent returned null (this can't happen)", event);
             return event.getParcelableData();
         } catch (TimeoutException e) {
-            fail(message);
+            fail(message.get());
             return null;
         }
     }
@@ -557,15 +558,12 @@
                 log("Hierarchy before swiping up to home");
                 dumpViewHierarchy();
                 log(action = "swiping up to home from " + getVisibleStateMessage());
-                final int finalState = mDevice.hasObject(By.pkg(getLauncherPackageName()))
-                        || isFallbackOverview()
-                        ? NORMAL_STATE_ORDINAL : BACKGROUND_APP_STATE_ORDINAL;
 
                 try (LauncherInstrumentation.Closable c = addContextLayer(action)) {
                     swipeToState(
                             displaySize.x / 2, displaySize.y - 1,
                             displaySize.x / 2, 0,
-                            ZERO_BUTTON_STEPS_FROM_BACKGROUND_TO_HOME, finalState);
+                            ZERO_BUTTON_STEPS_FROM_BACKGROUND_TO_HOME, NORMAL_STATE_ORDINAL);
                 }
             }
         } else {
@@ -576,7 +574,7 @@
                         waitForSystemUiObject("home").click();
                     },
                     event -> true,
-                    "Pressing Home didn't produce any events");
+                    () -> "Pressing Home didn't produce any events");
             mDevice.waitForIdle();
         }
         try (LauncherInstrumentation.Closable c = addContextLayer(
@@ -771,14 +769,38 @@
         return mDevice;
     }
 
+    private static String eventListToString(List<Integer> actualEvents) {
+        if (actualEvents.isEmpty()) return "no events";
+
+        return "["
+                + actualEvents.stream()
+                .map(state -> TestProtocol.stateOrdinalToString(state))
+                .collect(Collectors.joining(", "))
+                + "]";
+    }
+
+    void runToState(Runnable command, int expectedState) {
+        final List<Integer> actualEvents = new ArrayList<>();
+        executeAndWaitForEvent(
+                command,
+                event -> isSwitchToStateEvent(event, expectedState, actualEvents),
+                () -> "Failed to receive an event for the swipe end: expected "
+                        + TestProtocol.stateOrdinalToString(expectedState)
+                        + ", actual: " + eventListToString(actualEvents));
+    }
+
+    private boolean isSwitchToStateEvent(
+            AccessibilityEvent event, int expectedState, List<Integer> actualEvents) {
+        if (!TestProtocol.SWITCHED_TO_STATE_MESSAGE.equals(event.getClassName())) return false;
+
+        final Bundle parcel = (Bundle) event.getParcelableData();
+        final int actualState = parcel.getInt(TestProtocol.STATE_FIELD);
+        actualEvents.add(actualState);
+        return actualState == expectedState;
+    }
+
     void swipeToState(int startX, int startY, int endX, int endY, int steps, int expectedState) {
-        final Bundle parcel = (Bundle) executeAndWaitForEvent(
-                () -> linearGesture(startX, startY, endX, endY, steps, false),
-                event -> TestProtocol.SWITCHED_TO_STATE_MESSAGE.equals(event.getClassName()),
-                "Swipe failed to receive an event for the swipe end");
-        assertEquals("Swipe switched launcher to a wrong state;",
-                TestProtocol.stateOrdinalToString(expectedState),
-                TestProtocol.stateOrdinalToString(parcel.getInt(TestProtocol.STATE_FIELD)));
+        runToState(() -> linearGesture(startX, startY, endX, endY, steps, false), expectedState);
     }
 
     int getBottomGestureSize() {
@@ -864,7 +886,7 @@
         executeAndWaitForEvent(
                 () -> linearGesture(startX, startY, endX, endY, steps, slowDown),
                 event -> TestProtocol.SCROLL_FINISHED_MESSAGE.equals(event.getClassName()),
-                "Didn't receive a scroll end message: " + startX + ", " + startY
+                () -> "Didn't receive a scroll end message: " + startX + ", " + startY
                         + ", " + endX + ", " + endY);
     }
 
diff --git a/tests/tapl/com/android/launcher3/tapl/OverviewTask.java b/tests/tapl/com/android/launcher3/tapl/OverviewTask.java
index 2ee424b..46f8ba5 100644
--- a/tests/tapl/com/android/launcher3/tapl/OverviewTask.java
+++ b/tests/tapl/com/android/launcher3/tapl/OverviewTask.java
@@ -67,8 +67,8 @@
             mLauncher.executeAndWaitForEvent(
                     () -> mTask.click(),
                     event -> event.getEventType() == TYPE_WINDOW_STATE_CHANGED,
-                    "Launching task didn't open a new window: " +
-                            mTask.getParent().getContentDescription());
+                    () -> "Launching task didn't open a new window: "
+                            + mTask.getParent().getContentDescription());
         }
         return new Background(mLauncher);
     }
diff --git a/tests/tapl/com/android/launcher3/tapl/Workspace.java b/tests/tapl/com/android/launcher3/tapl/Workspace.java
index 81d343d..8a53ef1 100644
--- a/tests/tapl/com/android/launcher3/tapl/Workspace.java
+++ b/tests/tapl/com/android/launcher3/tapl/Workspace.java
@@ -17,6 +17,8 @@
 package com.android.launcher3.tapl;
 
 import static com.android.launcher3.testing.TestProtocol.ALL_APPS_STATE_ORDINAL;
+import static com.android.launcher3.testing.TestProtocol.NORMAL_STATE_ORDINAL;
+import static com.android.launcher3.testing.TestProtocol.SPRING_LOADED_STATE_ORDINAL;
 
 import static junit.framework.TestCase.assertTrue;
 
@@ -165,14 +167,21 @@
         LauncherInstrumentation.log("dragIconToWorkspace: begin");
         final Point launchableCenter = launchable.getObject().getVisibleCenter();
         final long downTime = SystemClock.uptimeMillis();
-        launcher.sendPointer(downTime, downTime, MotionEvent.ACTION_DOWN, launchableCenter);
-        LauncherInstrumentation.log("dragIconToWorkspace: sent down");
-        launcher.waitForLauncherObject(longPressIndicator);
-        LauncherInstrumentation.log("dragIconToWorkspace: indicator");
-        launcher.movePointer(launchableCenter, dest, 10, downTime, true);
+        launcher.runToState(
+                () -> {
+                    launcher.sendPointer(downTime, downTime, MotionEvent.ACTION_DOWN,
+                            launchableCenter);
+                    LauncherInstrumentation.log("dragIconToWorkspace: sent down");
+                    launcher.waitForLauncherObject(longPressIndicator);
+                    LauncherInstrumentation.log("dragIconToWorkspace: indicator");
+                    launcher.movePointer(launchableCenter, dest, 10, downTime, true);
+                },
+                SPRING_LOADED_STATE_ORDINAL);
         LauncherInstrumentation.log("dragIconToWorkspace: moved pointer");
-        launcher.sendPointer(
-                downTime, SystemClock.uptimeMillis(), MotionEvent.ACTION_UP, dest);
+        launcher.runToState(
+                () -> launcher.sendPointer(
+                        downTime, SystemClock.uptimeMillis(), MotionEvent.ACTION_UP, dest),
+                NORMAL_STATE_ORDINAL);
         LauncherInstrumentation.log("dragIconToWorkspace: end");
         launcher.waitUntilGone("drop_target_bar");
     }