Snap for 7196551 from ded87e65056428ff6991e2509b7a17cd8b322643 to sc-v2-release

Change-Id: Ifcac331151fafb997ea3748d9f4fca237f1e3189
diff --git a/go/src/com/android/launcher3/model/WidgetsModel.java b/go/src/com/android/launcher3/model/WidgetsModel.java
index a202095..cc5e1cb 100644
--- a/go/src/com/android/launcher3/model/WidgetsModel.java
+++ b/go/src/com/android/launcher3/model/WidgetsModel.java
@@ -30,6 +30,7 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 /**
@@ -53,10 +54,15 @@
      *
      * @see com.android.launcher3.widget.picker.WidgetsListAdapter#setWidgets(List)
      */
-    public synchronized ArrayList<WidgetsListBaseEntry> getWidgetsList(Context context) {
+    public synchronized ArrayList<WidgetsListBaseEntry> getWidgetsListForPicker(Context context) {
         return EMPTY_WIDGET_LIST;
     }
 
+    /** Returns a mapping of packages to their widgets without static shortcuts. */
+    public synchronized Map<PackageUserKey, List<WidgetItem>> getAllWidgetsWithoutShortcuts() {
+        return Map.of();
+    }
+
     /**
      * @param packageUser If null, all widgets and shortcuts are updated and returned, otherwise
      *                    only widgets and shortcuts associated with the package/user are.
diff --git a/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java b/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
new file mode 100644
index 0000000..70a143e
--- /dev/null
+++ b/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
@@ -0,0 +1,263 @@
+/*
+ * Copyright (C) 2021 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 static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
+import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
+import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.robolectric.Shadows.shadowOf;
+
+import android.app.prediction.AppTarget;
+import android.app.prediction.AppTargetId;
+import android.appwidget.AppWidgetManager;
+import android.appwidget.AppWidgetProviderInfo;
+import android.content.ComponentName;
+import android.content.Context;
+import android.os.Process;
+import android.os.UserHandle;
+
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.icons.ComponentWithLabel;
+import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.model.BgDataModel.FixedContainerItems;
+import com.android.launcher3.model.QuickstepModelDelegate.PredictorState;
+import com.android.launcher3.model.data.AppInfo;
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.model.data.LauncherAppWidgetInfo;
+import com.android.launcher3.model.data.WorkspaceItemInfo;
+import com.android.launcher3.util.ComponentKey;
+import com.android.launcher3.util.IntArray;
+import com.android.launcher3.util.ItemInfoMatcher;
+import com.android.launcher3.util.LauncherModelHelper;
+import com.android.launcher3.util.ViewOnDrawExecutor;
+import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.widget.PendingAddWidgetInfo;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.shadows.ShadowAppWidgetManager;
+import org.robolectric.shadows.ShadowPackageManager;
+import org.robolectric.util.ReflectionHelpers;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.stream.Collectors;
+
+@RunWith(RobolectricTestRunner.class)
+public final class WidgetsPredicationUpdateTaskTest {
+
+    private AppWidgetProviderInfo mApp1Provider1 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp1Provider2 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp2Provider1 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp4Provider1 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp4Provider2 = new AppWidgetProviderInfo();
+    private AppWidgetProviderInfo mApp5Provider1 = new AppWidgetProviderInfo();
+
+    private FakeBgDataModelCallback mCallback = new FakeBgDataModelCallback();
+    private Context mContext;
+    private LauncherModelHelper mModelHelper;
+    private UserHandle mUserHandle;
+    private InvariantDeviceProfile mTestProfile;
+
+    @Mock
+    private IconCache mIconCache;
+
+    @Before
+    public void setup() throws Exception {
+        MockitoAnnotations.initMocks(this);
+        doAnswer(invocation -> {
+            ComponentWithLabel componentWithLabel = invocation.getArgument(0);
+            return componentWithLabel.getComponent().getShortClassName();
+        }).when(mIconCache).getTitleNoCache(any());
+
+        mContext = RuntimeEnvironment.application;
+        mModelHelper = new LauncherModelHelper();
+        mUserHandle = Process.myUserHandle();
+        mTestProfile = new InvariantDeviceProfile();
+        // 2 widgets, app4/provider1 & app5/provider1, have already been added to the workspace.
+        mModelHelper.initializeData("/widgets_predication_update_task_data.txt");
+
+        ShadowPackageManager packageManager = shadowOf(mContext.getPackageManager());
+        mApp1Provider1.provider = ComponentName.createRelative("app1", "provider1");
+        ReflectionHelpers.setField(mApp1Provider1, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp1Provider1.provider));
+        mApp1Provider2.provider = ComponentName.createRelative("app1", "provider2");
+        ReflectionHelpers.setField(mApp1Provider2, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp1Provider2.provider));
+        mApp2Provider1.provider = ComponentName.createRelative("app2", "provider1");
+        ReflectionHelpers.setField(mApp2Provider1, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp2Provider1.provider));
+        mApp4Provider1.provider = ComponentName.createRelative("app4", "provider1");
+        ReflectionHelpers.setField(mApp4Provider1, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp4Provider1.provider));
+        mApp4Provider2.provider = ComponentName.createRelative("app4", ".provider2");
+        ReflectionHelpers.setField(mApp4Provider2, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp4Provider2.provider));
+        mApp5Provider1.provider = ComponentName.createRelative("app5", "provider1");
+        ReflectionHelpers.setField(mApp5Provider1, "providerInfo",
+                packageManager.addReceiverIfNotPresent(mApp5Provider1.provider));
+
+        ShadowAppWidgetManager shadowAppWidgetManager =
+                shadowOf(mContext.getSystemService(AppWidgetManager.class));
+        shadowAppWidgetManager.addInstalledProvider(mApp1Provider1);
+        shadowAppWidgetManager.addInstalledProvider(mApp1Provider2);
+        shadowAppWidgetManager.addInstalledProvider(mApp2Provider1);
+        shadowAppWidgetManager.addInstalledProvider(mApp4Provider1);
+        shadowAppWidgetManager.addInstalledProvider(mApp4Provider2);
+        shadowAppWidgetManager.addInstalledProvider(mApp5Provider1);
+
+        mModelHelper.getModel().addCallbacks(mCallback);
+
+        MODEL_EXECUTOR.post(() -> mModelHelper.getBgDataModel().widgetsModel.update(
+                LauncherAppState.getInstance(mContext), /* packageUser= */ null));
+        waitUntilIdle();
+    }
+
+
+    @Test
+    public void widgetsRecommendationRan_shouldOnlyReturnNotAddedWidgetsInAppPredictionOrder()
+            throws Exception {
+        // WHEN newPredicationTask is executed with app predication of 5 apps.
+        AppTarget app1 = new AppTarget(new AppTargetId("app1"), "app1", "className",
+                mUserHandle);
+        AppTarget app2 = new AppTarget(new AppTargetId("app2"), "app2", "className",
+                mUserHandle);
+        AppTarget app3 = new AppTarget(new AppTargetId("app3"), "app3", "className",
+                mUserHandle);
+        AppTarget app4 = new AppTarget(new AppTargetId("app4"), "app4", "className",
+                mUserHandle);
+        AppTarget app5 = new AppTarget(new AppTargetId("app5"), "app5", "className",
+                mUserHandle);
+        mModelHelper.executeTaskForTest(
+                newWidgetsPredicationTask(List.of(app5, app3, app2, app4, app1)))
+                .forEach(Runnable::run);
+
+        // THEN only 3 widgets are returned because
+        // 1. app5/provider1 & app4/provider1 have already been added to workspace. They are
+        //    excluded from the result.
+        // 2. app3 doesn't have a widget.
+        // 3. only 1 widget is picked from app1 because we only want to promote one widget per app.
+        List<PendingAddWidgetInfo> recommendedWidgets = mCallback.mRecommendedWidgets.items
+                .stream()
+                .map(itemInfo -> (PendingAddWidgetInfo) itemInfo)
+                .collect(Collectors.toList());
+        assertThat(recommendedWidgets).hasSize(3);
+        assertWidgetInfo(recommendedWidgets.get(0).info, mApp2Provider1);
+        assertWidgetInfo(recommendedWidgets.get(1).info, mApp4Provider2);
+        assertWidgetInfo(recommendedWidgets.get(2).info, mApp1Provider1);
+    }
+
+    private void assertWidgetInfo(
+            LauncherAppWidgetProviderInfo actual, AppWidgetProviderInfo expected) {
+        assertThat(actual.provider).isEqualTo(expected.provider);
+        assertThat(actual.getUser()).isEqualTo(expected.getProfile());
+    }
+
+    private void waitUntilIdle() {
+        shadowOf(MODEL_EXECUTOR.getLooper()).idle();
+        shadowOf(MAIN_EXECUTOR.getLooper()).idle();
+    }
+
+    private WidgetsPredictionUpdateTask newWidgetsPredicationTask(List<AppTarget> appTargets) {
+        return new WidgetsPredictionUpdateTask(
+                new PredictorState(CONTAINER_WIDGETS_PREDICTION, "test_widgets_prediction"),
+                appTargets);
+    }
+
+    private final class FakeBgDataModelCallback implements BgDataModel.Callbacks {
+
+        private FixedContainerItems mRecommendedWidgets = null;
+
+        @Override
+        public void bindExtraContainerItems(FixedContainerItems item) {
+            mRecommendedWidgets = item;
+        }
+
+        @Override
+        public int getPageToBindSynchronously() {
+            return 0;
+        }
+
+        @Override
+        public void clearPendingBinds() { }
+
+        @Override
+        public void startBinding() { }
+
+        @Override
+        public void bindItems(List<ItemInfo> shortcuts, boolean forceAnimateIcons) { }
+
+        @Override
+        public void bindScreens(IntArray orderedScreenIds) { }
+
+        @Override
+        public void finishFirstPageBind(ViewOnDrawExecutor executor) { }
+
+        @Override
+        public void finishBindingItems(int pageBoundFirst) { }
+
+        @Override
+        public void preAddApps() { }
+
+        @Override
+        public void bindAppsAdded(IntArray newScreens, ArrayList<ItemInfo> addNotAnimated,
+                ArrayList<ItemInfo> addAnimated) { }
+
+        @Override
+        public void bindIncrementalDownloadProgressUpdated(AppInfo app) { }
+
+        @Override
+        public void bindWorkspaceItemsChanged(List<WorkspaceItemInfo> updated) { }
+
+        @Override
+        public void bindWidgetsRestored(ArrayList<LauncherAppWidgetInfo> widgets) { }
+
+        @Override
+        public void bindRestoreItemsChange(HashSet<ItemInfo> updates) { }
+
+        @Override
+        public void bindWorkspaceComponentsRemoved(ItemInfoMatcher matcher) { }
+
+        @Override
+        public void bindAllWidgets(List<WidgetsListBaseEntry> widgets) { }
+
+        @Override
+        public void onPageBoundSynchronously(int page) { }
+
+        @Override
+        public void executeOnNextDraw(ViewOnDrawExecutor executor) { }
+
+        @Override
+        public void bindDeepShortcutMap(HashMap<ComponentKey, Integer> deepShortcutMap) { }
+
+        @Override
+        public void bindAllApplications(AppInfo[] apps, int flags) { }
+    }
+}
diff --git a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
index 2ea8bd2..b1b4d70 100644
--- a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
+++ b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
@@ -21,6 +21,7 @@
 import static com.android.launcher3.InvariantDeviceProfile.CHANGE_FLAG_GRID;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT_PREDICTION;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_PREDICTION;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT;
 import static com.android.launcher3.Utilities.getDevicePrefs;
@@ -80,6 +81,8 @@
             new PredictorState(CONTAINER_PREDICTION, "all_apps_predictions");
     private final PredictorState mHotseatState =
             new PredictorState(CONTAINER_HOTSEAT_PREDICTION, "hotseat_predictions");
+    private final PredictorState mWidgetsRecommendationState =
+            new PredictorState(CONTAINER_WIDGETS_PREDICTION, "widgets_prediction");
 
     private final InvariantDeviceProfile mIDP;
     private final AppEventProducer mAppEventProducer;
@@ -111,6 +114,9 @@
         mHotseatState.items.setItems(
                 mHotseatState.storage.read(mApp.getContext(), hotseatFactory, ums.allUsers::get));
         mDataModel.extraItems.put(CONTAINER_HOTSEAT_PREDICTION, mHotseatState.items);
+
+        // Widgets prediction isn't used frequently. And thus, it is not persisted on disk.
+        mDataModel.extraItems.put(CONTAINER_WIDGETS_PREDICTION, mWidgetsRecommendationState.items);
         mActive = true;
     }
 
@@ -161,6 +167,9 @@
         if (mAllAppsState.predictor != null) {
             mAllAppsState.predictor.requestPredictionUpdate();
         }
+        if (mWidgetsRecommendationState.predictor != null) {
+            mWidgetsRecommendationState.predictor.requestPredictionUpdate();
+        }
     }
 
     @Override
@@ -176,6 +185,7 @@
     private void destroyPredictors() {
         mAllAppsState.destroyPredictor();
         mHotseatState.destroyPredictor();
+        mWidgetsRecommendationState.destroyPredictor();
     }
 
     @WorkerThread
@@ -203,6 +213,12 @@
                         .setPredictedTargetCount(mIDP.numHotseatIcons)
                         .setExtras(convertDataModelToAppTargetBundle(context, mDataModel))
                         .build()));
+
+        registerWidgetsPredictor(apm.createAppPredictionSession(
+                new AppPredictionContext.Builder(context)
+                        .setUiSurface("widgets")
+                        .setPredictedTargetCount(mIDP.numColumns)
+                        .build()));
     }
 
     private void registerPredictor(PredictorState state, AppPredictor predictor) {
@@ -220,6 +236,20 @@
         mApp.getModel().enqueueModelUpdateTask(new PredictionUpdateTask(state, targets));
     }
 
+    private void registerWidgetsPredictor(AppPredictor predictor) {
+        mWidgetsRecommendationState.predictor = predictor;
+        mWidgetsRecommendationState.predictor.registerPredictionUpdates(
+                Executors.MODEL_EXECUTOR, targets -> {
+                    if (mWidgetsRecommendationState.setTargets(targets)) {
+                        // No diff, skip
+                        return;
+                    }
+                    mApp.getModel().enqueueModelUpdateTask(
+                            new WidgetsPredictionUpdateTask(mWidgetsRecommendationState, targets));
+                });
+        mWidgetsRecommendationState.predictor.requestPredictionUpdate();
+    }
+
     @Override
     public void onIdpChanged(int changeFlags, InvariantDeviceProfile profile) {
         if ((changeFlags & CHANGE_FLAG_GRID) != 0) {
diff --git a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
new file mode 100644
index 0000000..a29ac1a
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
@@ -0,0 +1,82 @@
+/*
+ * Copyright (C) 2021 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 android.app.prediction.AppTarget;
+
+import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.model.BgDataModel.FixedContainerItems;
+import com.android.launcher3.model.QuickstepModelDelegate.PredictorState;
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.util.ComponentKey;
+import com.android.launcher3.util.PackageUserKey;
+import com.android.launcher3.widget.PendingAddWidgetInfo;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/** Task to update model as a result of predicted widgets update */
+public final class WidgetsPredictionUpdateTask extends BaseModelUpdateTask {
+    private final PredictorState mPredictorState;
+    private final List<AppTarget> mTargets;
+
+    WidgetsPredictionUpdateTask(PredictorState predictorState, List<AppTarget> targets) {
+        mPredictorState = predictorState;
+        mTargets = targets;
+    }
+
+    /**
+     * Uses the app predication result to infer widgets that the user may want to use.
+     *
+     * <p>The algorithm uses the app prediction ranking to create a widgets ranking which only
+     * includes one widget per app and excludes widgets that have already been added to the
+     * workspace.
+     */
+    @Override
+    public void execute(LauncherAppState appState, BgDataModel dataModel, AllAppsList apps) {
+        Set<ComponentKey> widgetsInWorkspace = dataModel.appWidgets.stream().map(
+                widget -> new ComponentKey(widget.providerName, widget.user)).collect(
+                Collectors.toSet());
+        Map<PackageUserKey, List<WidgetItem>> allWidgets =
+                dataModel.widgetsModel.getAllWidgetsWithoutShortcuts();
+
+        ArrayList<ItemInfo> recommendedWidgetsInDescendingOrder = new ArrayList<>();
+        for (AppTarget app : mTargets) {
+            PackageUserKey packageUserKey = new PackageUserKey(app.getPackageName(), app.getUser());
+            if (allWidgets.containsKey(packageUserKey)) {
+                List<WidgetItem> notAddedWidgets = allWidgets.get(packageUserKey).stream()
+                        .filter(item ->
+                                !widgetsInWorkspace.contains(
+                                        new ComponentKey(item.componentName, item.user)))
+                        .collect(Collectors.toList());
+                if (notAddedWidgets.size() > 0) {
+                    // Even an apps have more than one widgets, we only include one widget.
+                    recommendedWidgetsInDescendingOrder.add(
+                            new PendingAddWidgetInfo(notAddedWidgets.get(0).widgetInfo));
+                }
+            }
+        }
+        FixedContainerItems fixedContainerItems = mPredictorState.items;
+        fixedContainerItems.items.clear();
+        fixedContainerItems.items.addAll(recommendedWidgetsInDescendingOrder);
+        bindExtraContainerItems(fixedContainerItems);
+
+        // Don't store widgets prediction to disk because it is not used frequently.
+    }
+}
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarController.java
index 74a82ab..544835c 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarController.java
@@ -85,6 +85,7 @@
     // Contains all loaded Hotseat items.
     private ItemInfo[] mLatestLoadedHotseatItems;
 
+    private @Nullable Animator mAnimator;
     private boolean mIsAnimatingToLauncher;
 
     public TaskbarController(BaseQuickstepLauncher launcher,
@@ -252,6 +253,10 @@
         mTaskbarVisibilityController.cleanup();
         mHotseatController.cleanup();
         mRecentsController.cleanup();
+
+        if (mAnimator != null) {
+            mAnimator.end();
+        }
     }
 
     private void removeFromWindowManager() {
@@ -294,13 +299,21 @@
      */
     public void onLauncherResumedOrPaused(boolean isResumed) {
         long duration = QuickstepAppTransitionManagerImpl.CONTENT_ALPHA_DURATION;
-        final Animator anim;
-        if (isResumed) {
-            anim = createAnimToLauncher(null, duration);
-        } else {
-            anim = createAnimToApp(duration);
+        if (mAnimator != null) {
+            mAnimator.cancel();
         }
-        anim.start();
+        if (isResumed) {
+            mAnimator = createAnimToLauncher(null, duration);
+        } else {
+            mAnimator = createAnimToApp(duration);
+        }
+        mAnimator.addListener(new AnimatorListenerAdapter() {
+            @Override
+            public void onAnimationEnd(Animator animation) {
+                mAnimator = null;
+            }
+        });
+        mAnimator.start();
     }
 
     /**
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index 55dde45..b009629 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -218,6 +218,8 @@
                     .setPredictedApps(item.items);
         } else if (item.containerId == Favorites.CONTAINER_HOTSEAT_PREDICTION) {
             mHotseatPredictionController.setPredictedItems(item);
+        } else if (item.containerId == Favorites.CONTAINER_WIDGETS_PREDICTION) {
+            getPopupDataProvider().setRecommendedWidgets(item.items);
         }
     }
 
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 7f2af6b..e0c041e 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -934,7 +934,9 @@
                     duration = Math.min(MAX_SWIPE_DURATION, 2 * baseDuration);
             }
         }
-        Interpolator interpolator = endTarget == RECENTS ? OVERSHOOT_1_2 : DEACCEL;
+        Interpolator interpolator =
+                endTarget == RECENTS ? (mDp.isTablet && FeatureFlags.ENABLE_OVERVIEW_GRID.get()
+                        ? ACCEL_DEACCEL : OVERSHOOT_1_2) : DEACCEL;
 
         if (endTarget.isLauncher) {
             mInputConsumerProxy.enable();
diff --git a/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java b/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java
index 6a32dc4..a749f9a 100644
--- a/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java
+++ b/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java
@@ -273,7 +273,6 @@
         Point size = display.realSize;
         int rotation = display.rotation;
         int touchHeight = mNavBarGesturalHeight;
-        int largerGesturalHeight = mNavBarLargerGesturalHeight;
         OrientationRectF orientationRectF =
                 new OrientationRectF(0, 0, size.x, size.y, rotation);
         if (mMode == SysUINavigationMode.Mode.NO_BUTTON) {
diff --git a/quickstep/src/com/android/quickstep/inputconsumers/OneHandedModeInputConsumer.java b/quickstep/src/com/android/quickstep/inputconsumers/OneHandedModeInputConsumer.java
index b10bdde..cd69cf1 100644
--- a/quickstep/src/com/android/quickstep/inputconsumers/OneHandedModeInputConsumer.java
+++ b/quickstep/src/com/android/quickstep/inputconsumers/OneHandedModeInputConsumer.java
@@ -21,14 +21,18 @@
 import static android.view.MotionEvent.ACTION_MOVE;
 import static android.view.MotionEvent.ACTION_UP;
 
+import static com.android.launcher3.ResourceUtils.NAVBAR_BOTTOM_GESTURE_SIZE;
 import static com.android.launcher3.Utilities.squaredHypot;
 
 import android.content.Context;
+import android.graphics.Point;
 import android.graphics.PointF;
 import android.view.MotionEvent;
 
 import com.android.launcher3.R;
+import com.android.launcher3.ResourceUtils;
 import com.android.launcher3.Utilities;
+import com.android.launcher3.util.DisplayController;
 import com.android.quickstep.InputConsumer;
 import com.android.quickstep.RecentsAnimationDeviceState;
 import com.android.quickstep.SystemUiProxy;
@@ -45,11 +49,15 @@
     private static final int ANGLE_MIN = 30;
 
     private final Context mContext;
+    private final DisplayController.DisplayHolder mDisplayHolder;
+    private final Point mDisplaySize;
     private final RecentsAnimationDeviceState mDeviceState;
 
     private final float mDragDistThreshold;
     private final float mSquaredSlop;
 
+    private final int mNavBarSize;
+
     private final PointF mDownPos = new PointF();
     private final PointF mLastPos = new PointF();
 
@@ -60,10 +68,14 @@
             InputConsumer delegate, InputMonitorCompat inputMonitor) {
         super(delegate, inputMonitor);
         mContext = context;
+        mDisplayHolder = DisplayController.getDefaultDisplay(mContext);
         mDeviceState = deviceState;
         mDragDistThreshold = context.getResources().getDimensionPixelSize(
                 R.dimen.gestures_onehanded_drag_threshold);
         mSquaredSlop = Utilities.squaredTouchSlop(context);
+        mDisplaySize = mDisplayHolder.getInfo().realSize;
+        mNavBarSize = ResourceUtils.getNavbarSize(NAVBAR_BOTTOM_GESTURE_SIZE,
+                mContext.getResources());
     }
 
     @Override
@@ -96,7 +108,8 @@
                                 mDownPos.x - mLastPos.x, mDownPos.y - mLastPos.y))
                                 || (mDeviceState.isOneHandedModeActive() && isValidExitAngle(
                                 mDownPos.x - mLastPos.x, mDownPos.y - mLastPos.y))) {
-                            mPassedSlop = true;
+                            // To avoid mis-trigger when motion not touch system gesture region.
+                            mPassedSlop = isInSystemGestureRegion(mLastPos);
                             setActive(ev);
                         } else {
                             mState = STATE_DELEGATE_ACTIVE;
@@ -154,6 +167,11 @@
         SystemUiProxy.INSTANCE.get(mContext).stopOneHandedMode();
     }
 
+    private boolean isInSystemGestureRegion(PointF lastPos) {
+        final int navBarUpperBound = mDisplaySize.y - mNavBarSize;
+        return mDeviceState.isGesturalNavMode() && lastPos.y > navBarUpperBound;
+    }
+
     private boolean isValidStartAngle(float deltaX, float deltaY) {
         final float angle = (float) Math.toDegrees(Math.atan2(deltaY, deltaX));
         return angle > -(ANGLE_MAX) && angle < -(ANGLE_MIN);
diff --git a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
index a762cb7..66c24c8 100644
--- a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
+++ b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
@@ -31,6 +31,7 @@
 
 import android.content.Context;
 import android.util.Log;
+import android.view.View;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.WorkerThread;
@@ -55,6 +56,7 @@
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.util.Executors;
 import com.android.launcher3.util.LogConfig;
+import com.android.systemui.shared.system.InteractionJankMonitorWrapper;
 import com.android.systemui.shared.system.SysUiStatsLog;
 
 import java.util.Optional;
@@ -94,7 +96,7 @@
 
     @Override
     protected StatsLogger createLogger() {
-        return new StatsCompatLogger();
+        return new StatsCompatLogger(mContext);
     }
 
     /**
@@ -136,6 +138,7 @@
 
         private static final ItemInfo DEFAULT_ITEM_INFO = new ItemInfo();
 
+        private Context mContext;
         private ItemInfo mItemInfo = DEFAULT_ITEM_INFO;
         private InstanceId mInstanceId = DEFAULT_INSTANCE_ID;
         private OptionalInt mRank = OptionalInt.empty();
@@ -147,6 +150,10 @@
         private Optional<String> mEditText = Optional.empty();
         private SliceItem mSliceItem;
 
+        StatsCompatLogger(Context context) {
+            mContext = context;
+        }
+
         @Override
         public StatsLogger withItemInfo(ItemInfo itemInfo) {
             if (mContainerInfo.isPresent()) {
@@ -220,7 +227,6 @@
             if (!Utilities.ATLEAST_R) {
                 return;
             }
-
             LauncherAppState appState = LauncherAppState.getInstanceNoCreate();
 
             if (mSliceItem != null) {
@@ -256,6 +262,26 @@
             }
         }
 
+        @Override
+        public void sendToInteractionJankMonitor(EventEnum event, View view) {
+            if (!(event instanceof LauncherEvent)) {
+                return;
+            }
+            switch ((LauncherEvent) event) {
+                case LAUNCHER_ALLAPPS_VERTICAL_SWIPE_BEGIN:
+                    InteractionJankMonitorWrapper.begin(
+                            view,
+                            InteractionJankMonitorWrapper.CUJ_ALL_APPS_SCROLL);
+                    break;
+                case LAUNCHER_ALLAPPS_VERTICAL_SWIPE_END:
+                    InteractionJankMonitorWrapper.end(
+                            InteractionJankMonitorWrapper.CUJ_ALL_APPS_SCROLL);
+                    break;
+                default:
+                    break;
+            }
+        }
+
         private LauncherAtom.ItemInfo applyOverwrites(LauncherAtom.ItemInfo atomInfo) {
             LauncherAtom.ItemInfo.Builder itemInfoBuilder = atomInfo.toBuilder();
 
diff --git a/robolectric_tests/resources/widgets_predication_update_task_data.txt b/robolectric_tests/resources/widgets_predication_update_task_data.txt
new file mode 100644
index 0000000..941d195
--- /dev/null
+++ b/robolectric_tests/resources/widgets_predication_update_task_data.txt
@@ -0,0 +1,24 @@
+# Model data used by WidgetsPredictionUpdateTasksTest
+
+classMap s com.android.launcher3.model.data.WorkspaceItemInfo
+classMap w com.android.launcher3.model.data.LauncherAppWidgetInfo
+
+# Items for the BgDataModel
+
+# App shortcuts
+bgItem s itemType=0 title=app1-class1 intent=component=app1/class1 id=1
+bgItem s itemType=0 title=app1-class2 intent=component=app1/class2 id=2
+bgItem s itemType=0 title=app2-class1 intent=component=app2/class1 id=3
+bgItem s itemType=0 title=app2-class2 intent=component=app2/class2 id=4
+
+# Promise icons for app3
+bgItem s itemType=0 status=2 title=app3-class1 intent=component=app3/class1 id=5
+bgItem s itemType=0 status=2 title=app3-class2 intent=component=app3/class2 id=6
+bgItem s itemType=1 status=1 title=app3-shrt intent=component=app3/class3 id=7
+
+# Promise icon for app4
+bgItem s itemType=1 status=1 title=app4-shrt intent=component=app4/class1 id=8
+
+# Widget
+bgItem w providerName=app4/provider1 id=9
+bgItem w providerName=app5/provider1 id=10
\ No newline at end of file
diff --git a/src/com/android/launcher3/LauncherSettings.java b/src/com/android/launcher3/LauncherSettings.java
index dfdc53c..22c257a 100644
--- a/src/com/android/launcher3/LauncherSettings.java
+++ b/src/com/android/launcher3/LauncherSettings.java
@@ -193,6 +193,7 @@
         public static final int CONTAINER_DESKTOP = -100;
         public static final int CONTAINER_HOTSEAT = -101;
         public static final int CONTAINER_PREDICTION = -102;
+        public static final int CONTAINER_WIDGETS_PREDICTION = -111;
         public static final int CONTAINER_HOTSEAT_PREDICTION = -103;
         public static final int CONTAINER_ALL_APPS = -104;
         public static final int CONTAINER_WIDGETS_TRAY = -105;
diff --git a/src/com/android/launcher3/allapps/AllAppsGridAdapter.java b/src/com/android/launcher3/allapps/AllAppsGridAdapter.java
index 5030c5e..bb175ea 100644
--- a/src/com/android/launcher3/allapps/AllAppsGridAdapter.java
+++ b/src/com/android/launcher3/allapps/AllAppsGridAdapter.java
@@ -240,20 +240,18 @@
         @Override
         public int getSpanSize(int position) {
             int viewType = mApps.getAdapterItems().get(position).viewType;
+            int totalSpans = mGridLayoutMgr.getSpanCount();
             if (isIconViewType(viewType)) {
-                return 1 * SPAN_MULTIPLIER;
+                return totalSpans / mAppsPerRow;
             } else if (mSearchAdapterProvider.isSearchView(viewType)) {
-                return mSearchAdapterProvider.getGridSpanSize(viewType, mAppsPerRow);
+                return totalSpans / mSearchAdapterProvider.getItemsPerRow(viewType, mAppsPerRow);
             } else {
                 // Section breaks span the full width
-                return mAppsPerRow * SPAN_MULTIPLIER;
+                return totalSpans;
             }
         }
     }
 
-    // multiplier to support adapter item column count that is not mAppsPerRow.
-    public static final int SPAN_MULTIPLIER = 3;
-
     private final BaseDraggingActivity mLauncher;
     private final LayoutInflater mLayoutInflater;
     private final AlphabeticalAppsList mApps;
@@ -285,14 +283,19 @@
 
         mOnIconClickListener = launcher.getItemOnClickListener();
 
-        setAppsPerRow(mLauncher.getDeviceProfile().inv.numAllAppsColumns);
-
         mSearchAdapterProvider = searchAdapterProvider;
+        setAppsPerRow(mLauncher.getDeviceProfile().inv.numAllAppsColumns);
     }
 
     public void setAppsPerRow(int appsPerRow) {
         mAppsPerRow = appsPerRow;
-        mGridLayoutMgr.setSpanCount(mAppsPerRow * SPAN_MULTIPLIER);
+        int totalSpans = mAppsPerRow;
+        for (int itemPerRow : mSearchAdapterProvider.getSupportedItemsPerRowArray()) {
+            if (totalSpans % itemPerRow != 0) {
+                totalSpans *= itemPerRow;
+            }
+        }
+        mGridLayoutMgr.setSpanCount(totalSpans);
     }
 
     /**
diff --git a/src/com/android/launcher3/allapps/AllAppsRecyclerView.java b/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
index ace9938..179cb77 100644
--- a/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
+++ b/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
@@ -19,6 +19,9 @@
 import static android.view.View.MeasureSpec.UNSPECIFIED;
 import static android.view.View.MeasureSpec.makeMeasureSpec;
 
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_VERTICAL_SWIPE_BEGIN;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_VERTICAL_SWIPE_END;
+
 import android.content.Context;
 import android.content.res.Resources;
 import android.graphics.Canvas;
@@ -36,6 +39,7 @@
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.R;
+import com.android.launcher3.logging.StatsLogManager;
 import com.android.launcher3.views.RecyclerViewFastScroller;
 
 import java.util.ArrayList;
@@ -176,6 +180,23 @@
     }
 
     @Override
+    public void onScrollStateChanged(int state) {
+        super.onScrollStateChanged(state);
+
+        StatsLogManager mgr = BaseDraggingActivity.fromContext(getContext()).getStatsLogManager();
+        switch (state) {
+            case SCROLL_STATE_DRAGGING:
+                mgr.logger().sendToInteractionJankMonitor(
+                        LAUNCHER_ALLAPPS_VERTICAL_SWIPE_BEGIN, this);
+                break;
+            case SCROLL_STATE_IDLE:
+                mgr.logger().sendToInteractionJankMonitor(
+                        LAUNCHER_ALLAPPS_VERTICAL_SWIPE_END, this);
+                break;
+        }
+    }
+
+    @Override
     public boolean onInterceptTouchEvent(MotionEvent e) {
         boolean result = super.onInterceptTouchEvent(e);
         if (!result && e.getAction() == MotionEvent.ACTION_DOWN
diff --git a/src/com/android/launcher3/allapps/search/SearchAdapterProvider.java b/src/com/android/launcher3/allapps/search/SearchAdapterProvider.java
index a79ec43..fdacd3d 100644
--- a/src/com/android/launcher3/allapps/search/SearchAdapterProvider.java
+++ b/src/com/android/launcher3/allapps/search/SearchAdapterProvider.java
@@ -51,10 +51,17 @@
             ViewGroup parent, int viewType);
 
     /**
+     * Returns supported item per row combinations supported
+     */
+    public int[] getSupportedItemsPerRowArray() {
+        return new int[]{};
+    }
+
+    /**
      * Returns how many cells a view should span
      */
-    public int getGridSpanSize(int viewType, int appsPerRow) {
-        return appsPerRow * AllAppsGridAdapter.SPAN_MULTIPLIER;
+    public int getItemsPerRow(int viewType, int appsPerRow) {
+        return appsPerRow;
     }
 
     /**
diff --git a/src/com/android/launcher3/logging/StatsLogManager.java b/src/com/android/launcher3/logging/StatsLogManager.java
index d554bb9..c5ffff9 100644
--- a/src/com/android/launcher3/logging/StatsLogManager.java
+++ b/src/com/android/launcher3/logging/StatsLogManager.java
@@ -21,6 +21,7 @@
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_OVERVIEW_GESTURE;
 
 import android.content.Context;
+import android.view.View;
 
 import androidx.annotation.Nullable;
 import androidx.slice.SliceItem;
@@ -407,6 +408,12 @@
 
         @UiEvent(doc = "User switched to AllApps Work tab by tapping on it.")
         LAUNCHER_ALLAPPS_TAP_ON_WORK_TAB(722),
+
+        @UiEvent(doc = "All apps vertical fling started.")
+        LAUNCHER_ALLAPPS_VERTICAL_SWIPE_BEGIN(724),
+
+        @UiEvent(doc = "All apps vertical fling ended.")
+        LAUNCHER_ALLAPPS_VERTICAL_SWIPE_END(725)
         ;
 
         // ADD MORE
@@ -525,6 +532,13 @@
          */
         default void log(EventEnum event) {
         }
+
+        /**
+         * Builds the final message and logs it to two different atoms, one for
+         * event tracking and the other for jank tracking.
+         */
+        default void sendToInteractionJankMonitor(EventEnum event, View v) {
+        }
     }
 
     /**
diff --git a/src/com/android/launcher3/model/BaseModelUpdateTask.java b/src/com/android/launcher3/model/BaseModelUpdateTask.java
index 543ee2a..ad553d5 100644
--- a/src/com/android/launcher3/model/BaseModelUpdateTask.java
+++ b/src/com/android/launcher3/model/BaseModelUpdateTask.java
@@ -124,7 +124,7 @@
 
     public void bindUpdatedWidgets(BgDataModel dataModel) {
         final ArrayList<WidgetsListBaseEntry> widgets =
-                dataModel.widgetsModel.getWidgetsList(mApp.getContext());
+                dataModel.widgetsModel.getWidgetsListForPicker(mApp.getContext());
         scheduleCallbackTask(c -> c.bindAllWidgets(widgets));
     }
 
diff --git a/src/com/android/launcher3/popup/PopupDataProvider.java b/src/com/android/launcher3/popup/PopupDataProvider.java
index 1853632..7780894 100644
--- a/src/com/android/launcher3/popup/PopupDataProvider.java
+++ b/src/com/android/launcher3/popup/PopupDataProvider.java
@@ -58,8 +58,11 @@
     private HashMap<ComponentKey, Integer> mDeepShortcutMap = new HashMap<>();
     /** Maps packages to their DotInfo's . */
     private Map<PackageUserKey, DotInfo> mPackageUserToDotInfos = new HashMap<>();
-    /** Maps packages to their Widgets */
+
+    /** All installed widgets. */
     private List<WidgetsListBaseEntry> mAllWidgets = List.of();
+    /** Widgets that can be recommended to the users. */
+    private List<ItemInfo> mRecommendedWidgets = List.of();
 
     private PopupDataChangeListener mChangeListener = PopupDataChangeListener.INSTANCE;
 
@@ -187,6 +190,15 @@
         notificationListener.cancelNotificationFromLauncher(notificationKey);
     }
 
+    /**
+     * Sets a list of recommended widgets ordered by their order of appearance in the widgets
+     * recommendation UI.
+     */
+    public void setRecommendedWidgets(List<ItemInfo> recommendedWidgets) {
+        mRecommendedWidgets = recommendedWidgets;
+        mChangeListener.onRecommendedWidgetsBound();
+    }
+
     public void setAllWidgets(List<WidgetsListBaseEntry> allWidgets) {
         mAllWidgets = allWidgets;
         mChangeListener.onWidgetsBound();
@@ -200,6 +212,21 @@
         return mAllWidgets;
     }
 
+    /** Returns a list of recommended widgets. */
+    public List<WidgetItem> getRecommendedWidgets() {
+        HashMap<ComponentKey, WidgetItem> allWidgetItems = new HashMap<>();
+        mAllWidgets.stream()
+                .filter(entry -> entry instanceof WidgetsListContentEntry)
+                .forEach(entry -> ((WidgetsListContentEntry) entry).mWidgets
+                        .forEach(widget -> allWidgetItems.put(
+                                new ComponentKey(widget.componentName, widget.user), widget)));
+        return mRecommendedWidgets.stream()
+                .map(recommendedWidget -> allWidgetItems.get(
+                        new ComponentKey(recommendedWidget.getTargetComponent(),
+                                recommendedWidget.user)))
+                .collect(Collectors.toList());
+    }
+
     public List<WidgetItem> getWidgetsForPackageUser(PackageUserKey packageUserKey) {
         return mAllWidgets.stream()
                 .filter(row -> row instanceof WidgetsListContentEntry
@@ -244,5 +271,8 @@
         default void trimNotifications(Map<PackageUserKey, DotInfo> updatedDots) { }
 
         default void onWidgetsBound() { }
+
+        /** A callback to get notified when recommended widgets are bound. */
+        default void onRecommendedWidgetsBound() { }
     }
 }
diff --git a/src_shortcuts_overrides/com/android/launcher3/model/LoaderResults.java b/src_shortcuts_overrides/com/android/launcher3/model/LoaderResults.java
index 73b1601..abce2a2 100644
--- a/src_shortcuts_overrides/com/android/launcher3/model/LoaderResults.java
+++ b/src_shortcuts_overrides/com/android/launcher3/model/LoaderResults.java
@@ -48,7 +48,7 @@
     @Override
     public void bindWidgets() {
         final List<WidgetsListBaseEntry> widgets =
-                mBgDataModel.widgetsModel.getWidgetsList(mApp.getContext());
+                mBgDataModel.widgetsModel.getWidgetsListForPicker(mApp.getContext());
         executeCallbacksTask(c -> c.bindAllWidgets(widgets), mUiExecutor);
     }
 }
diff --git a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
index a8c7c48..3ea4766 100644
--- a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
+++ b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
@@ -5,6 +5,8 @@
 
 import static com.android.launcher3.pm.ShortcutConfigActivityInfo.queryList;
 
+import static java.util.stream.Collectors.toList;
+
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ComponentName;
 import android.content.Context;
@@ -67,7 +69,7 @@
      *
      * @see com.android.launcher3.widget.picker.WidgetsListAdapter#setWidgets(List)
      */
-    public synchronized ArrayList<WidgetsListBaseEntry> getWidgetsList(Context context) {
+    public synchronized ArrayList<WidgetsListBaseEntry> getWidgetsListForPicker(Context context) {
         ArrayList<WidgetsListBaseEntry> result = new ArrayList<>();
         AlphabeticIndexCompat indexer = new AlphabeticIndexCompat(context);
 
@@ -82,6 +84,22 @@
         return result;
     }
 
+    /** Returns a mapping of packages to their widgets without static shortcuts. */
+    public synchronized Map<PackageUserKey, List<WidgetItem>> getAllWidgetsWithoutShortcuts() {
+        Map<PackageUserKey, List<WidgetItem>> packagesToWidgets = new HashMap<>();
+        mWidgetsList.forEach((packageItemInfo, widgetsAndShortcuts) -> {
+            List<WidgetItem> widgets = widgetsAndShortcuts.stream()
+                        .filter(item -> item.widgetInfo != null)
+                        .collect(toList());
+            if (widgets.size() > 0) {
+                packagesToWidgets.put(
+                        new PackageUserKey(packageItemInfo.packageName, packageItemInfo.user),
+                        widgets);
+            }
+        });
+        return packagesToWidgets;
+    }
+
     /**
      * @param packageUser If null, all widgets and shortcuts are updated and returned, otherwise
      *                    only widgets and shortcuts associated with the package/user are.
diff --git a/tests/src/com/android/launcher3/ui/WorkTabTest.java b/tests/src/com/android/launcher3/ui/WorkTabTest.java
index 63220ad..a6c7c03 100644
--- a/tests/src/com/android/launcher3/ui/WorkTabTest.java
+++ b/tests/src/com/android/launcher3/ui/WorkTabTest.java
@@ -18,6 +18,7 @@
 import static com.android.launcher3.LauncherState.ALL_APPS;
 import static com.android.launcher3.LauncherState.NORMAL;
 import static com.android.launcher3.allapps.AllAppsStore.DEFER_UPDATES_TEST;
+import static com.android.launcher3.tapl.LauncherInstrumentation.LONG_WAIT_TIME_MS;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
@@ -140,9 +141,13 @@
                 WorkEduView.KEY_WORK_EDU_STEP).remove(
                 WorkEduView.KEY_LEGACY_WORK_EDU_SEEN).commit());
 
-        waitForLauncherCondition("Work tab not setup",
-                launcher -> launcher.getAppsView().getContentView() instanceof AllAppsPagedView,
-                60000);
+        waitForLauncherCondition("Work tab not setup", launcher -> {
+            if (launcher.getAppsView().getContentView() instanceof AllAppsPagedView) {
+                launcher.getAppsView().getAppsStore().enableDeferUpdates(DEFER_UPDATES_TEST);
+                return true;
+            }
+            return false;
+        }, LONG_WAIT_TIME_MS);
 
         executeOnLauncher(launcher -> launcher.getStateManager().goToState(ALL_APPS));
         WorkEduView workEduView = getEduView();
@@ -153,16 +158,6 @@
             workEduView.findViewById(R.id.proceed).callOnClick();
         });
 
-        executeOnLauncher(launcher -> Log.d(TestProtocol.WORK_PROFILE_REMOVED,
-                "work profile status (" + mProfileUserId + ") :"
-                        + launcher.getAppsView().isWorkTabVisible()));
-
-
-        executeOnLauncher(launcher -> {
-            launcher.getAppsView().getAppsStore().enableDeferUpdates(DEFER_UPDATES_TEST);
-            Log.d(TestProtocol.WORK_PROFILE_REMOVED, "Defer all apps update");
-        });
-
         AtomicInteger attempt = new AtomicInteger(0);
         // verify work edu is seen next
         waitForLauncherCondition("Launcher did not show the next edu screen", l -> {
@@ -178,7 +173,7 @@
             }
             return ((TextView) workEduView.findViewById(R.id.content_text)).getText().equals(
                     l.getResources().getString(R.string.work_profile_edu_work_apps));
-        }, 60000);
+        });
     }
 
     @Test