diff --git a/go/src/com/android/launcher3/model/WidgetsModel.java b/go/src/com/android/launcher3/model/WidgetsModel.java
index f8448da..1aa5d03 100644
--- a/go/src/com/android/launcher3/model/WidgetsModel.java
+++ b/go/src/com/android/launcher3/model/WidgetsModel.java
@@ -79,12 +79,13 @@
     }
 
     public WidgetItem getWidgetProviderInfoByProviderName(
-            ComponentName providerName) {
+            ComponentName providerName, UserHandle user) {
         return null;
     }
 
     /** Returns {@link PackageItemInfo} of a pending widget. */
-    public static PackageItemInfo newPendingItemInfo(ComponentName provider) {
-        return new PackageItemInfo(provider.getPackageName());
+    public static PackageItemInfo newPendingItemInfo(
+            Context context, ComponentName provider, UserHandle userHandle) {
+        return new PackageItemInfo(provider.getPackageName(), userHandle);
     }
 }
\ No newline at end of file
diff --git a/res/values/attrs.xml b/res/values/attrs.xml
index 720d736..a9493dc 100644
--- a/res/values/attrs.xml
+++ b/res/values/attrs.xml
@@ -258,4 +258,18 @@
     <declare-styleable name="WidgetsListRowHeader">
         <attr name="appIconSize" format="dimension" />
     </declare-styleable>
+
+    <declare-styleable name="WidgetSections">
+        <!-- Component name of an app widget provider. -->
+        <attr name="provider" format="string" />
+        <!-- If true, keep the app widget under its app listing in addition to the widget category
+             in the widget picker. Defaults to false if not specified. -->
+        <attr name="alsoKeepInApp" format="boolean" />
+        <!-- The category of an app widget provider. Defaults to -1 if not specified. -->
+        <attr name="category" format="integer" />
+        <!-- The title name of a widget category. -->
+        <attr name="sectionTitle" format="reference" />
+        <!-- The icon drawable of a widget category. -->
+        <attr name="sectionDrawable" format="reference" />
+    </declare-styleable>
 </resources>
diff --git a/res/xml/widget_sections.xml b/res/xml/widget_sections.xml
new file mode 100644
index 0000000..d755de6
--- /dev/null
+++ b/res/xml/widget_sections.xml
@@ -0,0 +1,25 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+     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.
+-->
+
+<widget-sections xmlns:launcher="http://schemas.android.com/apk/res-auto">
+    <section
+        launcher:category="0"
+        launcher:sectionDrawable="@drawable/ic_conversations_widget_category"
+        launcher:sectionTitle="@string/widget_category_conversations">
+        <widget launcher:provider="com.android.systemui/.people.widget.PeopleSpaceWidgetProvider" />
+    </section>
+</widget-sections>
\ No newline at end of file
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index b4907e8..57bb909 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -349,7 +349,7 @@
 
     private void inflateAndAddWidgets(LauncherAppWidgetInfo info, WidgetsModel widgetsModel) {
         WidgetItem widgetItem = widgetsModel.getWidgetProviderInfoByProviderName(
-                info.providerName);
+                info.providerName, info.user);
         if (widgetItem == null) {
             return;
         }
diff --git a/src/com/android/launcher3/icons/IconCache.java b/src/com/android/launcher3/icons/IconCache.java
index 60d6e83..936eeb9 100644
--- a/src/com/android/launcher3/icons/IconCache.java
+++ b/src/com/android/launcher3/icons/IconCache.java
@@ -18,6 +18,7 @@
 
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
+import static com.android.launcher3.widget.WidgetSections.NO_CATEGORY;
 
 import static java.util.stream.Collectors.groupingBy;
 
@@ -46,7 +47,6 @@
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherFiles;
-import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.icons.ComponentWithLabel.ComponentCachingLogic;
@@ -63,6 +63,8 @@
 import com.android.launcher3.util.InstantAppResolver;
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.util.Preconditions;
+import com.android.launcher3.widget.WidgetSections;
+import com.android.launcher3.widget.WidgetSections.WidgetSection;
 
 import java.util.Collections;
 import java.util.List;
@@ -275,7 +277,8 @@
             getTitleAndIcon(appInfo, false);
             return appInfo.bitmap;
         } else {
-            PackageItemInfo pkgInfo = new PackageItemInfo(shortcutInfo.getPackage());
+            PackageItemInfo pkgInfo = new PackageItemInfo(shortcutInfo.getPackage(),
+                    shortcutInfo.getUserHandle());
             getTitleAndIconForApp(pkgInfo, false);
             return pkgInfo.bitmap;
         }
@@ -409,8 +412,10 @@
         CacheEntry entry = getEntryForPackageLocked(
                 infoInOut.packageName, infoInOut.user, useLowResIcon);
         applyCacheEntry(entry, infoInOut);
-        if (infoInOut.category == PackageItemInfo.CONVERSATIONS) {
-            infoInOut.title = mContext.getString(R.string.widget_category_conversations);
+        if (infoInOut.widgetCategory != NO_CATEGORY) {
+            WidgetSection widgetSection = WidgetSections.getWidgetSections(mContext)
+                    .get(infoInOut.widgetCategory);
+            infoInOut.title = mContext.getString(widgetSection.mSectionTitle);
             infoInOut.contentDescription = mPackageManager.getUserBadgedLabel(
                     infoInOut.title, infoInOut.user);
         }
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 1249606..976c1e9 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -804,8 +804,9 @@
                                 if (appWidgetInfo.restoreStatus !=
                                         LauncherAppWidgetInfo.RESTORE_COMPLETED) {
                                     appWidgetInfo.pendingItemInfo = WidgetsModel.newPendingItemInfo(
-                                            appWidgetInfo.providerName);
-                                    appWidgetInfo.pendingItemInfo.user = appWidgetInfo.user;
+                                            mApp.getContext(),
+                                            appWidgetInfo.providerName,
+                                            appWidgetInfo.user);
                                     mIconCache.getTitleAndIconForApp(
                                             appWidgetInfo.pendingItemInfo, false);
                                 }
diff --git a/src/com/android/launcher3/model/data/PackageItemInfo.java b/src/com/android/launcher3/model/data/PackageItemInfo.java
index a81fe6a..0055763 100644
--- a/src/com/android/launcher3/model/data/PackageItemInfo.java
+++ b/src/com/android/launcher3/model/data/PackageItemInfo.java
@@ -16,47 +16,41 @@
 
 package com.android.launcher3.model.data;
 
-import androidx.annotation.IntDef;
+import static com.android.launcher3.widget.WidgetSections.NO_CATEGORY;
+
+import android.os.UserHandle;
 
 import com.android.launcher3.LauncherSettings;
 
-import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
 import java.util.Objects;
 
 /**
  * Represents a {@link Package} in the widget tray section.
  */
 public class PackageItemInfo extends ItemInfoWithIcon {
-    @Retention(RetentionPolicy.SOURCE)
-    @IntDef({NO_CATEGORY, CONVERSATIONS})
-    public @interface Category{}
-    /** The package is not categorized in the widget tray. */
-    public static final int NO_CATEGORY = 0;
-    /** The package is categorized to conversations widget in the widget tray. */
-    public static final int CONVERSATIONS = 1;
-
     /**
      * Package name of the {@link PackageItemInfo}.
      */
     public final String packageName;
 
     /** Represents a widget category shown in the widget tray section. */
-    @Category public final int category;
+    public final int widgetCategory;
 
-    public PackageItemInfo(String packageName) {
-        this(packageName, NO_CATEGORY);
+    public PackageItemInfo(String packageName, UserHandle user) {
+        this(packageName, NO_CATEGORY, user);
     }
 
-    public PackageItemInfo(String packageName, @Category int category) {
+    public PackageItemInfo(String packageName, int widgetCategory, UserHandle user) {
         this.packageName = packageName;
-        this.category = category;
+        this.widgetCategory = widgetCategory;
+        this.user = user;
         this.itemType = LauncherSettings.Favorites.ITEM_TYPE_NON_ACTIONABLE;
     }
 
     public PackageItemInfo(PackageItemInfo copy) {
         this.packageName = copy.packageName;
-        this.category = copy.category;
+        this.widgetCategory = copy.widgetCategory;
+        this.user = copy.user;
         this.itemType = LauncherSettings.Favorites.ITEM_TYPE_NON_ACTIONABLE;
     }
 
@@ -75,11 +69,13 @@
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         PackageItemInfo that = (PackageItemInfo) o;
-        return Objects.equals(packageName, that.packageName);
+        return Objects.equals(packageName, that.packageName)
+                && Objects.equals(user, that.user)
+                && widgetCategory == that.widgetCategory;
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(packageName, user);
+        return Objects.hash(packageName, user, widgetCategory);
     }
 }
diff --git a/src/com/android/launcher3/util/PackageUserKey.java b/src/com/android/launcher3/util/PackageUserKey.java
index 3a3b5a2..92d9737 100644
--- a/src/com/android/launcher3/util/PackageUserKey.java
+++ b/src/com/android/launcher3/util/PackageUserKey.java
@@ -1,19 +1,24 @@
 package com.android.launcher3.util;
 
+import static com.android.launcher3.widget.WidgetSections.NO_CATEGORY;
+
 import android.os.UserHandle;
 import android.service.notification.StatusBarNotification;
+import android.text.TextUtils;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
 import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.model.data.PackageItemInfo;
 
-import java.util.Arrays;
+import java.util.Objects;
 
-/** Creates a hash key based on package name and user. */
+/** Creates a hash key based on package name, widget category, and user. */
 public class PackageUserKey {
 
     public String mPackageName;
+    public int mWidgetCategory;
     public UserHandle mUser;
     private int mHashCode;
 
@@ -27,14 +32,31 @@
         return new PackageUserKey(notification.getPackageName(), notification.getUser());
     }
 
+    /** Creates a {@link PackageUserKey} from {@link PackageItemInfo}. */
+    public static PackageUserKey fromPackageItemInfo(PackageItemInfo info) {
+        if (TextUtils.isEmpty(info.packageName) && info.widgetCategory != NO_CATEGORY) {
+            return new PackageUserKey(info.widgetCategory, info.user);
+        }
+        return new PackageUserKey(info.packageName, info.user);
+    }
+
     public PackageUserKey(String packageName, UserHandle user) {
         update(packageName, user);
     }
 
+    public PackageUserKey(int widgetCategory, UserHandle user) {
+        update(/* packageName= */ "", widgetCategory, user);
+    }
+
     public void update(String packageName, UserHandle user) {
+        update(packageName, NO_CATEGORY, user);
+    }
+
+    private void update(String packageName, int widgetCategory, UserHandle user) {
         mPackageName = packageName;
+        mWidgetCategory = widgetCategory;
         mUser = user;
-        mHashCode = Arrays.hashCode(new Object[] {packageName, user});
+        mHashCode = Objects.hash(packageName, widgetCategory, user);
     }
 
     /**
@@ -59,12 +81,14 @@
     public boolean equals(Object obj) {
         if (!(obj instanceof PackageUserKey)) return false;
         PackageUserKey otherKey = (PackageUserKey) obj;
-        return mPackageName.equals(otherKey.mPackageName) && mUser.equals(otherKey.mUser);
+        return Objects.equals(mPackageName, otherKey.mPackageName)
+                && mWidgetCategory == otherKey.mWidgetCategory
+                && Objects.equals(mUser, otherKey.mUser);
     }
 
     @NonNull
     @Override
     public String toString() {
-        return mPackageName + "#" + mUser;
+        return mPackageName + "#" + mUser + ",category=" + mWidgetCategory;
     }
 }
diff --git a/src/com/android/launcher3/widget/PendingAppWidgetHostView.java b/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
index 57a6d3f..286305f 100644
--- a/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
+++ b/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
@@ -17,7 +17,7 @@
 package com.android.launcher3.widget;
 
 import static com.android.launcher3.graphics.PreloadIconDrawable.newPendingIcon;
-import static com.android.launcher3.model.data.PackageItemInfo.CONVERSATIONS;
+import static com.android.launcher3.widget.WidgetSections.getWidgetSections;
 
 import android.content.Context;
 import android.graphics.Canvas;
@@ -89,8 +89,8 @@
         setOnClickListener(ItemClickHandler.INSTANCE);
 
         if (info.pendingItemInfo == null) {
-            info.pendingItemInfo = new PackageItemInfo(info.providerName.getPackageName());
-            info.pendingItemInfo.user = info.user;
+            info.pendingItemInfo = new PackageItemInfo(info.providerName.getPackageName(),
+                    info.user);
             cache.updateIconInBackground(this, info.pendingItemInfo);
         } else {
             reapplyItemInfo(info.pendingItemInfo);
@@ -338,10 +338,11 @@
      */
     @Nullable
     private Drawable getWidgetCategoryIcon() {
-        switch (mInfo.pendingItemInfo.category) {
-            case CONVERSATIONS:
-                return getContext().getDrawable(R.drawable.ic_conversations_widget_category);
+        if (mInfo.pendingItemInfo.widgetCategory == WidgetSections.NO_CATEGORY) {
+            return null;
         }
-        return null;
+        Context context = getContext();
+        return context.getDrawable(getWidgetSections(context).get(
+                mInfo.pendingItemInfo.widgetCategory).mSectionDrawable);
     }
 }
diff --git a/src/com/android/launcher3/widget/WidgetSections.java b/src/com/android/launcher3/widget/WidgetSections.java
new file mode 100644
index 0000000..c45b095
--- /dev/null
+++ b/src/com/android/launcher3/widget/WidgetSections.java
@@ -0,0 +1,134 @@
+/*
+ * 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.widget;
+
+import static android.content.res.Resources.ID_NULL;
+
+import android.content.ComponentName;
+import android.content.Context;
+import android.content.res.TypedArray;
+import android.content.res.XmlResourceParser;
+import android.util.ArrayMap;
+import android.util.AttributeSet;
+import android.util.SparseArray;
+import android.util.Xml;
+
+import androidx.annotation.DrawableRes;
+import androidx.annotation.StringRes;
+
+import com.android.launcher3.R;
+import com.android.launcher3.util.IntSet;
+
+import org.xmlpull.v1.XmlPullParser;
+import org.xmlpull.v1.XmlPullParserException;
+
+import java.io.IOException;
+import java.util.Map;
+
+/** A helper class to parse widget sections (categories) resource overlay. */
+public final class WidgetSections {
+    /** The package is not categorized in the widget tray. */
+    public static final int NO_CATEGORY = -1;
+
+    private static final String TAG_SECTION_NAME = "section";
+    private static final String TAG_WIDGET_NAME = "widget";
+
+    private static SparseArray<WidgetSection> sWidgetSections;
+    private static Map<ComponentName, IntSet> sWidgetsToCategories;
+
+    /** Returns a list of widget sections that are shown in the widget picker. */
+    public static synchronized SparseArray<WidgetSection> getWidgetSections(Context context) {
+        if (sWidgetSections != null) {
+            return sWidgetSections;
+        }
+        parseWidgetSectionsXml(context);
+        return sWidgetSections;
+    }
+
+    /** Returns a map which maps app widget providers to app widget categories. */
+    public static synchronized Map<ComponentName, IntSet> getWidgetsToCategory(
+            Context context) {
+        if (sWidgetsToCategories != null) {
+            return sWidgetsToCategories;
+        }
+        parseWidgetSectionsXml(context);
+        return sWidgetsToCategories;
+    }
+
+    private static synchronized void parseWidgetSectionsXml(Context context) {
+        SparseArray<WidgetSection> widgetSections = new SparseArray();
+        Map<ComponentName, IntSet> widgetsToCategories = new ArrayMap<>();
+        try (XmlResourceParser parser = context.getResources().getXml(R.xml.widget_sections)) {
+            final int depth = parser.getDepth();
+            int type;
+            while (((type = parser.next()) != XmlPullParser.END_TAG
+                    || parser.getDepth() > depth) && type != XmlPullParser.END_DOCUMENT) {
+                if ((type == XmlPullParser.START_TAG)
+                        && TAG_SECTION_NAME.equals(parser.getName())) {
+                    AttributeSet sectionAttributes = Xml.asAttributeSet(parser);
+                    WidgetSection section = new WidgetSection(context, sectionAttributes);
+                    final int sectionDepth = parser.getDepth();
+                    while (((type = parser.next()) != XmlPullParser.END_TAG
+                                    || parser.getDepth() > sectionDepth)
+                            && type != XmlPullParser.END_DOCUMENT) {
+                        if ((type == XmlPullParser.START_TAG)
+                                && TAG_WIDGET_NAME.equals(parser.getName())) {
+                            TypedArray a = context.obtainStyledAttributes(
+                                    Xml.asAttributeSet(parser), R.styleable.WidgetSections);
+                            ComponentName provider = ComponentName.unflattenFromString(
+                                    a.getString(R.styleable.WidgetSections_provider));
+                            boolean alsoKeepInApp = a.getBoolean(
+                                    R.styleable.WidgetSections_alsoKeepInApp,
+                                    /* defValue= */ false);
+                            final IntSet categories;
+                            if (widgetsToCategories.containsKey(provider)) {
+                                categories = widgetsToCategories.get(provider);
+                            } else {
+                                categories = new IntSet();
+                                widgetsToCategories.put(provider, categories);
+                            }
+                            if (alsoKeepInApp) {
+                                categories.add(NO_CATEGORY);
+                            }
+                            categories.add(section.mCategory);
+                        }
+                    }
+                    widgetSections.put(section.mCategory, section);
+                }
+            }
+            sWidgetSections = widgetSections;
+            sWidgetsToCategories = widgetsToCategories;
+        } catch (IOException | XmlPullParserException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    /** A data class which contains a widget section's information. */
+    public static final class WidgetSection {
+        public final int mCategory;
+        @StringRes
+        public final int mSectionTitle;
+        @DrawableRes
+        public final int mSectionDrawable;
+
+        public WidgetSection(Context context, AttributeSet attrs) {
+            TypedArray a = context.obtainStyledAttributes(attrs, R.styleable.WidgetSections);
+            mCategory = a.getInt(R.styleable.WidgetSections_category, NO_CATEGORY);
+            mSectionTitle = a.getResourceId(R.styleable.WidgetSections_sectionTitle, ID_NULL);
+            mSectionDrawable = a.getResourceId(R.styleable.WidgetSections_sectionDrawable, ID_NULL);
+        }
+    }
+}
diff --git a/src/com/android/launcher3/widget/model/WidgetListSpaceEntry.java b/src/com/android/launcher3/widget/model/WidgetListSpaceEntry.java
index e62425f..7f24905 100644
--- a/src/com/android/launcher3/widget/model/WidgetListSpaceEntry.java
+++ b/src/com/android/launcher3/widget/model/WidgetListSpaceEntry.java
@@ -16,6 +16,8 @@
 
 package com.android.launcher3.widget.model;
 
+import android.os.Process;
+
 import com.android.launcher3.model.data.PackageItemInfo;
 
 import java.util.Collections;
@@ -26,7 +28,9 @@
 public class WidgetListSpaceEntry extends WidgetsListBaseEntry {
 
     public WidgetListSpaceEntry() {
-        super(new PackageItemInfo(""), "", Collections.EMPTY_LIST);
+        super(new PackageItemInfo(/* packageName= */ "", Process.myUserHandle()),
+                /* titleSectionName= */ "",
+                Collections.EMPTY_LIST);
         mPkgItem.title = "";
     }
 
diff --git a/src/com/android/launcher3/widget/picker/OnHeaderClickListener.java b/src/com/android/launcher3/widget/picker/OnHeaderClickListener.java
index 7372751..35f11bd 100644
--- a/src/com/android/launcher3/widget/picker/OnHeaderClickListener.java
+++ b/src/com/android/launcher3/widget/picker/OnHeaderClickListener.java
@@ -24,5 +24,5 @@
     /**
      * Calls when a header is clicked to show / hide widgets for a package.
      */
-    void onHeaderClicked(boolean showWidgets, PackageUserKey packageUserKey);
+    void onHeaderClicked(boolean showWidgets, PackageUserKey key);
 }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
index d52134c..0e5a7d7 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
@@ -98,7 +98,7 @@
     private Predicate<WidgetsListBaseEntry> mHeaderAndSelectedContentFilter = entry ->
             entry instanceof WidgetsListHeaderEntry
                     || entry instanceof WidgetsListSearchHeaderEntry
-                    || new PackageUserKey(entry.mPkgItem.packageName, entry.mPkgItem.user)
+                    || PackageUserKey.fromPackageItemInfo(entry.mPkgItem)
                             .equals(mWidgetsContentVisiblePackageUserKey);
     @Nullable private Predicate<WidgetsListBaseEntry> mFilter = null;
     @Nullable private RecyclerView mRecyclerView;
@@ -252,10 +252,11 @@
         return entry instanceof WidgetsListBaseEntry.Header && matchesKey(entry, key);
     }
 
-    private static boolean matchesKey(
-            @NonNull WidgetsListBaseEntry entry, @Nullable PackageUserKey key) {
+    private static boolean matchesKey(@NonNull WidgetsListBaseEntry entry,
+            @Nullable PackageUserKey key) {
         if (key == null) return false;
         return entry.mPkgItem.packageName.equals(key.mPackageName)
+                && entry.mPkgItem.widgetCategory == key.mWidgetCategory
                 && entry.mPkgItem.user.equals(key.mUser);
     }
 
@@ -434,11 +435,10 @@
                         .filter(entry -> entry instanceof WidgetsListHeaderEntry)
                         .map(entry -> entry.mPkgItem)
                         .collect(Collectors.toMap(
-                                entry -> new PackageUserKey(entry.packageName, entry.user),
+                                entry -> PackageUserKey.fromPackageItemInfo(entry),
                                 entry -> entry));
         for (WidgetsListBaseEntry visibleEntry: mVisibleEntries) {
-            PackageUserKey key = new PackageUserKey(visibleEntry.mPkgItem.packageName,
-                    visibleEntry.mPkgItem.user);
+            PackageUserKey key = PackageUserKey.fromPackageItemInfo(visibleEntry.mPkgItem);
             PackageItemInfo packageItemInfo = packagesInfo.get(key);
             if (packageItemInfo != null
                     && !visibleEntry.mPkgItem.title.equals(packageItemInfo.title)) {
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListHeader.java b/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
index ebd2d10..932e06d 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.widget.picker;
 
+import static com.android.launcher3.widget.WidgetSections.NO_CATEGORY;
+
 import android.content.Context;
 import android.content.res.Resources;
 import android.content.res.TypedArray;
@@ -41,6 +43,8 @@
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.util.PluralMessageFormat;
 import com.android.launcher3.views.ActivityContext;
+import com.android.launcher3.widget.WidgetSections;
+import com.android.launcher3.widget.WidgetSections.WidgetSection;
 import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
 import com.android.launcher3.widget.model.WidgetsListSearchHeaderEntry;
 
@@ -173,12 +177,12 @@
 
     private void setIcon(PackageItemInfo info) {
         Drawable icon;
-        switch (info.category) {
-            case PackageItemInfo.CONVERSATIONS:
-                icon = getContext().getDrawable(R.drawable.ic_conversations_widget_category);
-                break;
-            default:
-                icon = info.newIcon(getContext());
+        if (info.widgetCategory == NO_CATEGORY) {
+            icon = info.newIcon(getContext());
+        } else {
+            WidgetSection widgetSection = WidgetSections.getWidgetSections(getContext())
+                    .get(info.widgetCategory);
+            icon = getContext().getDrawable(widgetSection.mSectionDrawable);
         }
         applyDrawables(icon);
         mIconDrawable = icon;
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java
index fadb637..c6a7285 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinder.java
@@ -62,9 +62,7 @@
                         (position & POSITION_LAST) != 0,
                         /* isExpanded= */ data.isWidgetListShown()));
         widgetsListHeader.setOnExpandChangeListener(isExpanded ->
-                mOnHeaderClickListener.onHeaderClicked(
-                        isExpanded,
-                        new PackageUserKey(data.mPkgItem.packageName, data.mPkgItem.user)
-                ));
+                mOnHeaderClickListener.onHeaderClicked(isExpanded,
+                        PackageUserKey.fromPackageItemInfo(data.mPkgItem)));
     }
 }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinder.java
index bff43c1..2b27fc2 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinder.java
@@ -64,6 +64,6 @@
                         /* isExpanded= */ data.isWidgetListShown()));
         widgetsListHeader.setOnExpandChangeListener(isExpanded ->
                 mOnHeaderClickListener.onHeaderClicked(isExpanded,
-                        new PackageUserKey(data.mPkgItem.packageName, data.mPkgItem.user)));
+                        PackageUserKey.fromPackageItemInfo(data.mPkgItem)));
     }
 }
diff --git a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
index 12e9e1e..702f343 100644
--- a/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
+++ b/src_shortcuts_overrides/com/android/launcher3/model/WidgetsModel.java
@@ -4,7 +4,10 @@
 import static android.appwidget.AppWidgetProviderInfo.WIDGET_FEATURE_HIDE_FROM_PICKER;
 
 import static com.android.launcher3.pm.ShortcutConfigActivityInfo.queryList;
+import static com.android.launcher3.widget.WidgetSections.NO_CATEGORY;
 
+import static java.util.stream.Collectors.groupingBy;
+import static java.util.stream.Collectors.mapping;
 import static java.util.stream.Collectors.toList;
 
 import android.appwidget.AppWidgetProviderInfo;
@@ -13,6 +16,7 @@
 import android.content.pm.PackageManager;
 import android.os.UserHandle;
 import android.util.Log;
+import android.util.Pair;
 
 import androidx.annotation.Nullable;
 import androidx.collection.ArrayMap;
@@ -27,10 +31,12 @@
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.pm.ShortcutConfigActivityInfo;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.util.Preconditions;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.WidgetManagerHelper;
+import com.android.launcher3.widget.WidgetSections;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
 import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
@@ -40,12 +46,12 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Set;
 import java.util.function.Predicate;
-import java.util.stream.Collectors;
 
 /**
  * Widgets data model that is used by the adapters of the widget views and controllers.
@@ -61,9 +67,6 @@
     private static final String TAG = "WidgetsModel";
     private static final boolean DEBUG = false;
 
-    private static final ComponentName CONVERSATION_WIDGET = ComponentName.createRelative(
-            "com.android.systemui", ".people.widget.PeopleSpaceWidgetProvider");
-
     /* Map of widgets and shortcuts that are tracked per package. */
     private final Map<PackageItemInfo, List<WidgetItem>> mWidgetsList = new HashMap<>();
 
@@ -168,16 +171,15 @@
             mWidgetsList.clear();
         } else {
             // Otherwise, only clear the widgets and shortcuts for the changed package.
-            mWidgetsList.remove(
-                    packageItemInfoCache.getOrCreate(new WidgetPackageOrCategoryKey(packageUser)));
+            mWidgetsList.remove(packageItemInfoCache.getOrCreate(packageUser));
         }
 
         // add and update.
         mWidgetsList.putAll(rawWidgetsShortcuts.stream()
                 .filter(new WidgetValidityCheck(app))
-                .collect(Collectors.groupingBy(item ->
-                        packageItemInfoCache.getOrCreate(getWidgetPackageOrCategoryKey(item))
-                )));
+                .flatMap(widgetItem -> getPackageUserKeys(app.getContext(), widgetItem).stream()
+                        .map(key -> new Pair<>(packageItemInfoCache.getOrCreate(key), widgetItem)))
+                .collect(groupingBy(pair -> pair.first, mapping(pair -> pair.second, toList()))));
 
         // Update each package entry
         IconCache iconCache = app.getIconCache();
@@ -209,9 +211,9 @@
     }
 
     public WidgetItem getWidgetProviderInfoByProviderName(
-            ComponentName providerName) {
+            ComponentName providerName, UserHandle user) {
         List<WidgetItem> widgetsList = mWidgetsList.get(
-                new PackageItemInfo(providerName.getPackageName()));
+                new PackageItemInfo(providerName.getPackageName(), user));
         if (widgetsList == null) {
             return null;
         }
@@ -225,18 +227,40 @@
     }
 
     /** Returns {@link PackageItemInfo} of a pending widget. */
-    public static PackageItemInfo newPendingItemInfo(ComponentName provider) {
-        if (CONVERSATION_WIDGET.equals(provider)) {
-            return new PackageItemInfo(provider.getPackageName(), PackageItemInfo.CONVERSATIONS);
+    public static PackageItemInfo newPendingItemInfo(Context context, ComponentName provider,
+            UserHandle user) {
+        Map<ComponentName, IntSet> widgetsToCategories =
+                WidgetSections.getWidgetsToCategory(context);
+        if (widgetsToCategories.containsKey(provider)) {
+            Iterator<Integer> categoriesIterator = widgetsToCategories.get(provider).iterator();
+            int firstCategory = NO_CATEGORY;
+            while (categoriesIterator.hasNext() && firstCategory == NO_CATEGORY) {
+                firstCategory = categoriesIterator.next();
+            }
+            return new PackageItemInfo(provider.getPackageName(), firstCategory, user);
         }
-        return new PackageItemInfo(provider.getPackageName());
+        return new PackageItemInfo(provider.getPackageName(), user);
     }
 
-    private WidgetPackageOrCategoryKey getWidgetPackageOrCategoryKey(WidgetItem item) {
-        if (CONVERSATION_WIDGET.equals(item.componentName)) {
-            return new WidgetPackageOrCategoryKey(PackageItemInfo.CONVERSATIONS, item.user);
+    private List<PackageUserKey> getPackageUserKeys(Context context, WidgetItem item) {
+        Map<ComponentName, IntSet> widgetsToCategories =
+                WidgetSections.getWidgetsToCategory(context);
+        IntSet categories = widgetsToCategories.get(item.componentName);
+        if (categories == null || categories.isEmpty()) {
+            return Arrays.asList(
+                    new PackageUserKey(item.componentName.getPackageName(), item.user));
         }
-        return new WidgetPackageOrCategoryKey(item.componentName.getPackageName(), item.user);
+        List<PackageUserKey> packageUserKeys = new ArrayList<>();
+        categories.forEach(category -> {
+            if (category == NO_CATEGORY) {
+                packageUserKeys.add(
+                        new PackageUserKey(item.componentName.getPackageName(),
+                                item.user));
+            } else {
+                packageUserKeys.add(new PackageUserKey(category, item.user));
+            }
+        });
+        return packageUserKeys;
     }
 
     private static class WidgetValidityCheck implements Predicate<WidgetItem> {
@@ -279,53 +303,13 @@
         }
     }
 
-    /** A hash key for grouping widgets by package name or category. */
-    private static class WidgetPackageOrCategoryKey {
-        /**
-         * The package name of the widget provider.
-         *
-         * <p>This shouldn't be empty if {@link #mCategory} has a value,
-         * {@link PackageItemInfo#NO_CATEGORY}.
-         */
-        public final String mPackage;
-        /** A widget category. */
-        @PackageItemInfo.Category public final int mCategory;
-        public final UserHandle mUser;
-        private final int mHashCode;
-
-        WidgetPackageOrCategoryKey(PackageUserKey key) {
-            this(key.mPackageName, key.mUser);
-        }
-
-        WidgetPackageOrCategoryKey(String packageName, UserHandle user) {
-            this(packageName,  PackageItemInfo.NO_CATEGORY, user);
-        }
-
-        WidgetPackageOrCategoryKey(@PackageItemInfo.Category int category, UserHandle user) {
-            this("", category, user);
-        }
-
-        private WidgetPackageOrCategoryKey(String packageName,
-                @PackageItemInfo.Category int category, UserHandle user) {
-            mPackage = packageName;
-            mCategory = category;
-            mUser = user;
-            mHashCode = Arrays.hashCode(new Object[]{mPackage, mCategory, mUser});
-        }
-
-        @Override
-        public int hashCode() {
-            return mHashCode;
-        }
-    }
-
     private static final class PackageItemInfoCache {
-        private final Map<WidgetPackageOrCategoryKey, PackageItemInfo> mMap = new ArrayMap<>();
+        private final Map<PackageUserKey, PackageItemInfo> mMap = new ArrayMap<>();
 
-        PackageItemInfo getOrCreate(WidgetPackageOrCategoryKey key) {
+        PackageItemInfo getOrCreate(PackageUserKey key) {
             PackageItemInfo pInfo = mMap.get(key);
             if (pInfo == null) {
-                pInfo = new PackageItemInfo(key.mPackage, key.mCategory);
+                pInfo = new PackageItemInfo(key.mPackageName, key.mWidgetCategory, key.mUser);
                 pInfo.user = key.mUser;
                 mMap.put(key,  pInfo);
             }
diff --git a/tests/src/com/android/launcher3/util/PackageUserKeyTest.java b/tests/src/com/android/launcher3/util/PackageUserKeyTest.java
new file mode 100644
index 0000000..99490eb
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/PackageUserKeyTest.java
@@ -0,0 +1,162 @@
+/*
+ * 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.util;
+
+import static com.android.launcher3.widget.WidgetSections.NO_CATEGORY;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.os.UserHandle;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.launcher3.model.data.PackageItemInfo;
+
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.junit.runner.RunWith;
+
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public final class PackageUserKeyTest {
+    @Rule
+    public ExpectedException exception = ExpectedException.none();
+
+    private static final String TEST_PACKAGE = "com.android.test.package";
+    private static final int CONVERSATIONS = 0;
+    private static final int WEATHER = 1;
+
+    @Test
+    public void fromPackageItemInfo_shouldCreateExpectedObject() {
+        PackageUserKey packageUserKey = PackageUserKey.fromPackageItemInfo(
+                new PackageItemInfo(TEST_PACKAGE, UserHandle.CURRENT));
+
+        assertThat(packageUserKey.mPackageName).isEqualTo(TEST_PACKAGE);
+        assertThat(packageUserKey.mWidgetCategory).isEqualTo(NO_CATEGORY);
+        assertThat(packageUserKey.mUser).isEqualTo(UserHandle.CURRENT);
+    }
+
+    @Test
+    public void constructor_packageNameAndUserHandle_shouldCreateExpectedObject() {
+        PackageUserKey packageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.CURRENT);
+
+        assertThat(packageUserKey.mPackageName).isEqualTo(TEST_PACKAGE);
+        assertThat(packageUserKey.mWidgetCategory).isEqualTo(NO_CATEGORY);
+        assertThat(packageUserKey.mUser).isEqualTo(UserHandle.CURRENT);
+    }
+
+    @Test
+    public void constructor_widgetCategoryAndUserHandle_shouldCreateExpectedObject() {
+        PackageUserKey packageUserKey = new PackageUserKey(CONVERSATIONS, UserHandle.CURRENT);
+
+        assertThat(packageUserKey.mPackageName).isEqualTo("");
+        assertThat(packageUserKey.mWidgetCategory).isEqualTo(CONVERSATIONS);
+        assertThat(packageUserKey.mUser).isEqualTo(UserHandle.CURRENT);
+    }
+
+    @Test
+    public void equals_sameObject_shouldReturnTrue() {
+        PackageUserKey packageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.CURRENT);
+        PackageUserKey otherPackageUserKey = packageUserKey;
+
+        assertThat(packageUserKey).isEqualTo(otherPackageUserKey);
+    }
+
+    @Test
+    public void equals_differentObjectSameContent_shouldReturnTrue() {
+        PackageUserKey packageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.CURRENT);
+        PackageUserKey otherPackageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.CURRENT);
+
+        assertThat(packageUserKey).isEqualTo(otherPackageUserKey);
+    }
+
+    @Test
+    public void equals_compareAgainstNull_shouldReturnFalse() {
+        PackageUserKey packageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.CURRENT);
+
+        assertThat(packageUserKey).isNotEqualTo(null);
+    }
+
+    @Test
+    public void equals_differentPackage_shouldReturnFalse() {
+        PackageUserKey packageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.CURRENT);
+        PackageUserKey otherPackageUserKey = new PackageUserKey(TEST_PACKAGE + "1",
+                UserHandle.CURRENT);
+
+        assertThat(packageUserKey).isNotEqualTo(otherPackageUserKey);
+    }
+
+
+    @Test
+    public void equals_differentCategory_shouldReturnFalse() {
+        PackageUserKey packageUserKey = new PackageUserKey(WEATHER, UserHandle.CURRENT);
+        PackageUserKey otherPackageUserKey = new PackageUserKey(CONVERSATIONS, UserHandle.CURRENT);
+
+        assertThat(packageUserKey).isNotEqualTo(otherPackageUserKey);
+    }
+
+    @Test
+    public void equals_differentUser_shouldReturnFalse() {
+        PackageUserKey packageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.of(1));
+        PackageUserKey otherPackageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.of(2));
+
+        assertThat(packageUserKey).isNotEqualTo(otherPackageUserKey);
+    }
+
+    @Test
+    public void hashCode_sameObject_shouldBeTheSame() {
+        PackageUserKey packageUserKey = new PackageUserKey(WEATHER, UserHandle.CURRENT);
+        PackageUserKey otherPackageUserKey = packageUserKey;
+
+        assertThat(packageUserKey.hashCode()).isEqualTo(otherPackageUserKey.hashCode());
+    }
+
+    @Test
+    public void hashCode_differentObjectSameContent_shouldBeTheSame() {
+        PackageUserKey packageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.CURRENT);
+        PackageUserKey otherPackageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.CURRENT);
+
+        assertThat(packageUserKey.hashCode()).isEqualTo(otherPackageUserKey.hashCode());
+    }
+
+    @Test
+    public void hashCode_differentPackage_shouldBeDifferent() {
+        PackageUserKey packageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.CURRENT);
+        PackageUserKey otherPackageUserKey = new PackageUserKey(TEST_PACKAGE + "1",
+                UserHandle.CURRENT);
+
+        assertThat(packageUserKey.hashCode()).isNotEqualTo(otherPackageUserKey.hashCode());
+    }
+
+
+    @Test
+    public void hashCode_differentCategory_shouldBeDifferent() {
+        PackageUserKey packageUserKey = new PackageUserKey(WEATHER, UserHandle.CURRENT);
+        PackageUserKey otherPackageUserKey = new PackageUserKey(CONVERSATIONS, UserHandle.CURRENT);
+
+        assertThat(packageUserKey.hashCode()).isNotEqualTo(otherPackageUserKey.hashCode());
+    }
+
+    @Test
+    public void hashCode_differentUser_shouldBeDifferent() {
+        PackageUserKey packageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.of(1));
+        PackageUserKey otherPackageUserKey = new PackageUserKey(TEST_PACKAGE, UserHandle.of(2));
+
+        assertThat(packageUserKey.hashCode()).isNotEqualTo(otherPackageUserKey.hashCode());
+    }
+}
diff --git a/tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java b/tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
index 6232938..b480a4c 100644
--- a/tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/WidgetsDiffReporterTest.java
@@ -288,9 +288,8 @@
 
     private PackageItemInfo createPackageItemInfo(String packageName, String appName,
             UserHandle userHandle) {
-        PackageItemInfo pInfo = new PackageItemInfo(packageName);
+        PackageItemInfo pInfo = new PackageItemInfo(packageName, userHandle);
         pInfo.title = appName;
-        pInfo.user = userHandle;
         pInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
         return pInfo;
     }
diff --git a/tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java b/tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
index 44d6964..4e0bdda 100644
--- a/tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/WidgetsListAdapterTest.java
@@ -215,14 +215,23 @@
     @Test
     public void setWidgetsOnSearch_expandedApp_shouldResetExpandedApp() {
         // GIVEN a list of widgets entries:
-        // [com.google.test0, com.google.test0 content,
-        //  com.google.test1, com.google.test1 content,
-        //  com.google.test2, com.google.test2 content]
-        // The visible widgets entries: [com.google.test0, com.google.test1, com.google.test2].
-        ArrayList<WidgetsListBaseEntry> allEntries = generateSampleMap(2);
+        // [Empty item
+        //  com.google.test0,
+        //  com.google.test0 content,
+        //  com.google.test1,
+        //  com.google.test1 content,
+        //  com.google.test2,
+        //  com.google.test2 content]
+        // The visible widgets entries:
+        // [Empty item,
+        // com.google.test0,
+        // com.google.test1,
+        // com.google.test2].
+        ArrayList<WidgetsListBaseEntry> allEntries = generateSampleMap(3);
         mAdapter.setWidgetsOnSearch(allEntries);
         // GIVEN com.google.test.1 header is expanded. The visible entries list becomes:
-        // [com.google.test0, com.google.test1, com.google.test1 content, com.google.test2]
+        // [Empty item, com.google.test0, com.google.test1, com.google.test1 content,
+        // com.google.test2]
         mAdapter.onHeaderClicked(/* showWidgets= */ true,
                 new PackageUserKey(TEST_PACKAGE_PLACEHOLDER + 1, mUserHandle));
         Mockito.reset(mListener);
@@ -231,7 +240,7 @@
         mAdapter.setWidgetsOnSearch(allEntries);
 
         // THEN expanded app is reset and the visible entries list becomes:
-        // [com.google.test0, com.google.test1, com.google.test2]
+        // [Empty item, com.google.test0, com.google.test1, com.google.test2]
         verify(mListener).onItemRangeChanged(eq(2), eq(1), isNull());
         verify(mListener).onItemRangeRemoved(/* positionStart= */ 3, /* itemCount= */ 1);
     }
@@ -257,9 +266,8 @@
 
             List<WidgetItem> widgetItems = generateWidgetItems(packageName, /* numOfWidgets= */ 1);
 
-            PackageItemInfo pInfo = new PackageItemInfo(packageName);
+            PackageItemInfo pInfo = new PackageItemInfo(packageName, widgetItems.get(0).user);
             pInfo.title = pInfo.packageName;
-            pInfo.user = widgetItems.get(0).user;
             pInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
 
             result.add(new WidgetsListHeaderEntry(pInfo, /* titleSectionName= */ "", widgetItems));
diff --git a/tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java b/tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
index 969c12a..211318c 100644
--- a/tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
@@ -30,6 +30,7 @@
 import android.content.ComponentName;
 import android.content.Context;
 import android.graphics.Bitmap;
+import android.os.UserHandle;
 import android.view.LayoutInflater;
 import android.widget.FrameLayout;
 import android.widget.TextView;
@@ -124,12 +125,12 @@
         widgetsListHeader.callOnClick();
 
         verify(mOnHeaderClickListener).onHeaderClicked(eq(true),
-                eq(new PackageUserKey(entry.mPkgItem.packageName, entry.mPkgItem.user)));
+                eq(PackageUserKey.fromPackageItemInfo(entry.mPkgItem)));
     }
 
     private WidgetsListHeaderEntry generateSampleAppHeader(String appName, String packageName,
             int numOfWidgets) {
-        PackageItemInfo appInfo = new PackageItemInfo(packageName);
+        PackageItemInfo appInfo = new PackageItemInfo(packageName, UserHandle.CURRENT);
         appInfo.title = appName;
         appInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
 
diff --git a/tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java b/tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java
index 453f4fb..66c2f36 100644
--- a/tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/WidgetsListSearchHeaderViewHolderBinderTest.java
@@ -30,6 +30,7 @@
 import android.content.ComponentName;
 import android.content.Context;
 import android.graphics.Bitmap;
+import android.os.UserHandle;
 import android.view.LayoutInflater;
 import android.widget.FrameLayout;
 import android.widget.TextView;
@@ -124,12 +125,12 @@
         widgetsListHeader.callOnClick();
 
         verify(mOnHeaderClickListener).onHeaderClicked(eq(true),
-                eq(new PackageUserKey(entry.mPkgItem.packageName, entry.mPkgItem.user)));
+                eq(PackageUserKey.fromPackageItemInfo(entry.mPkgItem)));
     }
 
     private WidgetsListSearchHeaderEntry generateSampleSearchHeader(String appName,
             String packageName, int numOfWidgets) {
-        PackageItemInfo appInfo = new PackageItemInfo(packageName);
+        PackageItemInfo appInfo = new PackageItemInfo(packageName, UserHandle.CURRENT);
         appInfo.title = appName;
         appInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
 
diff --git a/tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java b/tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
index 5816b77..7ec4d20 100644
--- a/tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
@@ -28,6 +28,7 @@
 import android.content.ComponentName;
 import android.content.Context;
 import android.graphics.Bitmap;
+import android.os.UserHandle;
 import android.view.LayoutInflater;
 import android.view.View;
 import android.view.View.OnClickListener;
@@ -125,7 +126,7 @@
 
     private WidgetsListContentEntry generateSampleAppWithWidgets(String appName, String packageName,
             int numOfWidgets) {
-        PackageItemInfo appInfo = new PackageItemInfo(packageName);
+        PackageItemInfo appInfo = new PackageItemInfo(packageName, UserHandle.CURRENT);
         appInfo.title = appName;
         appInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
 
diff --git a/tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java b/tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
index 4b61b2c..d8f1f14 100644
--- a/tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
@@ -26,6 +26,7 @@
 
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ComponentName;
+import android.os.UserHandle;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
@@ -53,8 +54,10 @@
 public final class WidgetsListContentEntryTest {
     private static final String PACKAGE_NAME = "com.android.test";
     private static final String PACKAGE_NAME_2 = "com.android.test2";
-    private final PackageItemInfo mPackageItemInfo1 = new PackageItemInfo(PACKAGE_NAME);
-    private final PackageItemInfo mPackageItemInfo2 = new PackageItemInfo(PACKAGE_NAME_2);
+    private final PackageItemInfo mPackageItemInfo1 = new PackageItemInfo(PACKAGE_NAME,
+            UserHandle.CURRENT);
+    private final PackageItemInfo mPackageItemInfo2 = new PackageItemInfo(PACKAGE_NAME_2,
+            UserHandle.CURRENT);
     private final ComponentName mWidget1 = ComponentName.createRelative(PACKAGE_NAME, ".mWidget1");
     private final ComponentName mWidget2 = ComponentName.createRelative(PACKAGE_NAME, ".mWidget2");
     private final ComponentName mWidget3 = ComponentName.createRelative(PACKAGE_NAME, ".mWidget3");
diff --git a/tests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java b/tests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
index c862d6b..d812ab0 100644
--- a/tests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
+++ b/tests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
@@ -189,9 +189,8 @@
 
     private PackageItemInfo createPackageItemInfo(String packageName, String appName,
             UserHandle userHandle) {
-        PackageItemInfo pInfo = new PackageItemInfo(packageName);
+        PackageItemInfo pInfo = new PackageItemInfo(packageName, userHandle);
         pInfo.title = appName;
-        pInfo.user = userHandle;
         pInfo.bitmap = BitmapInfo.of(Bitmap.createBitmap(10, 10, Bitmap.Config.ALPHA_8), 0);
         return pInfo;
     }