Merge "[automerger skipped] Merge Android 24Q1 Release (ab/11220357) am: 6810c19646 -s ours" into main
diff --git a/aconfig/launcher.aconfig b/aconfig/launcher.aconfig
index defb0e6..7032d5f 100644
--- a/aconfig/launcher.aconfig
+++ b/aconfig/launcher.aconfig
@@ -140,3 +140,10 @@
     description: "Tie unfold animation with state animation"
     bug: "297057373"
 }
+
+flag {
+  name: "enable_categorized_widget_suggestions"
+  namespace: "launcher"
+  description: "Enables widget suggestions in widget picker to be displayed in categories"
+  bug: "318410881"
+}
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index 0991fce..32d10b0 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -43,6 +43,7 @@
 import static com.android.launcher3.popup.SystemShortcut.DONT_SUGGEST_APP;
 import static com.android.launcher3.popup.SystemShortcut.INSTALL;
 import static com.android.launcher3.popup.SystemShortcut.PRIVATE_PROFILE_INSTALL;
+import static com.android.launcher3.popup.SystemShortcut.UNINSTALL_APP;
 import static com.android.launcher3.popup.SystemShortcut.WIDGETS;
 import static com.android.launcher3.taskbar.LauncherTaskbarUIController.ALL_APPS_PAGE_PROGRESS_INDEX;
 import static com.android.launcher3.taskbar.LauncherTaskbarUIController.MINUS_ONE_PAGE_PROGRESS_INDEX;
@@ -437,6 +438,9 @@
         if (Flags.enableShortcutDontSuggestApp()) {
             shortcuts.add(DONT_SUGGEST_APP);
         }
+        if (Flags.enablePrivateSpace()) {
+            shortcuts.add(UNINSTALL_APP);
+        }
         return shortcuts.stream();
     }
 
diff --git a/res/values/config.xml b/res/values/config.xml
index 5bdd7ebb..29c4e66 100644
--- a/res/values/config.xml
+++ b/res/values/config.xml
@@ -194,6 +194,11 @@
 
     <string-array name="filtered_components" ></string-array>
 
+    <!-- Widget component names to be included in weather category of widget suggestions. -->
+    <string-array name="weather_recommendations"></string-array>
+    <!-- Widget component names to be included in fitness category of widget suggestions. -->
+    <string-array name="fitness_recommendations"></string-array>
+
     <!-- Name of the class used to generate colors from the wallpaper colors. Must be implementing the LauncherAppWidgetHostView.ColorGenerator interface. -->
     <string name="color_generator_class" translatable="false"/>
 
@@ -252,6 +257,9 @@
     <!--  Used for custom widgets  -->
     <array name="custom_widget_providers"/>
 
+    <!--  Used for determining category of a widget presented in widget recommendations. -->
+    <string name="widget_recommendation_category_provider_class" translatable="false"></string>
+
     <!-- Embed parameters -->
     <dimen name="activity_split_ratio"  format="float">0.5</dimen>
     <integer name="min_width_split">720</integer>
diff --git a/res/values/strings.xml b/res/values/strings.xml
index b9fb024..5cc4616 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -71,6 +71,12 @@
     <!-- Widget suggestions header title in the full widgets picker for large screen devices
     in landscape mode. [CHAR_LIMIT=50] -->
     <string name="suggested_widgets_header_title">Suggestions</string>
+    <string name="productivity_widget_recommendation_category_label">Boost your day</string>
+    <string name="news_widget_recommendation_category_label">News For You</string>
+    <string name="social_and_entertainment_widget_recommendation_category_label">Your Chill Zone</string>
+    <string name="fitness_widget_recommendation_category_label">Reach Your Fitness Goals</string>
+    <string name="weather_widget_recommendation_category_label">Stay Ahead of the Weather</string>
+    <string name="others_widget_recommendation_category_label">You Might Also Like</string>
     <!-- Label for showing the number of widgets an app has in the full widgets picker.
          [CHAR_LIMIT=25][ICU SYNTAX] -->
     <string name="widgets_count">
diff --git a/src/com/android/launcher3/SecondaryDropTarget.java b/src/com/android/launcher3/SecondaryDropTarget.java
index 0b92c28..1362586 100644
--- a/src/com/android/launcher3/SecondaryDropTarget.java
+++ b/src/com/android/launcher3/SecondaryDropTarget.java
@@ -34,6 +34,8 @@
 import android.view.View;
 import android.widget.Toast;
 
+import androidx.annotation.Nullable;
+
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.dragndrop.DragOptions;
 import com.android.launcher3.logging.FileLog;
@@ -43,6 +45,7 @@
 import com.android.launcher3.logging.StatsLogManager.StatsLogger;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.ItemInfoWithIcon;
+import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.util.PackageManagerHelper;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 
@@ -176,6 +179,10 @@
         if (uninstallDisabled) {
             return INVALID;
         }
+        if (Flags.enablePrivateSpace() && UserCache.getInstance(getContext()).getUserInfo(
+                info.user).isPrivate()) {
+            return INVALID;
+        }
 
         if (info instanceof ItemInfoWithIcon) {
             ItemInfoWithIcon iconInfo = (ItemInfoWithIcon) info;
@@ -184,7 +191,7 @@
                 return INVALID;
             }
         }
-        if (getUninstallTarget(info) == null) {
+        if (getUninstallTarget(getContext(), info) == null) {
             return INVALID;
         }
         return UNINSTALL;
@@ -193,7 +200,7 @@
     /**
      * @return the component name that should be uninstalled or null.
      */
-    private ComponentName getUninstallTarget(ItemInfo item) {
+    public static ComponentName getUninstallTarget(Context context, ItemInfo item) {
         Intent intent = null;
         UserHandle user = null;
         if (item != null &&
@@ -202,7 +209,7 @@
             user = item.user;
         }
         if (intent != null) {
-            LauncherActivityInfo info = getContext().getSystemService(LauncherApps.class)
+            LauncherActivityInfo info = context.getSystemService(LauncherApps.class)
                     .resolveActivity(intent, user);
             if (info != null
                     && (info.getApplicationInfo().flags & ApplicationInfo.FLAG_SYSTEM) == 0) {
@@ -280,32 +287,41 @@
             if (FeatureFlags.ENABLE_DISMISS_PREDICTION_UNDO.get()) {
                 CharSequence announcement = getContext().getString(R.string.item_removed);
                 mDropTargetHandler
-                        .dismissPrediction(announcement, () -> {}, () -> {
-                    mStatsLogManager.logger()
-                            .withInstanceId(instanceId)
-                            .withItemInfo(info)
-                            .log(LAUNCHER_DISMISS_PREDICTION_UNDO);
-                });
+                        .dismissPrediction(announcement, () -> {
+                        }, () -> {
+                            mStatsLogManager.logger()
+                                    .withInstanceId(instanceId)
+                                    .withItemInfo(info)
+                                    .log(LAUNCHER_DISMISS_PREDICTION_UNDO);
+                        });
             }
             return null;
         }
 
-        ComponentName cn = getUninstallTarget(info);
+        return performUninstall(getContext(), getUninstallTarget(getContext(), info), info);
+    }
+
+    /**
+     * Performs uninstall and returns the target component for the {@link ItemInfo} or null if
+     * the uninstall was not performed.
+     */
+    public static ComponentName performUninstall(Context context, @Nullable ComponentName cn,
+            ItemInfo info) {
         if (cn == null) {
             // System applications cannot be installed. For now, show a toast explaining that.
             // We may give them the option of disabling apps this way.
             Toast.makeText(
-                    getContext(),
+                    context,
                     R.string.uninstall_system_app_text,
                     Toast.LENGTH_SHORT
-                ).show();
+            ).show();
             return null;
         }
         try {
-            Intent i = Intent.parseUri(getContext().getString(R.string.delete_package_intent), 0)
+            Intent i = Intent.parseUri(context.getString(R.string.delete_package_intent), 0)
                     .setData(Uri.fromParts("package", cn.getPackageName(), cn.getClassName()))
                     .putExtra(Intent.EXTRA_USER, info.user);
-            getContext().startActivity(i);
+            context.startActivity(i);
             FileLog.d(TAG, "start uninstall activity " + cn.getPackageName());
             return cn;
         } catch (URISyntaxException e) {
diff --git a/src/com/android/launcher3/logging/StatsLogManager.java b/src/com/android/launcher3/logging/StatsLogManager.java
index f9b7495..2a0f030 100644
--- a/src/com/android/launcher3/logging/StatsLogManager.java
+++ b/src/com/android/launcher3/logging/StatsLogManager.java
@@ -717,7 +717,10 @@
         LAUNCHER_PRIVATE_SPACE_INSTALL_APP_BUTTON_TAP(1605),
 
         @UiEvent(doc = "User attempted to create split screen with a widget")
-        LAUNCHER_SPLIT_WIDGET_ATTEMPT(1604)
+        LAUNCHER_SPLIT_WIDGET_ATTEMPT(1604),
+
+        @UiEvent(doc = "User tapped on private space uninstall system shortcut.")
+        LAUNCHER_PRIVATE_SPACE_UNINSTALL_SYSTEM_SHORTCUT_TAP(1608),
 
         // ADD MORE
         ;
diff --git a/src/com/android/launcher3/popup/SystemShortcut.java b/src/com/android/launcher3/popup/SystemShortcut.java
index 3030ed4..fbbfea9 100644
--- a/src/com/android/launcher3/popup/SystemShortcut.java
+++ b/src/com/android/launcher3/popup/SystemShortcut.java
@@ -1,6 +1,7 @@
 package com.android.launcher3.popup;
 
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_PRIVATE_SPACE_INSTALL_SYSTEM_SHORTCUT_TAP;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_PRIVATE_SPACE_UNINSTALL_SYSTEM_SHORTCUT_TAP;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_SYSTEM_SHORTCUT_APP_INFO_TAP;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_SYSTEM_SHORTCUT_DONT_SUGGEST_APP_TAP;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_SYSTEM_SHORTCUT_WIDGETS_TAP;
@@ -17,17 +18,21 @@
 import android.widget.ImageView;
 import android.widget.TextView;
 
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
 import com.android.launcher3.AbstractFloatingView;
 import com.android.launcher3.BaseDraggingActivity;
+import com.android.launcher3.Flags;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.R;
+import com.android.launcher3.SecondaryDropTarget;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.allapps.PrivateProfileManager;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
+import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.uioverrides.ApiWrapper;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.InstantAppResolver;
@@ -301,16 +306,11 @@
         }
     }
 
-    public static final Factory<Launcher> DONT_SUGGEST_APP = new Factory<Launcher>() {
-        @Nullable
-        @Override
-        public SystemShortcut<Launcher> getShortcut(Launcher activity, ItemInfo itemInfo,
-                View originalView) {
-            if (!itemInfo.isPredictedItem()) {
-                return null;
-            }
-            return new DontSuggestApp(activity, itemInfo, originalView);
+    public static final Factory<Launcher> DONT_SUGGEST_APP = (activity, itemInfo, originalView) -> {
+        if (!itemInfo.isPredictedItem()) {
+            return null;
         }
+        return new DontSuggestApp(activity, itemInfo, originalView);
     };
 
     private static class DontSuggestApp extends SystemShortcut<Launcher> {
@@ -329,6 +329,51 @@
         }
     }
 
+    public static final Factory<Launcher> UNINSTALL_APP = (activity, itemInfo, originalView) -> {
+        if (!Flags.enablePrivateSpace()) {
+            return null;
+        }
+        if (!UserCache.getInstance(activity.getApplicationContext()).getUserInfo(
+                itemInfo.user).isPrivate()) {
+            // If app is not Private Space app.
+            return null;
+        }
+        ComponentName cn = SecondaryDropTarget.getUninstallTarget(activity.getApplicationContext(),
+                itemInfo);
+        if (cn == null) {
+            // If component name is null, don't show uninstall shortcut.
+            // System apps will have component name as null.
+            return null;
+        }
+        return new UninstallApp(activity, itemInfo, originalView, cn);
+    };
+
+    private static class UninstallApp extends SystemShortcut<Launcher> {
+        private static final String TAG = "UninstallApp";
+        Context mContext;
+        @NonNull
+        ComponentName mComponentName;
+
+        UninstallApp(Launcher target, ItemInfo itemInfo, View originalView,
+                @NonNull ComponentName cn) {
+            super(R.drawable.ic_uninstall_no_shadow, R.string.uninstall_drop_target_label, target,
+                    itemInfo, originalView);
+            mContext = target.getApplicationContext();
+            mComponentName = cn;
+
+        }
+
+        @Override
+        public void onClick(View view) {
+            dismissTaskMenuView(mTarget);
+            SecondaryDropTarget.performUninstall(mContext, mComponentName, mItemInfo);
+            mTarget.getStatsLogManager()
+                    .logger()
+                    .withItemInfo(mItemInfo)
+                    .log(LAUNCHER_PRIVATE_SPACE_UNINSTALL_SYSTEM_SHORTCUT_TAP);
+        }
+    }
+
     public static <T extends Context & ActivityContext> void dismissTaskMenuView(T activity) {
         AbstractFloatingView.closeOpenViews(activity, true,
             AbstractFloatingView.TYPE_ALL & ~AbstractFloatingView.TYPE_REBIND_SAFE);
diff --git a/src/com/android/launcher3/widget/picker/WidgetRecommendationCategory.java b/src/com/android/launcher3/widget/picker/WidgetRecommendationCategory.java
new file mode 100644
index 0000000..072d1d5
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetRecommendationCategory.java
@@ -0,0 +1,62 @@
+/*
+ * 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.widget.picker;
+
+import androidx.annotation.Nullable;
+import androidx.annotation.StringRes;
+
+import java.util.Objects;
+
+/**
+ * A category of widget recommendations displayed in the widget picker (launched from "Widgets"
+ * option in the pop-up opened on long press of launcher workspace).
+ */
+public class WidgetRecommendationCategory implements Comparable<WidgetRecommendationCategory> {
+    /** Resource id that holds the user friendly label for the category. */
+    @StringRes
+    public final int categoryTitleRes;
+    /**
+     * Relative order of this category with respect to other categories.
+     *
+     * <p>Category with lowest order is displayed first in the recommendations section.</p>
+     */
+    public final int order;
+
+    public WidgetRecommendationCategory(@StringRes int categoryTitleRes, int order) {
+        this.categoryTitleRes = categoryTitleRes;
+        this.order = order;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(categoryTitleRes, order);
+    }
+
+    @Override
+    public boolean equals(@Nullable Object obj) {
+        if (!(obj instanceof WidgetRecommendationCategory category)) {
+            return false;
+        }
+        return categoryTitleRes == category.categoryTitleRes
+                && order == category.order;
+    }
+
+    @Override
+    public int compareTo(WidgetRecommendationCategory widgetRecommendationCategory) {
+        return order - widgetRecommendationCategory.order;
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProvider.java b/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProvider.java
new file mode 100644
index 0000000..d99c5d1
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProvider.java
@@ -0,0 +1,126 @@
+/*
+ * 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.widget.picker;
+
+import android.content.Context;
+import android.content.pm.ApplicationInfo;
+import android.content.pm.PackageManager;
+import android.util.Log;
+
+import androidx.annotation.WorkerThread;
+
+import com.android.launcher3.R;
+import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.util.Preconditions;
+import com.android.launcher3.util.ResourceBasedOverride;
+
+/**
+ * A {@link ResourceBasedOverride} that categorizes widget recommendations.
+ *
+ * <p>Override the {@code widget_recommendation_category_provider_class} resource to provide your
+ * own implementation. Method {@code getWidgetRecommendationCategory} is called per widget to get
+ * the category.</p>
+ */
+public class WidgetRecommendationCategoryProvider implements ResourceBasedOverride {
+    private static final String TAG = "WidgetRecommendationCategoryProvider";
+
+    /**
+     * Retrieve instance of this object that can be overridden in runtime based on the build
+     * variant of the application.
+     */
+    public static WidgetRecommendationCategoryProvider newInstance(Context context) {
+        Preconditions.assertWorkerThread();
+        return Overrides.getObject(
+                WidgetRecommendationCategoryProvider.class, context.getApplicationContext(),
+                R.string.widget_recommendation_category_provider_class);
+    }
+
+    /**
+     * Returns a {@link WidgetRecommendationCategory} for the provided widget item that can be used
+     * to display the recommendation grouped by categories.
+     */
+    @WorkerThread
+    public WidgetRecommendationCategory getWidgetRecommendationCategory(Context context,
+            WidgetItem item) {
+        // This is a default implementation that uses application category to derive the category to
+        // be displayed. The implementation can be overridden in individual launcher customization
+        // via the overridden WidgetRecommendationCategoryProvider resource.
+
+        Preconditions.assertWorkerThread();
+        PackageManager pm = context.getPackageManager();
+        if (item.widgetInfo != null && item.widgetInfo.getComponent() != null) {
+            String widgetComponentName = item.widgetInfo.getComponent().getClassName();
+            try {
+                int predictionCategory = pm.getApplicationInfo(
+                        item.widgetInfo.getComponent().getPackageName(), 0 /* flags */).category;
+                return getCategoryFromApplicationCategory(context, predictionCategory,
+                        widgetComponentName);
+            } catch (PackageManager.NameNotFoundException e) {
+                Log.e(TAG, "Failed to retrieve application category when determining the "
+                        + "widget category for " + widgetComponentName, e);
+            }
+        }
+        return null;
+    }
+
+    /** Maps application category to an appropriate displayable category. */
+    private static WidgetRecommendationCategory getCategoryFromApplicationCategory(
+            Context context, int applicationCategory, String componentName) {
+        if (applicationCategory == ApplicationInfo.CATEGORY_PRODUCTIVITY) {
+            return new WidgetRecommendationCategory(
+                    R.string.productivity_widget_recommendation_category_label, /*order=*/0);
+        }
+
+        if (applicationCategory == ApplicationInfo.CATEGORY_NEWS) {
+            return new WidgetRecommendationCategory(
+                    R.string.news_widget_recommendation_category_label, /*order=*/1);
+        }
+
+        if (applicationCategory == ApplicationInfo.CATEGORY_SOCIAL
+                || applicationCategory == ApplicationInfo.CATEGORY_AUDIO
+                || applicationCategory == ApplicationInfo.CATEGORY_VIDEO
+                || applicationCategory == ApplicationInfo.CATEGORY_IMAGE) {
+            return new WidgetRecommendationCategory(
+                    R.string.social_and_entertainment_widget_recommendation_category_label,
+                    /*order=*/4);
+        }
+
+        // Fitness & weather categories don't map to a specific application category, so, we
+        // maintain an allowlist.
+        String[] weatherRecommendationAllowlist =
+                context.getResources().getStringArray(R.array.weather_recommendations);
+        for (String allowedWeatherComponentName : weatherRecommendationAllowlist) {
+            if (componentName.equalsIgnoreCase(allowedWeatherComponentName)) {
+                return new WidgetRecommendationCategory(
+                        R.string.weather_widget_recommendation_category_label, /*order=*/2);
+            }
+        }
+
+        String[] fitnessRecommendationAllowlist =
+                context.getResources().getStringArray(R.array.fitness_recommendations);
+        for (String allowedFitnessComponentName : fitnessRecommendationAllowlist) {
+            if (componentName.equalsIgnoreCase(allowedFitnessComponentName)) {
+                return new WidgetRecommendationCategory(
+                        R.string.fitness_widget_recommendation_category_label, /*order=*/3);
+            }
+        }
+
+        return new WidgetRecommendationCategory(
+                R.string.others_widget_recommendation_category_label, /*order=*/5);
+    }
+
+}
diff --git a/tests/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProviderTest.java b/tests/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProviderTest.java
new file mode 100644
index 0000000..c807771
--- /dev/null
+++ b/tests/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProviderTest.java
@@ -0,0 +1,156 @@
+/*
+ * 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.widget.picker;
+
+import static android.content.pm.ApplicationInfo.CATEGORY_AUDIO;
+import static android.content.pm.ApplicationInfo.CATEGORY_IMAGE;
+import static android.content.pm.ApplicationInfo.CATEGORY_NEWS;
+import static android.content.pm.ApplicationInfo.CATEGORY_PRODUCTIVITY;
+import static android.content.pm.ApplicationInfo.CATEGORY_SOCIAL;
+import static android.content.pm.ApplicationInfo.CATEGORY_UNDEFINED;
+import static android.content.pm.ApplicationInfo.CATEGORY_VIDEO;
+
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.when;
+
+import android.appwidget.AppWidgetManager;
+import android.appwidget.AppWidgetProviderInfo;
+import android.content.ComponentName;
+import android.content.Context;
+import android.content.ContextWrapper;
+import android.content.pm.ApplicationInfo;
+import android.content.pm.PackageManager;
+import android.os.Process;
+
+import androidx.test.core.content.pm.ApplicationInfoBuilder;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.R;
+import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.model.WidgetItem;
+import com.android.launcher3.util.Executors;
+import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
+
+import com.google.common.collect.ImmutableMap;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.util.Map;
+
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class WidgetRecommendationCategoryProviderTest {
+    private static final String TEST_PACKAGE = "com.foo.test";
+    private static final String TEST_APP_NAME = "foo";
+    public static final WidgetRecommendationCategory SOCIAL_AND_ENTERTAINMENT_CATEGORY =
+            new WidgetRecommendationCategory(
+                    R.string.social_and_entertainment_widget_recommendation_category_label,
+                    /*order=*/4);
+    private final ApplicationInfo mTestAppInfo = ApplicationInfoBuilder.newBuilder().setPackageName(
+            TEST_PACKAGE).setName(TEST_APP_NAME).build();
+    private Context mContext;
+    @Mock
+    private IconCache mIconCache;
+
+    private WidgetItem mTestWidgetItem;
+    @Mock
+    private PackageManager mPackageManager;
+    private InvariantDeviceProfile mTestProfile;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mContext = new ContextWrapper(getInstrumentation().getTargetContext()) {
+            @Override
+            public PackageManager getPackageManager() {
+                return mPackageManager;
+            }
+        };
+        mTestProfile = new InvariantDeviceProfile();
+        mTestProfile.numRows = 5;
+        mTestProfile.numColumns = 5;
+        createTestWidgetItem();
+    }
+
+    @Test
+    public void getWidgetRecommendationCategory_returnsMappedCategory() throws Exception {
+        ImmutableMap<Integer, WidgetRecommendationCategory> testCategories = ImmutableMap.of(
+                CATEGORY_PRODUCTIVITY, new WidgetRecommendationCategory(
+                        R.string.productivity_widget_recommendation_category_label,
+                        /*order=*/
+                        0),
+                CATEGORY_NEWS, new WidgetRecommendationCategory(
+                        R.string.news_widget_recommendation_category_label, /*order=*/1),
+                CATEGORY_SOCIAL, SOCIAL_AND_ENTERTAINMENT_CATEGORY,
+                CATEGORY_AUDIO, SOCIAL_AND_ENTERTAINMENT_CATEGORY,
+                CATEGORY_IMAGE, SOCIAL_AND_ENTERTAINMENT_CATEGORY,
+                CATEGORY_VIDEO, SOCIAL_AND_ENTERTAINMENT_CATEGORY,
+                CATEGORY_UNDEFINED, new WidgetRecommendationCategory(
+                        R.string.others_widget_recommendation_category_label, /*order=*/5));
+
+        for (Map.Entry<Integer, WidgetRecommendationCategory> testCategory :
+                testCategories.entrySet()) {
+
+            mTestAppInfo.category = testCategory.getKey();
+            when(mPackageManager.getApplicationInfo(anyString(), anyInt())).thenReturn(
+                    mTestAppInfo);
+
+            WidgetRecommendationCategory category = Executors.MODEL_EXECUTOR.submit(() ->
+                    new WidgetRecommendationCategoryProvider().getWidgetRecommendationCategory(
+                            mContext,
+                            mTestWidgetItem)).get();
+
+            assertThat(category).isEqualTo(testCategory.getValue());
+        }
+    }
+
+    private void createTestWidgetItem() {
+        String widgetLabel = "Foo Widget";
+        String widgetClassName = ".mWidget";
+
+        doAnswer(invocation -> widgetLabel).when(mIconCache).getTitleNoCache(any());
+
+        AppWidgetProviderInfo providerInfo = AppWidgetManager.getInstance(getApplicationContext())
+                .getInstalledProvidersForPackage(
+                        getInstrumentation().getContext().getPackageName(), Process.myUserHandle())
+                .get(0);
+        providerInfo.provider = ComponentName.createRelative(TEST_PACKAGE, widgetClassName);
+
+        LauncherAppWidgetProviderInfo launcherAppWidgetProviderInfo =
+                LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, providerInfo);
+        launcherAppWidgetProviderInfo.spanX = 2;
+        launcherAppWidgetProviderInfo.spanY = 2;
+        launcherAppWidgetProviderInfo.label = widgetLabel;
+        mTestWidgetItem = new WidgetItem(launcherAppWidgetProviderInfo, mTestProfile, mIconCache,
+                mContext
+        );
+    }
+}