Update the WidgetPickerActivity to display recommendations for hub host

- Accepts a ui_surface param of format "widgets{_hub}" and existing
widgets on the surface to be excluded from predictions
- Refactored the widgets prediction update task to extract reusable
logic that maps the predictions to widget items and reused it.

http://screencast/cast/NjE1MTA5MDI0NzU2NTMxMnwzMGE3NTMwNi1hZg

Bug: 326092660
Test: WidgetsPredictionHelperTest and see screencast above.
Flag: N/A
Change-Id: I6ceeb752c167893bab4ed496cedc5e8081e1b950
diff --git a/quickstep/src/com/android/launcher3/WidgetPickerActivity.java b/quickstep/src/com/android/launcher3/WidgetPickerActivity.java
index 8c4db4a..23cb8e9 100644
--- a/quickstep/src/com/android/launcher3/WidgetPickerActivity.java
+++ b/quickstep/src/com/android/launcher3/WidgetPickerActivity.java
@@ -35,23 +35,31 @@
 import android.view.WindowManager;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 
 import com.android.launcher3.dragndrop.SimpleDragLayer;
 import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.model.WidgetPredictionsRequester;
 import com.android.launcher3.model.WidgetsModel;
+import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.popup.PopupDataProvider;
+import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.widget.BaseWidgetSheet;
 import com.android.launcher3.widget.WidgetCell;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
 import com.android.launcher3.widget.picker.WidgetsFullSheet;
 
 import java.util.ArrayList;
+import java.util.List;
 import java.util.Locale;
+import java.util.Map;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 /** An Activity that can host Launcher's widget picker. */
 public class WidgetPickerActivity extends BaseActivity {
     private static final String TAG = "WidgetPickerActivity";
-
     /**
      * Name of the extra that indicates that a widget being dragged.
      *
@@ -64,14 +72,33 @@
     // the intent, then widgets will not be filtered for size.
     private static final String EXTRA_DESIRED_WIDGET_WIDTH = "desired_widget_width";
     private static final String EXTRA_DESIRED_WIDGET_HEIGHT = "desired_widget_height";
-
+    /**
+     * Widgets currently added by the user in the UI surface.
+     * <p>This allows widget picker to exclude existing widgets from suggestions.</p>
+     */
+    private static final String EXTRA_ADDED_APP_WIDGETS = "added_app_widgets";
+    /**
+     * A unique identifier of the surface hosting the widgets;
+     * <p>"widgets" is reserved for home screen surface.</p>
+     * <p>"widgets_hub" is reserved for glanceable hub surface.</p>
+     */
+    private static final String EXTRA_UI_SURFACE = "ui_surface";
+    private static final Pattern UI_SURFACE_PATTERN =
+            Pattern.compile("^(widgets|widgets_hub)$");
     private SimpleDragLayer<WidgetPickerActivity> mDragLayer;
     private WidgetsModel mModel;
+    private LauncherAppState mApp;
+    private WidgetPredictionsRequester mWidgetPredictionsRequester;
     private final PopupDataProvider mPopupDataProvider = new PopupDataProvider(i -> {});
 
     private int mDesiredWidgetWidth;
     private int mDesiredWidgetHeight;
     private int mWidgetCategoryFilter;
+    @Nullable
+    private String mUiSurface;
+    // Widgets existing on the host surface.
+    @NonNull
+    private List<AppWidgetProviderInfo> mAddedWidgets = new ArrayList<>();
 
     @Override
     protected void onCreate(Bundle savedInstanceState) {
@@ -80,9 +107,8 @@
         getWindow().addFlags(WindowManager.LayoutParams.FLAG_SHOW_WHEN_LOCKED);
         getWindow().clearFlags(WindowManager.LayoutParams.FLAG_SHOW_WALLPAPER);
 
-        LauncherAppState app = LauncherAppState.getInstance(this);
-        InvariantDeviceProfile idp = app.getInvariantDeviceProfile();
-
+        mApp = LauncherAppState.getInstance(this);
+        InvariantDeviceProfile idp = mApp.getInvariantDeviceProfile();
         mDeviceProfile = idp.getDeviceProfile(this);
         mModel = new WidgetsModel();
 
@@ -97,6 +123,11 @@
         widgetSheet.disableNavBarScrim(true);
         widgetSheet.addOnCloseListener(this::finish);
 
+        parseIntentExtras();
+        refreshAndBindWidgets();
+    }
+
+    private void parseIntentExtras() {
         // A value of 0 for either size means that no filtering will occur in that dimension. If
         // both values are 0, then no size filtering will occur.
         mDesiredWidgetWidth =
@@ -108,7 +139,15 @@
         mWidgetCategoryFilter =
                 getIntent().getIntExtra(AppWidgetManager.EXTRA_CATEGORY_FILTER, 0);
 
-        refreshAndBindWidgets();
+        String uiSurfaceParam = getIntent().getStringExtra(EXTRA_UI_SURFACE);
+        if (uiSurfaceParam != null && UI_SURFACE_PATTERN.matcher(uiSurfaceParam).matches()) {
+            mUiSurface = uiSurfaceParam;
+        }
+        ArrayList<AppWidgetProviderInfo> addedWidgets = getIntent().getParcelableArrayListExtra(
+                EXTRA_ADDED_APP_WIDGETS, AppWidgetProviderInfo.class);
+        if (addedWidgets != null) {
+            mAddedWidgets = addedWidgets;
+        }
     }
 
     @NonNull
@@ -179,11 +218,12 @@
         };
     }
 
+    /** Updates the model with widgets and provides them after applying the provided filter. */
     private void refreshAndBindWidgets() {
         MODEL_EXECUTOR.execute(() -> {
             LauncherAppState app = LauncherAppState.getInstance(this);
             mModel.update(app, null);
-            final ArrayList<WidgetsListBaseEntry> widgets =
+            final List<WidgetsListBaseEntry> allWidgets =
                     mModel.getFilteredWidgetsListForPicker(
                             app.getContext(),
                             /*widgetItemFilter=*/ widget -> {
@@ -193,10 +233,37 @@
                                 return verdict.isAcceptable;
                             }
                     );
-            MAIN_EXECUTOR.execute(() -> mPopupDataProvider.setAllWidgets(widgets));
+            bindWidgets(allWidgets);
+            if (mUiSurface != null) {
+                Map<PackageUserKey, List<WidgetItem>> allWidgetsMap = allWidgets.stream()
+                        .filter(WidgetsListHeaderEntry.class::isInstance)
+                        .collect(Collectors.toMap(
+                                entry -> PackageUserKey.fromPackageItemInfo(entry.mPkgItem),
+                                entry -> entry.mWidgets)
+                        );
+                mWidgetPredictionsRequester = new WidgetPredictionsRequester(app.getContext(),
+                        mUiSurface, allWidgetsMap);
+                mWidgetPredictionsRequester.request(mAddedWidgets, this::bindRecommendedWidgets);
+            }
         });
     }
 
+    private void bindWidgets(List<WidgetsListBaseEntry> widgets) {
+        MAIN_EXECUTOR.execute(() -> mPopupDataProvider.setAllWidgets(widgets));
+    }
+
+    private void bindRecommendedWidgets(List<ItemInfo> recommendedWidgets) {
+        MAIN_EXECUTOR.execute(() -> mPopupDataProvider.setRecommendedWidgets(recommendedWidgets));
+    }
+
+    @Override
+    protected void onDestroy() {
+        super.onDestroy();
+        if (mWidgetPredictionsRequester != null) {
+            mWidgetPredictionsRequester.clear();
+        }
+    }
+
     private WidgetAcceptabilityVerdict isWidgetAcceptable(WidgetItem widget) {
         final AppWidgetProviderInfo info = widget.widgetInfo;
         if (info == null) {
diff --git a/quickstep/src/com/android/launcher3/model/WidgetPredictionsRequester.java b/quickstep/src/com/android/launcher3/model/WidgetPredictionsRequester.java
new file mode 100644
index 0000000..8431396
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/model/WidgetPredictionsRequester.java
@@ -0,0 +1,233 @@
+/*
+ * Copyright (C) 2024 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.Flags.enableCategorizedWidgetSuggestions;
+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 android.app.prediction.AppPredictionContext;
+import android.app.prediction.AppPredictionManager;
+import android.app.prediction.AppPredictor;
+import android.app.prediction.AppTarget;
+import android.app.prediction.AppTargetEvent;
+import android.app.prediction.AppTargetId;
+import android.appwidget.AppWidgetProviderInfo;
+import android.content.ComponentName;
+import android.content.Context;
+import android.os.Bundle;
+import android.text.TextUtils;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
+import androidx.annotation.WorkerThread;
+
+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 com.android.launcher3.widget.picker.WidgetRecommendationCategoryProvider;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Consumer;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
+
+/**
+ * Works with app predictor to fetch and process widget predictions displayed in a standalone
+ * widget picker activity for a UI surface.
+ */
+public class WidgetPredictionsRequester {
+    private static final int NUM_OF_RECOMMENDED_WIDGETS_PREDICATION = 20;
+    private static final String BUNDLE_KEY_ADDED_APP_WIDGETS = "added_app_widgets";
+
+    @Nullable
+    private AppPredictor mAppPredictor;
+    private final Context mContext;
+    @NonNull
+    private final String mUiSurface;
+    @NonNull
+    private final Map<PackageUserKey, List<WidgetItem>> mAllWidgets;
+
+    public WidgetPredictionsRequester(Context context, @NonNull String uiSurface,
+            @NonNull Map<PackageUserKey, List<WidgetItem>> allWidgets) {
+        mContext = context;
+        mUiSurface = uiSurface;
+        mAllWidgets = Collections.unmodifiableMap(allWidgets);
+    }
+
+    /**
+     * Requests predictions from the app predictions manager and registers the provided callback to
+     * receive updates when predictions are available.
+     *
+     * @param existingWidgets widgets that are currently added to the surface;
+     * @param callback        consumer of prediction results to be called when predictions are
+     *                        available
+     */
+    public void request(List<AppWidgetProviderInfo> existingWidgets,
+            Consumer<List<ItemInfo>> callback) {
+        Bundle bundle = buildBundleForPredictionSession(existingWidgets, mUiSurface);
+        Predicate<WidgetItem> filter = notOnUiSurfaceFilter(existingWidgets);
+
+        MODEL_EXECUTOR.execute(() -> {
+            clear();
+            AppPredictionManager apm = mContext.getSystemService(AppPredictionManager.class);
+            if (apm == null) {
+                return;
+            }
+
+            mAppPredictor = apm.createAppPredictionSession(
+                    new AppPredictionContext.Builder(mContext)
+                            .setUiSurface(mUiSurface)
+                            .setExtras(bundle)
+                            .setPredictedTargetCount(NUM_OF_RECOMMENDED_WIDGETS_PREDICATION)
+                            .build());
+            mAppPredictor.registerPredictionUpdates(MODEL_EXECUTOR,
+                    targets -> bindPredictions(targets, filter, callback));
+            mAppPredictor.requestPredictionUpdate();
+        });
+    }
+
+    /**
+     * Returns a bundle that can be passed in a prediction session
+     *
+     * @param addedWidgets widgets that are already added by the user in the ui surface
+     * @param uiSurface    a unique identifier of the surface hosting widgets; format
+     *                     "widgets_xx"; note - "widgets" is reserved for home screen surface.
+     */
+    @VisibleForTesting
+    static Bundle buildBundleForPredictionSession(List<AppWidgetProviderInfo> addedWidgets,
+            String uiSurface) {
+        Bundle bundle = new Bundle();
+        ArrayList<AppTargetEvent> addedAppTargetEvents = new ArrayList<>();
+        for (AppWidgetProviderInfo info : addedWidgets) {
+            ComponentName componentName = info.provider;
+            AppTargetEvent appTargetEvent = buildAppTargetEvent(uiSurface, info, componentName);
+            addedAppTargetEvents.add(appTargetEvent);
+        }
+        bundle.putParcelableArrayList(BUNDLE_KEY_ADDED_APP_WIDGETS, addedAppTargetEvents);
+        return bundle;
+    }
+
+    /**
+     * Builds the AppTargetEvent for added widgets in a form that can be passed to the widget
+     * predictor.
+     * Also see {@link PredictionHelper}
+     */
+    private static AppTargetEvent buildAppTargetEvent(String uiSurface, AppWidgetProviderInfo info,
+            ComponentName componentName) {
+        AppTargetId appTargetId = new AppTargetId("widget:" + componentName.getPackageName());
+        AppTarget appTarget = new AppTarget.Builder(appTargetId, componentName.getPackageName(),
+                /*user=*/ info.getProfile()).setClassName(componentName.getClassName()).build();
+        return new AppTargetEvent.Builder(appTarget, AppTargetEvent.ACTION_PIN)
+                .setLaunchLocation(uiSurface).build();
+    }
+
+    /**
+     * Returns a filter to match {@link WidgetItem}s that don't exist on the UI surface.
+     */
+    @NonNull
+    @VisibleForTesting
+    static Predicate<WidgetItem> notOnUiSurfaceFilter(
+            List<AppWidgetProviderInfo> existingWidgets) {
+        Set<ComponentKey> existingComponentKeys = existingWidgets.stream().map(
+                widget -> new ComponentKey(widget.provider, widget.getProfile())).collect(
+                Collectors.toSet());
+        return widgetItem -> !existingComponentKeys.contains(widgetItem);
+    }
+
+    /** Provides the predictions returned by the predictor to the registered callback. */
+    @WorkerThread
+    private void bindPredictions(List<AppTarget> targets, Predicate<WidgetItem> filter,
+            Consumer<List<ItemInfo>> callback) {
+        List<WidgetItem> filteredPredictions = filterPredictions(targets, mAllWidgets, filter);
+        List<ItemInfo> mappedPredictions = mapWidgetItemsToItemInfo(filteredPredictions);
+
+        MAIN_EXECUTOR.execute(() -> callback.accept(mappedPredictions));
+    }
+
+    /**
+     * Applies the provided filter (e.g. widgets not on workspace) on the predictions returned by
+     * the predictor.
+     */
+    @VisibleForTesting
+    static List<WidgetItem> filterPredictions(List<AppTarget> predictions,
+            Map<PackageUserKey, List<WidgetItem>> allWidgets, Predicate<WidgetItem> filter) {
+        List<WidgetItem> servicePredictedItems = new ArrayList<>();
+        List<WidgetItem> localFilteredWidgets = new ArrayList<>();
+
+        for (AppTarget prediction : predictions) {
+            List<WidgetItem> widgetsInPackage = allWidgets.get(
+                    new PackageUserKey(prediction.getPackageName(), prediction.getUser()));
+            if (widgetsInPackage == null || widgetsInPackage.isEmpty()) {
+                continue;
+            }
+            String className = prediction.getClassName();
+            if (!TextUtils.isEmpty(className)) {
+                WidgetItem item = widgetsInPackage.stream()
+                        .filter(w -> className.equals(w.componentName.getClassName()))
+                        .filter(filter)
+                        .findFirst().orElse(null);
+                if (item != null) {
+                    servicePredictedItems.add(item);
+                    continue;
+                }
+            }
+            // No widget was added by the service, try local filtering
+            widgetsInPackage.stream().filter(filter).findFirst()
+                    .ifPresent(localFilteredWidgets::add);
+        }
+        if (servicePredictedItems.isEmpty()) {
+            servicePredictedItems.addAll(localFilteredWidgets);
+        }
+
+        return servicePredictedItems;
+    }
+
+    /**
+     * Converts the list of {@link WidgetItem}s to the list of {@link ItemInfo}s.
+     */
+    private List<ItemInfo> mapWidgetItemsToItemInfo(List<WidgetItem> widgetItems) {
+        List<ItemInfo> items;
+        if (enableCategorizedWidgetSuggestions()) {
+            WidgetRecommendationCategoryProvider categoryProvider =
+                    WidgetRecommendationCategoryProvider.newInstance(mContext);
+            items = widgetItems.stream()
+                    .map(it -> new PendingAddWidgetInfo(it.widgetInfo, CONTAINER_WIDGETS_PREDICTION,
+                            categoryProvider.getWidgetRecommendationCategory(mContext, it)))
+                    .collect(Collectors.toList());
+        } else {
+            items = widgetItems.stream().map(it -> new PendingAddWidgetInfo(it.widgetInfo,
+                    CONTAINER_WIDGETS_PREDICTION)).collect(Collectors.toList());
+        }
+        return items;
+    }
+
+    /** Cleans up any open prediction sessions. */
+    public void clear() {
+        if (mAppPredictor != null) {
+            mAppPredictor.destroy();
+            mAppPredictor = null;
+        }
+    }
+}
diff --git a/quickstep/tests/src/com/android/launcher3/model/WidgetsPredictionsRequesterTest.kt b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredictionsRequesterTest.kt
new file mode 100644
index 0000000..5c7b4ab
--- /dev/null
+++ b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredictionsRequesterTest.kt
@@ -0,0 +1,221 @@
+/*
+ * Copyright (C) 2024 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 android.app.prediction.AppTargetEvent
+import android.app.prediction.AppTargetId
+import android.appwidget.AppWidgetProviderInfo
+import android.content.ComponentName
+import android.content.Context
+import android.os.Process.myUserHandle
+import android.os.UserHandle
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.launcher3.DeviceProfile
+import com.android.launcher3.InvariantDeviceProfile
+import com.android.launcher3.LauncherAppState
+import com.android.launcher3.icons.IconCache
+import com.android.launcher3.model.WidgetPredictionsRequester.buildBundleForPredictionSession
+import com.android.launcher3.model.WidgetPredictionsRequester.filterPredictions
+import com.android.launcher3.model.WidgetPredictionsRequester.notOnUiSurfaceFilter
+import com.android.launcher3.util.ActivityContextWrapper
+import com.android.launcher3.util.PackageUserKey
+import com.android.launcher3.util.WidgetUtils.createAppWidgetProviderInfo
+import com.android.launcher3.widget.LauncherAppWidgetProviderInfo
+import com.google.common.truth.Truth.assertThat
+import java.util.function.Predicate
+import junit.framework.Assert.assertNotNull
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.MockitoAnnotations
+
+@RunWith(AndroidJUnit4::class)
+class WidgetsPredictionsRequesterTest {
+
+    private lateinit var mUserHandle: UserHandle
+    private lateinit var context: Context
+    private lateinit var deviceProfile: DeviceProfile
+    private lateinit var testInvariantProfile: InvariantDeviceProfile
+
+    private lateinit var widget1aInfo: AppWidgetProviderInfo
+    private lateinit var widget1bInfo: AppWidgetProviderInfo
+    private lateinit var widget2Info: AppWidgetProviderInfo
+
+    private lateinit var widgetItem1a: WidgetItem
+    private lateinit var widgetItem1b: WidgetItem
+    private lateinit var widgetItem2: WidgetItem
+
+    private lateinit var allWidgets: Map<PackageUserKey, List<WidgetItem>>
+
+    @Mock private lateinit var iconCache: IconCache
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+        mUserHandle = myUserHandle()
+        context = ActivityContextWrapper(ApplicationProvider.getApplicationContext())
+        testInvariantProfile = LauncherAppState.getIDP(context)
+        deviceProfile = testInvariantProfile.getDeviceProfile(context).copy(context)
+
+        widget1aInfo =
+            createAppWidgetProviderInfo(
+                ComponentName.createRelative(APP_1_PACKAGE_NAME, APP_1_PROVIDER_A_CLASS_NAME)
+            )
+        widget1bInfo =
+            createAppWidgetProviderInfo(
+                ComponentName.createRelative(APP_1_PACKAGE_NAME, APP_1_PROVIDER_B_CLASS_NAME)
+            )
+        widgetItem1a = createWidgetItem(widget1aInfo)
+        widgetItem1b = createWidgetItem(widget1bInfo)
+
+        widget2Info =
+            createAppWidgetProviderInfo(
+                ComponentName.createRelative(APP_2_PACKAGE_NAME, APP_2_PROVIDER_1_CLASS_NAME)
+            )
+        widgetItem2 = createWidgetItem(widget2Info)
+
+        allWidgets =
+            mapOf(
+                PackageUserKey(APP_1_PACKAGE_NAME, mUserHandle) to
+                    listOf(widgetItem1a, widgetItem1b),
+                PackageUserKey(APP_2_PACKAGE_NAME, mUserHandle) to listOf(widgetItem2),
+            )
+    }
+
+    @Test
+    fun buildBundleForPredictionSession_includesAddedAppWidgets() {
+        val existingWidgets = arrayListOf(widget1aInfo, widget1bInfo, widget2Info)
+
+        val bundle = buildBundleForPredictionSession(existingWidgets, TEST_UI_SURFACE)
+        val addedWidgetsBundleExtra =
+            bundle.getParcelableArrayList(BUNDLE_KEY_ADDED_APP_WIDGETS, AppTarget::class.java)
+
+        assertNotNull(addedWidgetsBundleExtra)
+        assertThat(addedWidgetsBundleExtra)
+            .containsExactly(
+                buildExpectedAppTargetEvent(
+                    /*pkg=*/ APP_1_PACKAGE_NAME,
+                    /*providerClassName=*/ APP_1_PROVIDER_A_CLASS_NAME,
+                    /*user=*/ mUserHandle
+                ),
+                buildExpectedAppTargetEvent(
+                    /*pkg=*/ APP_1_PACKAGE_NAME,
+                    /*providerClassName=*/ APP_1_PROVIDER_B_CLASS_NAME,
+                    /*user=*/ mUserHandle
+                ),
+                buildExpectedAppTargetEvent(
+                    /*pkg=*/ APP_2_PACKAGE_NAME,
+                    /*providerClassName=*/ APP_2_PROVIDER_1_CLASS_NAME,
+                    /*user=*/ mUserHandle
+                )
+            )
+    }
+
+    @Test
+    fun filterPredictions_notOnUiSurfaceFilter_returnsOnlyEligiblePredictions() {
+        val widgetsAlreadyOnSurface = arrayListOf(widget1bInfo)
+        val filter: Predicate<WidgetItem> = notOnUiSurfaceFilter(widgetsAlreadyOnSurface)
+
+        val predictions =
+            listOf(
+                // already on surface
+                AppTarget(
+                    AppTargetId(APP_1_PACKAGE_NAME),
+                    APP_1_PACKAGE_NAME,
+                    APP_1_PROVIDER_B_CLASS_NAME,
+                    mUserHandle
+                ),
+                // eligible
+                AppTarget(
+                    AppTargetId(APP_2_PACKAGE_NAME),
+                    APP_2_PACKAGE_NAME,
+                    APP_2_PROVIDER_1_CLASS_NAME,
+                    mUserHandle
+                )
+            )
+
+        // only 2 was eligible
+        assertThat(filterPredictions(predictions, allWidgets, filter)).containsExactly(widgetItem2)
+    }
+
+    @Test
+    fun filterPredictions_appPredictions_returnsWidgetFromPackage() {
+        val widgetsAlreadyOnSurface = arrayListOf(widget1bInfo)
+        val filter: Predicate<WidgetItem> = notOnUiSurfaceFilter(widgetsAlreadyOnSurface)
+
+        val predictions =
+            listOf(
+                AppTarget(
+                    AppTargetId(APP_1_PACKAGE_NAME),
+                    APP_1_PACKAGE_NAME,
+                    "$APP_1_PACKAGE_NAME.SomeActivity",
+                    mUserHandle
+                ),
+                AppTarget(
+                    AppTargetId(APP_2_PACKAGE_NAME),
+                    APP_2_PACKAGE_NAME,
+                    "$APP_2_PACKAGE_NAME.SomeActivity2",
+                    mUserHandle
+                ),
+            )
+
+        assertThat(filterPredictions(predictions, allWidgets, filter))
+            .containsExactly(widgetItem1a, widgetItem2)
+    }
+
+    private fun createWidgetItem(
+        providerInfo: AppWidgetProviderInfo,
+    ): WidgetItem {
+        val widgetInfo = LauncherAppWidgetProviderInfo.fromProviderInfo(context, providerInfo)
+        return WidgetItem(widgetInfo, testInvariantProfile, iconCache, context)
+    }
+
+    companion object {
+        const val TEST_UI_SURFACE = "widgets_test"
+        const val BUNDLE_KEY_ADDED_APP_WIDGETS = "added_app_widgets"
+
+        const val APP_1_PACKAGE_NAME = "com.example.app1"
+        const val APP_1_PROVIDER_A_CLASS_NAME = "app1Provider1"
+        const val APP_1_PROVIDER_B_CLASS_NAME = "app1Provider2"
+
+        const val APP_2_PACKAGE_NAME = "com.example.app2"
+        const val APP_2_PROVIDER_1_CLASS_NAME = "app2Provider1"
+
+        const val TEST_PACKAGE = "pkg"
+
+        private fun buildExpectedAppTargetEvent(
+            pkg: String,
+            providerClassName: String,
+            userHandle: UserHandle
+        ): AppTargetEvent {
+            val appTarget =
+                AppTarget.Builder(
+                        /*id=*/ AppTargetId("widget:$pkg"),
+                        /*packageName=*/ pkg,
+                        /*user=*/ userHandle
+                    )
+                    .setClassName(providerClassName)
+                    .build()
+            return AppTargetEvent.Builder(appTarget, AppTargetEvent.ACTION_PIN)
+                .setLaunchLocation(TEST_UI_SURFACE)
+                .build()
+        }
+    }
+}