Merge "Add a flag to enable custom local filter for recommended widgets" into sc-dev
diff --git a/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java b/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
index 70a143e..7c97b93 100644
--- a/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
+++ b/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
@@ -36,6 +36,7 @@
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.icons.ComponentWithLabel;
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
@@ -44,6 +45,7 @@
 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.shadows.ShadowDeviceFlag;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.ItemInfoMatcher;
@@ -60,6 +62,7 @@
 import org.mockito.MockitoAnnotations;
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.RuntimeEnvironment;
+import org.robolectric.shadow.api.Shadow;
 import org.robolectric.shadows.ShadowAppWidgetManager;
 import org.robolectric.shadows.ShadowPackageManager;
 import org.robolectric.util.ReflectionHelpers;
@@ -174,6 +177,41 @@
         assertWidgetInfo(recommendedWidgets.get(2).info, mApp1Provider1);
     }
 
+    @Test
+    public void widgetsRecommendationRan_localFilterDisabled_shouldReturnWidgetsInPredicationOrder()
+            throws Exception {
+        ShadowDeviceFlag shadowDeviceFlag = Shadow.extract(
+                FeatureFlags.ENABLE_LOCAL_RECOMMENDED_WIDGETS_FILTER);
+        shadowDeviceFlag.setValue(false);
+
+        // WHEN newPredicationTask is executed with 5 predicated widgets.
+        AppTarget widget1 = new AppTarget(new AppTargetId("app1"), "app1", "provider1",
+                mUserHandle);
+        AppTarget widget2 = new AppTarget(new AppTargetId("app1"), "app1", "provider2",
+                mUserHandle);
+        // Not installed app
+        AppTarget widget3 = new AppTarget(new AppTargetId("app2"), "app3", "provider1",
+                mUserHandle);
+        // Not installed widget
+        AppTarget widget4 = new AppTarget(new AppTargetId("app4"), "app4", "provider3",
+                mUserHandle);
+        AppTarget widget5 = new AppTarget(new AppTargetId("app5"), "app5", "provider1",
+                mUserHandle);
+        mModelHelper.executeTaskForTest(
+                newWidgetsPredicationTask(List.of(widget5, widget3, widget2, widget4, widget1)))
+                .forEach(Runnable::run);
+
+        // THEN only 3 widgets are returned because the launcher only filters out non-exist widgets.
+        List<PendingAddWidgetInfo> recommendedWidgets = mCallback.mRecommendedWidgets.items
+                .stream()
+                .map(itemInfo -> (PendingAddWidgetInfo) itemInfo)
+                .collect(Collectors.toList());
+        assertThat(recommendedWidgets).hasSize(3);
+        assertWidgetInfo(recommendedWidgets.get(0).info, mApp5Provider1);
+        assertWidgetInfo(recommendedWidgets.get(1).info, mApp1Provider2);
+        assertWidgetInfo(recommendedWidgets.get(2).info, mApp1Provider1);
+    }
+
     private void assertWidgetInfo(
             LauncherAppWidgetProviderInfo actual, AppWidgetProviderInfo expected) {
         assertThat(actual.provider).isEqualTo(expected.provider);
diff --git a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
index a29ac1a..22a8c9b 100644
--- a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
+++ b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
@@ -16,16 +16,17 @@
 package com.android.launcher3.model;
 
 import android.app.prediction.AppTarget;
+import android.content.ComponentName;
+import android.text.TextUtils;
 
 import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.config.FeatureFlags;
 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;
@@ -56,25 +57,43 @@
         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();
+
+        if (FeatureFlags.ENABLE_LOCAL_RECOMMENDED_WIDGETS_FILTER.get()) {
+            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.
+                        fixedContainerItems.items.add(
+                                new PendingAddWidgetInfo(notAddedWidgets.get(0).widgetInfo));
+                    }
+                }
+            }
+        } else {
+            Map<ComponentKey, WidgetItem> widgetItems =
+                    allWidgets.values().stream().flatMap(List::stream)
+                            .collect(Collectors.toMap(widget -> (ComponentKey) widget,
+                                    widget -> widget));
+            for (AppTarget app : mTargets) {
+                if (TextUtils.isEmpty(app.getClassName())) {
+                    continue;
+                }
+                ComponentKey targetWidget = new ComponentKey(
+                        new ComponentName(app.getPackageName(), app.getClassName()), app.getUser());
+                if (widgetItems.containsKey(targetWidget)) {
+                    fixedContainerItems.items.add(
+                            new PendingAddWidgetInfo(widgetItems.get(targetWidget).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/uioverrides/DeviceFlag.java b/quickstep/src/com/android/launcher3/uioverrides/DeviceFlag.java
index bb1f6fc..c115bbb 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/DeviceFlag.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/DeviceFlag.java
@@ -68,6 +68,12 @@
         mListeners.remove(r);
     }
 
+    @Override
+    public boolean get() {
+        // Override this method in order to let Robolectric ShadowDeviceFlag to stub it.
+        return super.get();
+    }
+
     private void registerDeviceConfigChangedListener(Context context) {
         DeviceConfig.addOnPropertiesChangedListener(
                 NAMESPACE_LAUNCHER,
diff --git a/robolectric_tests/src/com/android/launcher3/shadows/ShadowDeviceFlag.java b/robolectric_tests/src/com/android/launcher3/shadows/ShadowDeviceFlag.java
index 344f532..b58e4b7 100644
--- a/robolectric_tests/src/com/android/launcher3/shadows/ShadowDeviceFlag.java
+++ b/robolectric_tests/src/com/android/launcher3/shadows/ShadowDeviceFlag.java
@@ -18,11 +18,15 @@
 
 import android.content.Context;
 
+import androidx.annotation.Nullable;
+
 import com.android.launcher3.uioverrides.DeviceFlag;
 import com.android.launcher3.util.LooperExecutor;
 
 import org.robolectric.annotation.Implementation;
 import org.robolectric.annotation.Implements;
+import org.robolectric.annotation.RealObject;
+import org.robolectric.shadow.api.Shadow;
 
 /**
  * Shadow for {@link LooperExecutor} to provide reset functionality for static executors.
@@ -30,6 +34,9 @@
 @Implements(value = DeviceFlag.class, isInAndroidSdk = false)
 public class ShadowDeviceFlag {
 
+    @RealObject private DeviceFlag mRealObject;
+    @Nullable private Boolean mValue;
+
     /**
      * Mock change listener as it uses internal system classes not available to robolectric
      */
@@ -40,4 +47,16 @@
     protected static boolean getDeviceValue(String key, boolean defaultValue) {
         return defaultValue;
     }
+
+    @Implementation
+    public boolean get() {
+        if (mValue != null) {
+            return mValue;
+        }
+        return Shadow.directlyOn(mRealObject, DeviceFlag.class, "get");
+    }
+
+    public void setValue(boolean value) {
+        mValue = new Boolean(value);
+    }
 }
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index d935032..3b88a0b 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -230,6 +230,10 @@
     public static final BooleanFlag ENABLE_ENFORCED_ROUNDED_CORNERS = new DeviceFlag(
             "ENABLE_ENFORCED_ROUNDED_CORNERS", true, "Enforce rounded corners on all App Widgets");
 
+    public static final BooleanFlag ENABLE_LOCAL_RECOMMENDED_WIDGETS_FILTER = new DeviceFlag(
+            "ENABLE_LOCAL_RECOMMENDED_WIDGETS_FILTER", true,
+            "Enables a local filter for recommended widgets.");
+
     public static final BooleanFlag NOTIFY_CRASHES = getDebugFlag("NOTIFY_CRASHES", false,
             "Sends a notification whenever launcher encounters an uncaught exception.");