diff --git a/src/com/android/launcher3/popup/PopupDataProvider.java b/src/com/android/launcher3/popup/PopupDataProvider.java
index c773079..bb9c520 100644
--- a/src/com/android/launcher3/popup/PopupDataProvider.java
+++ b/src/com/android/launcher3/popup/PopupDataProvider.java
@@ -148,7 +148,7 @@
     }
 
     public String[] getNotificationKeysForItem(ItemInfo info) {
-        BadgeInfo badgeInfo = mPackageUserToBadgeInfos.get(PackageUserKey.fromItemInfo(info));
+        BadgeInfo badgeInfo = getBadgeInfoForItem(info);
         if (badgeInfo == null) { return new String[0]; }
         List<String> notificationKeys = badgeInfo.getNotificationKeys();
         return notificationKeys.toArray(new String[notificationKeys.size()]);