Add getShareTargets() to ShortcutServiceInternal and get caller's user ID from AppPredictionSessionId

Previously, the caller's user ID was got from Binder. But
AppPredictionManager clears caller's identity. The only way to get the
caller's user ID is to get it from session ID.

Bug: 146522621
Test: atest com.android.server.people.data.DataManagerTest
Test: atest com.android.server.people.data.ShareTargetPredictorTest
Change-Id: Ia0dc35879084110849e06e919ff3b3c7241c46fd

9 files changed, 81 insertions, 109 deletions

diff --git a/core/java/android/content/pm/ShortcutServiceInternal.java b/core/java/android/content/pm/ShortcutServiceInternal.java
index a69905eb3de4..3a934211c7e5 100644
--- a/core/java/android/content/pm/ShortcutServiceInternal.java
+++ b/core/java/android/content/pm/ShortcutServiceInternal.java
@@ -22,6 +22,7 @@ import android.annotation.UserIdInt;
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ComponentName;
 import android.content.Intent;
+import android.content.IntentFilter;
 import android.content.IntentSender;
 import android.content.LocusId;
 import android.content.pm.LauncherApps.ShortcutQuery;
@@ -92,4 +93,11 @@ public abstract class ShortcutServiceInternal {
     public abstract void uncacheShortcuts(int launcherUserId,
             @NonNull String callingPackage, @NonNull String packageName,
             @NonNull List<String> shortcutIds, int userId);
+
+    /**
+     * Retrieves all of the direct share targets that match the given IntentFilter for the specified
+     * user.
+     */
+    public abstract List<ShortcutManager.ShareShortcutInfo> getShareTargets(
+            @NonNull String callingPackage, @NonNull IntentFilter intentFilter, int userId);
 }
diff --git a/services/core/java/com/android/server/pm/ShortcutService.java b/services/core/java/com/android/server/pm/ShortcutService.java
index 377fd16d4e19..12f7d5c27459 100644
--- a/services/core/java/com/android/server/pm/ShortcutService.java
+++ b/services/core/java/com/android/server/pm/ShortcutService.java
@@ -2773,6 +2773,13 @@ public class ShortcutService extends IShortcutService.Stub {
                     userId, /* doCache= */ false);
         }
 
+        @Override
+        public List<ShortcutManager.ShareShortcutInfo> getShareTargets(
+                @NonNull String callingPackage, @NonNull IntentFilter intentFilter, int userId) {
+            return ShortcutService.this.getShareTargets(
+                    callingPackage, intentFilter, userId).getList();
+        }
+
         private void updateCachedShortcutsInternal(int launcherUserId,
                 @NonNull String callingPackage, @NonNull String packageName,
                 @NonNull List<String> shortcutIds, int userId, boolean doCache) {
diff --git a/services/people/java/com/android/server/people/PeopleService.java b/services/people/java/com/android/server/people/PeopleService.java
index 663bf4f30708..2499614a3738 100644
--- a/services/people/java/com/android/server/people/PeopleService.java
+++ b/services/people/java/com/android/server/people/PeopleService.java
@@ -88,7 +88,7 @@ public class PeopleService extends SystemService {
         @Override
         public void onCreatePredictionSession(AppPredictionContext context,
                 AppPredictionSessionId sessionId) {
-            mSessions.put(sessionId, new SessionInfo(context, mDataManager));
+            mSessions.put(sessionId, new SessionInfo(context, mDataManager, sessionId.getUserId()));
         }
 
         @Override
diff --git a/services/people/java/com/android/server/people/SessionInfo.java b/services/people/java/com/android/server/people/SessionInfo.java
index eaa0781f12ef..28612f1dd49b 100644
--- a/services/people/java/com/android/server/people/SessionInfo.java
+++ b/services/people/java/com/android/server/people/SessionInfo.java
@@ -16,6 +16,7 @@
 
 package com.android.server.people;
 
+import android.annotation.UserIdInt;
 import android.app.prediction.AppPredictionContext;
 import android.app.prediction.AppTarget;
 import android.app.prediction.IPredictionCallback;
@@ -38,9 +39,10 @@ class SessionInfo {
     private final RemoteCallbackList<IPredictionCallback> mCallbacks =
             new RemoteCallbackList<>();
 
-    SessionInfo(AppPredictionContext predictionContext, DataManager dataManager) {
+    SessionInfo(AppPredictionContext predictionContext, DataManager dataManager,
+            @UserIdInt int callingUserId) {
         mAppTargetPredictor = AppTargetPredictor.create(predictionContext,
-                this::updatePredictions, dataManager);
+                this::updatePredictions, dataManager, callingUserId);
     }
 
     void addCallback(IPredictionCallback callback) {
diff --git a/services/people/java/com/android/server/people/data/DataManager.java b/services/people/java/com/android/server/people/data/DataManager.java
index dd9cbd00f6a2..6b97c98b0029 100644
--- a/services/people/java/com/android/server/people/data/DataManager.java
+++ b/services/people/java/com/android/server/people/data/DataManager.java
@@ -34,13 +34,11 @@ import android.content.IntentFilter;
 import android.content.pm.LauncherApps.ShortcutQuery;
 import android.content.pm.PackageManagerInternal;
 import android.content.pm.ShortcutInfo;
-import android.content.pm.ShortcutManager;
 import android.content.pm.ShortcutManager.ShareShortcutInfo;
 import android.content.pm.ShortcutServiceInternal;
 import android.content.pm.UserInfo;
 import android.database.ContentObserver;
 import android.net.Uri;
-import android.os.Binder;
 import android.os.CancellationSignal;
 import android.os.Handler;
 import android.os.Process;
@@ -83,7 +81,6 @@ import java.util.function.Function;
  */
 public class DataManager {
 
-    private static final String PLATFORM_PACKAGE_NAME = "android";
     private static final int MY_UID = Process.myUid();
     private static final int MY_PID = Process.myPid();
     private static final long QUERY_EVENTS_MAX_AGE_MS = DateUtils.DAY_IN_MILLIS;
@@ -106,7 +103,6 @@ public class DataManager {
 
     private ShortcutServiceInternal mShortcutServiceInternal;
     private PackageManagerInternal mPackageManagerInternal;
-    private ShortcutManager mShortcutManager;
     private UserManager mUserManager;
 
     public DataManager(Context context) {
@@ -125,7 +121,6 @@ public class DataManager {
     public void initialize() {
         mShortcutServiceInternal = LocalServices.getService(ShortcutServiceInternal.class);
         mPackageManagerInternal = LocalServices.getService(PackageManagerInternal.class);
-        mShortcutManager = mContext.getSystemService(ShortcutManager.class);
         mUserManager = mContext.getSystemService(UserManager.class);
 
         mShortcutServiceInternal.addListener(new ShortcutServiceListener());
@@ -171,8 +166,7 @@ public class DataManager {
         mNotificationListeners.put(userId, notificationListener);
         try {
             notificationListener.registerAsSystemService(mContext,
-                    new ComponentName(PLATFORM_PACKAGE_NAME, getClass().getCanonicalName()),
-                    userId);
+                    new ComponentName(mContext, getClass()), userId);
         } catch (RemoteException e) {
             // Should never occur for local calls.
         }
@@ -242,8 +236,8 @@ public class DataManager {
      * Iterates through all the {@link PackageData}s owned by the unlocked users who are in the
      * same profile group as the calling user.
      */
-    public void forAllPackages(Consumer<PackageData> consumer) {
-        List<UserInfo> users = mUserManager.getEnabledProfiles(mInjector.getCallingUserId());
+    void forPackagesInProfile(@UserIdInt int callingUserId, Consumer<PackageData> consumer) {
+        List<UserInfo> users = mUserManager.getEnabledProfiles(callingUserId);
         for (UserInfo userInfo : users) {
             UserData userData = getUnlockedUserData(userInfo.id);
             if (userData != null) {
@@ -275,8 +269,10 @@ public class DataManager {
      * Gets the {@link ShareShortcutInfo}s from all packages owned by the calling user that match
      * the specified {@link IntentFilter}.
      */
-    public List<ShareShortcutInfo> getShareShortcuts(@NonNull IntentFilter intentFilter) {
-        return mShortcutManager.getShareTargets(intentFilter);
+    public List<ShareShortcutInfo> getShareShortcuts(@NonNull IntentFilter intentFilter,
+            @UserIdInt int callingUserId) {
+        return mShortcutServiceInternal.getShareTargets(
+                mContext.getPackageName(), intentFilter, callingUserId);
     }
 
     /** Reports the {@link AppTargetEvent} from App Prediction Manager. */
@@ -361,7 +357,7 @@ public class DataManager {
         @ShortcutQuery.QueryFlags int queryFlags = ShortcutQuery.FLAG_MATCH_DYNAMIC
                 | ShortcutQuery.FLAG_MATCH_PINNED | ShortcutQuery.FLAG_MATCH_PINNED_BY_ANY_LAUNCHER;
         return mShortcutServiceInternal.getShortcuts(
-                mInjector.getCallingUserId(), /*callingPackage=*/ PLATFORM_PACKAGE_NAME,
+                UserHandle.USER_SYSTEM, mContext.getPackageName(),
                 /*changedSince=*/ 0, packageName, shortcutIds, /*locusIds=*/ null,
                 /*componentName=*/ null, queryFlags, userId, MY_PID, MY_UID);
     }
@@ -775,7 +771,7 @@ public class DataManager {
 
         @Override
         public void onReceive(Context context, Intent intent) {
-            forAllPackages(PackageData::saveToDisk);
+            forAllUnlockedUsers(userData -> userData.forAllPackages(PackageData::saveToDisk));
         }
     }
 
@@ -809,9 +805,5 @@ public class DataManager {
                 Function<String, PackageData> packageDataGetter) {
             return new UsageStatsQueryHelper(userId, packageDataGetter);
         }
-
-        int getCallingUserId() {
-            return Binder.getCallingUserHandle().getIdentifier();
-        }
     }
 }
diff --git a/services/people/java/com/android/server/people/prediction/AppTargetPredictor.java b/services/people/java/com/android/server/people/prediction/AppTargetPredictor.java
index 44f3e35833d9..19cf8af5d66b 100644
--- a/services/people/java/com/android/server/people/prediction/AppTargetPredictor.java
+++ b/services/people/java/com/android/server/people/prediction/AppTargetPredictor.java
@@ -18,6 +18,7 @@ package com.android.server.people.prediction;
 
 import android.annotation.MainThread;
 import android.annotation.NonNull;
+import android.annotation.UserIdInt;
 import android.annotation.WorkerThread;
 import android.app.prediction.AppPredictionContext;
 import android.app.prediction.AppTarget;
@@ -42,25 +43,28 @@ public class AppTargetPredictor {
     /** Creates a {@link AppTargetPredictor} instance based on the prediction context. */
     public static AppTargetPredictor create(@NonNull AppPredictionContext predictionContext,
             @NonNull Consumer<List<AppTarget>> updatePredictionsMethod,
-            @NonNull DataManager dataManager) {
+            @NonNull DataManager dataManager, @UserIdInt int callingUserId) {
         if (UI_SURFACE_SHARE.equals(predictionContext.getUiSurface())) {
             return new ShareTargetPredictor(
-                    predictionContext, updatePredictionsMethod, dataManager);
+                    predictionContext, updatePredictionsMethod, dataManager, callingUserId);
         }
-        return new AppTargetPredictor(predictionContext, updatePredictionsMethod, dataManager);
+        return new AppTargetPredictor(
+                predictionContext, updatePredictionsMethod, dataManager, callingUserId);
     }
 
     private final AppPredictionContext mPredictionContext;
     private final Consumer<List<AppTarget>> mUpdatePredictionsMethod;
     private final DataManager mDataManager;
+    final int mCallingUserId;
     private final ExecutorService mCallbackExecutor;
 
     AppTargetPredictor(@NonNull AppPredictionContext predictionContext,
             @NonNull Consumer<List<AppTarget>> updatePredictionsMethod,
-            @NonNull DataManager dataManager) {
+            @NonNull DataManager dataManager, @UserIdInt int callingUserId) {
         mPredictionContext = predictionContext;
         mUpdatePredictionsMethod = updatePredictionsMethod;
         mDataManager = dataManager;
+        mCallingUserId = callingUserId;
         mCallbackExecutor = Executors.newSingleThreadExecutor();
     }
 
diff --git a/services/people/java/com/android/server/people/prediction/ShareTargetPredictor.java b/services/people/java/com/android/server/people/prediction/ShareTargetPredictor.java
index 280ced3a07c5..90d821641149 100644
--- a/services/people/java/com/android/server/people/prediction/ShareTargetPredictor.java
+++ b/services/people/java/com/android/server/people/prediction/ShareTargetPredictor.java
@@ -19,6 +19,7 @@ package com.android.server.people.prediction;
 import android.annotation.MainThread;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.annotation.UserIdInt;
 import android.annotation.WorkerThread;
 import android.app.prediction.AppPredictionContext;
 import android.app.prediction.AppTarget;
@@ -45,8 +46,8 @@ class ShareTargetPredictor extends AppTargetPredictor {
 
     ShareTargetPredictor(@NonNull AppPredictionContext predictionContext,
             @NonNull Consumer<List<AppTarget>> updatePredictionsMethod,
-            @NonNull DataManager dataManager) {
-        super(predictionContext, updatePredictionsMethod, dataManager);
+            @NonNull DataManager dataManager, @UserIdInt int callingUserId) {
+        super(predictionContext, updatePredictionsMethod, dataManager, callingUserId);
         mIntentFilter = predictionContext.getExtras().getParcelable(
                 ChooserActivity.APP_PREDICTION_INTENT_FILTER_KEY);
     }
@@ -84,7 +85,7 @@ class ShareTargetPredictor extends AppTargetPredictor {
     List<ShareTarget> getShareTargets() {
         List<ShareTarget> shareTargets = new ArrayList<>();
         List<ShareShortcutInfo> shareShortcuts =
-                getDataManager().getShareShortcuts(mIntentFilter);
+                getDataManager().getShareShortcuts(mIntentFilter, mCallingUserId);
 
         for (ShareShortcutInfo shareShortcut : shareShortcuts) {
             ShortcutInfo shortcutInfo = shareShortcut.getShortcutInfo();
diff --git a/services/tests/servicestests/src/com/android/server/people/data/DataManagerTest.java b/services/tests/servicestests/src/com/android/server/people/data/DataManagerTest.java
index 3ecd3193058a..5e104a5210d7 100644
--- a/services/tests/servicestests/src/com/android/server/people/data/DataManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/people/data/DataManagerTest.java
@@ -54,7 +54,6 @@ import android.content.Intent;
 import android.content.IntentFilter;
 import android.content.pm.PackageManagerInternal;
 import android.content.pm.ShortcutInfo;
-import android.content.pm.ShortcutManager;
 import android.content.pm.ShortcutServiceInternal;
 import android.content.pm.UserInfo;
 import android.content.pm.parsing.AndroidPackage;
@@ -87,6 +86,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
@@ -110,7 +110,6 @@ public final class DataManagerTest {
     @Mock private ShortcutServiceInternal mShortcutServiceInternal;
     @Mock private UsageStatsManagerInternal mUsageStatsManagerInternal;
     @Mock private PackageManagerInternal mPackageManagerInternal;
-    @Mock private ShortcutManager mShortcutManager;
     @Mock private UserManager mUserManager;
     @Mock private TelephonyManager mTelephonyManager;
     @Mock private TelecomManager mTelecomManager;
@@ -123,7 +122,6 @@ public final class DataManagerTest {
 
     private NotificationChannel mNotificationChannel;
     private DataManager mDataManager;
-    private int mCallingUserId;
     private CancellationSignal mCancellationSignal;
     private TestInjector mInjector;
 
@@ -145,14 +143,11 @@ public final class DataManagerTest {
         }).when(mPackageManagerInternal).forEachInstalledPackage(any(Consumer.class), anyInt());
 
         when(mContext.getMainLooper()).thenReturn(Looper.getMainLooper());
+        when(mContext.getPackageName()).thenReturn("android");
 
         Context originalContext = getInstrumentation().getTargetContext();
         when(mContext.getApplicationInfo()).thenReturn(originalContext.getApplicationInfo());
 
-        when(mContext.getSystemService(Context.SHORTCUT_SERVICE)).thenReturn(mShortcutManager);
-        when(mContext.getSystemServiceName(ShortcutManager.class)).thenReturn(
-                Context.SHORTCUT_SERVICE);
-
         when(mContext.getSystemService(Context.USER_SERVICE)).thenReturn(mUserManager);
         when(mContext.getSystemServiceName(UserManager.class)).thenReturn(
                 Context.USER_SERVICE);
@@ -191,8 +186,6 @@ public final class DataManagerTest {
                 NOTIFICATION_CHANNEL_ID, "test channel", NotificationManager.IMPORTANCE_DEFAULT);
         mNotificationChannel.setConversationId("test", TEST_SHORTCUT_ID);
 
-        mCallingUserId = USER_ID_PRIMARY;
-
         mCancellationSignal = new CancellationSignal();
 
         mInjector = new TestInjector();
@@ -222,9 +215,7 @@ public final class DataManagerTest {
         mDataManager.onShortcutAddedOrUpdated(
                 buildShortcutInfo("pkg_3", USER_ID_SECONDARY, "sc_3", buildPerson()));
 
-        List<ConversationInfo> conversations = new ArrayList<>();
-        mDataManager.forAllPackages(
-                packageData -> packageData.forAllConversations(conversations::add));
+        List<ConversationInfo> conversations = getConversationsInPrimary();
 
         // USER_ID_SECONDARY is not in the same profile group as USER_ID_PRIMARY.
         assertEquals(2, conversations.size());
@@ -250,18 +241,14 @@ public final class DataManagerTest {
         mDataManager.onShortcutAddedOrUpdated(
                 buildShortcutInfo("pkg_2", USER_ID_PRIMARY_MANAGED, "sc_2", buildPerson()));
 
-        List<ConversationInfo> conversations = new ArrayList<>();
-        mDataManager.forAllPackages(
-                packageData -> packageData.forAllConversations(conversations::add));
+        List<ConversationInfo> conversations = getConversationsInPrimary();
 
         // USER_ID_PRIMARY_MANAGED is not locked, so only USER_ID_PRIMARY's conversation is stored.
         assertEquals(1, conversations.size());
         assertEquals("sc_1", conversations.get(0).getShortcutId());
 
         mDataManager.onUserStopped(USER_ID_PRIMARY);
-        conversations.clear();
-        mDataManager.forAllPackages(
-                packageData -> packageData.forAllConversations(conversations::add));
+        conversations = getConversationsInPrimary();
         assertTrue(conversations.isEmpty());
     }
 
@@ -289,12 +276,8 @@ public final class DataManagerTest {
         IntentFilter intentFilter = new IntentFilter(Intent.ACTION_SEND, "image/jpg");
         mDataManager.reportAppTargetEvent(appTargetEvent, intentFilter);
 
-        List<Range<Long>> activeShareTimeSlots = new ArrayList<>();
-        mDataManager.forAllPackages(packageData ->
-                activeShareTimeSlots.addAll(
-                        packageData.getEventHistory(TEST_SHORTCUT_ID)
-                                .getEventIndex(Event.TYPE_SHARE_IMAGE)
-                                .getActiveTimeSlots()));
+        List<Range<Long>> activeShareTimeSlots = getActiveSlotsForTestShortcut(
+                Event.SHARE_EVENT_TYPES);
         assertEquals(1, activeShareTimeSlots.size());
     }
 
@@ -315,9 +298,7 @@ public final class DataManagerTest {
                 USER_ID_PRIMARY);
         contentObserver.onChange(false, ContactsContract.Contacts.CONTENT_URI, USER_ID_PRIMARY);
 
-        List<ConversationInfo> conversations = new ArrayList<>();
-        mDataManager.forAllPackages(
-                packageData -> packageData.forAllConversations(conversations::add));
+        List<ConversationInfo> conversations = getConversationsInPrimary();
         assertEquals(1, conversations.size());
 
         assertEquals(TEST_SHORTCUT_ID, conversations.get(0).getShortcutId());
@@ -338,12 +319,8 @@ public final class DataManagerTest {
 
         listenerService.onNotificationPosted(mStatusBarNotification);
 
-        List<Range<Long>> activeNotificationOpenTimeSlots = new ArrayList<>();
-        mDataManager.forAllPackages(packageData ->
-                activeNotificationOpenTimeSlots.addAll(
-                        packageData.getEventHistory(TEST_SHORTCUT_ID)
-                                .getEventIndex(Event.TYPE_NOTIFICATION_POSTED)
-                                .getActiveTimeSlots()));
+        List<Range<Long>> activeNotificationOpenTimeSlots = getActiveSlotsForTestShortcut(
+                Event.NOTIFICATION_EVENT_TYPES);
         assertEquals(1, activeNotificationOpenTimeSlots.size());
     }
 
@@ -361,12 +338,8 @@ public final class DataManagerTest {
         listenerService.onNotificationRemoved(mStatusBarNotification, null,
                 NotificationListenerService.REASON_CLICK);
 
-        List<Range<Long>> activeNotificationOpenTimeSlots = new ArrayList<>();
-        mDataManager.forAllPackages(packageData ->
-                activeNotificationOpenTimeSlots.addAll(
-                        packageData.getEventHistory(TEST_SHORTCUT_ID)
-                                .getEventIndex(Event.TYPE_NOTIFICATION_OPENED)
-                                .getActiveTimeSlots()));
+        List<Range<Long>> activeNotificationOpenTimeSlots = getActiveSlotsForTestShortcut(
+                Event.NOTIFICATION_EVENT_TYPES);
         assertEquals(1, activeNotificationOpenTimeSlots.size());
     }
 
@@ -469,12 +442,7 @@ public final class DataManagerTest {
         mInjector.mCallLogQueryHelper.mEventConsumer.accept(PHONE_NUMBER,
                 new Event(currentTimestamp - MILLIS_PER_MINUTE * 5L, Event.TYPE_CALL_MISSED));
 
-        List<Range<Long>> activeTimeSlots = new ArrayList<>();
-        mDataManager.forAllPackages(packageData ->
-                activeTimeSlots.addAll(
-                        packageData.getEventHistory(TEST_SHORTCUT_ID)
-                                .getEventIndex(Event.CALL_EVENT_TYPES)
-                                .getActiveTimeSlots()));
+        List<Range<Long>> activeTimeSlots = getActiveSlotsForTestShortcut(Event.CALL_EVENT_TYPES);
         assertEquals(3, activeTimeSlots.size());
     }
 
@@ -498,12 +466,7 @@ public final class DataManagerTest {
         mInjector.mMmsQueryHelper.mEventConsumer.accept(PHONE_NUMBER, outgoingSmsEvent);
         mInjector.mSmsQueryHelper.mEventConsumer.accept(PHONE_NUMBER, incomingSmsEvent);
 
-        List<Range<Long>> activeTimeSlots = new ArrayList<>();
-        mDataManager.forAllPackages(packageData ->
-                activeTimeSlots.addAll(
-                        packageData.getEventHistory(TEST_SHORTCUT_ID)
-                                .getEventIndex(Event.SMS_EVENT_TYPES)
-                                .getActiveTimeSlots()));
+        List<Range<Long>> activeTimeSlots = getActiveSlotsForTestShortcut(Event.SMS_EVENT_TYPES);
         assertEquals(2, activeTimeSlots.size());
     }
 
@@ -551,22 +514,12 @@ public final class DataManagerTest {
         mInjector.mCallLogQueryHelper.mEventConsumer.accept(PHONE_NUMBER,
                 new Event(currentTimestamp - MILLIS_PER_MINUTE, Event.TYPE_CALL_OUTGOING));
 
-        List<Range<Long>> activeTimeSlots = new ArrayList<>();
-        mDataManager.forAllPackages(packageData ->
-                activeTimeSlots.addAll(
-                        packageData.getEventHistory(TEST_SHORTCUT_ID)
-                                .getEventIndex(Event.CALL_EVENT_TYPES)
-                                .getActiveTimeSlots()));
+        List<Range<Long>> activeTimeSlots = getActiveSlotsForTestShortcut(Event.CALL_EVENT_TYPES);
         assertEquals(1, activeTimeSlots.size());
 
         mDataManager.getUserDataForTesting(USER_ID_PRIMARY).setDefaultDialer(null);
         mDataManager.pruneDataForUser(USER_ID_PRIMARY, mCancellationSignal);
-        activeTimeSlots.clear();
-        mDataManager.forAllPackages(packageData ->
-                activeTimeSlots.addAll(
-                        packageData.getEventHistory(TEST_SHORTCUT_ID)
-                                .getEventIndex(Event.CALL_EVENT_TYPES)
-                                .getActiveTimeSlots()));
+        activeTimeSlots = getActiveSlotsForTestShortcut(Event.CALL_EVENT_TYPES);
         assertTrue(activeTimeSlots.isEmpty());
     }
 
@@ -583,22 +536,12 @@ public final class DataManagerTest {
         mInjector.mMmsQueryHelper.mEventConsumer.accept(PHONE_NUMBER,
                 new Event(currentTimestamp - MILLIS_PER_MINUTE, Event.TYPE_SMS_OUTGOING));
 
-        List<Range<Long>> activeTimeSlots = new ArrayList<>();
-        mDataManager.forAllPackages(packageData ->
-                activeTimeSlots.addAll(
-                        packageData.getEventHistory(TEST_SHORTCUT_ID)
-                                .getEventIndex(Event.SMS_EVENT_TYPES)
-                                .getActiveTimeSlots()));
+        List<Range<Long>> activeTimeSlots = getActiveSlotsForTestShortcut(Event.SMS_EVENT_TYPES);
         assertEquals(1, activeTimeSlots.size());
 
         mDataManager.getUserDataForTesting(USER_ID_PRIMARY).setDefaultSmsApp(null);
         mDataManager.pruneDataForUser(USER_ID_PRIMARY, mCancellationSignal);
-        activeTimeSlots.clear();
-        mDataManager.forAllPackages(packageData ->
-                activeTimeSlots.addAll(
-                        packageData.getEventHistory(TEST_SHORTCUT_ID)
-                                .getEventIndex(Event.SMS_EVENT_TYPES)
-                                .getActiveTimeSlots()));
+        activeTimeSlots = getActiveSlotsForTestShortcut(Event.SMS_EVENT_TYPES);
         assertTrue(activeTimeSlots.isEmpty());
     }
 
@@ -607,6 +550,24 @@ public final class DataManagerTest {
         LocalServices.addService(clazz, mock);
     }
 
+    private List<ConversationInfo> getConversationsInPrimary() {
+        List<ConversationInfo> conversations = new ArrayList<>();
+        mDataManager.forPackagesInProfile(USER_ID_PRIMARY,
+                packageData -> packageData.forAllConversations(conversations::add));
+        return conversations;
+    }
+
+    private List<Range<Long>> getActiveSlotsForTestShortcut(
+            Set<Integer> eventTypes) {
+        List<Range<Long>> activeSlots = new ArrayList<>();
+        mDataManager.forPackagesInProfile(USER_ID_PRIMARY, packageData ->
+                activeSlots.addAll(
+                        packageData.getEventHistory(TEST_SHORTCUT_ID)
+                                .getEventIndex(eventTypes)
+                                .getActiveTimeSlots()));
+        return activeSlots;
+    }
+
     private ShortcutInfo buildShortcutInfo(String packageName, int userId, String id,
             @Nullable Person person) {
         Context mockContext = mock(Context.class);
@@ -778,10 +739,5 @@ public final class DataManagerTest {
             mSmsQueryHelper = new TestSmsQueryHelper(context, eventConsumer);
             return mSmsQueryHelper;
         }
-
-        @Override
-        int getCallingUserId() {
-            return mCallingUserId;
-        }
     }
 }
diff --git a/services/tests/servicestests/src/com/android/server/people/prediction/ShareTargetPredictorTest.java b/services/tests/servicestests/src/com/android/server/people/prediction/ShareTargetPredictorTest.java
index 808906e3a06a..f498a9450c9e 100644
--- a/services/tests/servicestests/src/com/android/server/people/prediction/ShareTargetPredictorTest.java
+++ b/services/tests/servicestests/src/com/android/server/people/prediction/ShareTargetPredictorTest.java
@@ -20,6 +20,7 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -73,7 +74,7 @@ public final class ShareTargetPredictorTest {
     public void setUp() {
         MockitoAnnotations.initMocks(this);
 
-        when(mDataManager.getShareShortcuts(any())).thenReturn(mShareShortcuts);
+        when(mDataManager.getShareShortcuts(any(), anyInt())).thenReturn(mShareShortcuts);
         when(mDataManager.getPackage(PACKAGE_1, USER_ID)).thenReturn(mPackageData1);
         when(mDataManager.getPackage(PACKAGE_2, USER_ID)).thenReturn(mPackageData2);
 
@@ -82,7 +83,8 @@ public final class ShareTargetPredictorTest {
                 .setPredictedTargetCount(NUM_PREDICTED_TARGETS)
                 .setExtras(new Bundle())
                 .build();
-        mPredictor = new ShareTargetPredictor(predictionContext, targets -> { }, mDataManager);
+        mPredictor = new ShareTargetPredictor(
+                predictionContext, targets -> { }, mDataManager, USER_ID);
     }
 
     @Test