Merge "Extract shortcuts loading logic from ChooserActivity" into tm-qpr-dev

7 files changed, 674 insertions, 335 deletions

diff --git a/java/tests/src/com/android/intentresolver/ChooserActivityOverrideData.java b/java/tests/src/com/android/intentresolver/ChooserActivityOverrideData.java
index e474938b..dd78b69e 100644
--- a/java/tests/src/com/android/intentresolver/ChooserActivityOverrideData.java
+++ b/java/tests/src/com/android/intentresolver/ChooserActivityOverrideData.java
@@ -24,15 +24,17 @@ import android.content.res.Resources;
 import android.database.Cursor;
 import android.graphics.Bitmap;
 import android.os.UserHandle;
-import android.util.Pair;
 
 import com.android.intentresolver.chooser.TargetInfo;
+import com.android.intentresolver.shortcuts.ShortcutLoader;
 import com.android.internal.logging.MetricsLogger;
 
 import java.util.List;
-import java.util.function.BiFunction;
+import java.util.function.Consumer;
 import java.util.function.Function;
 
+import kotlin.jvm.functions.Function2;
+
 /**
  * Singleton providing overrides to be applied by any {@code IChooserWrapper} used in testing.
  * We cannot directly mock the activity created since instrumentation creates it, so instead we use
@@ -51,10 +53,8 @@ public class ChooserActivityOverrideData {
     @SuppressWarnings("Since15")
     public Function<PackageManager, PackageManager> createPackageManager;
     public Function<TargetInfo, Boolean> onSafelyStartCallback;
-    public Function<ChooserListAdapter, Void> onQueryDirectShareTargets;
-    public BiFunction<
-            IChooserWrapper, ChooserListAdapter, Pair<Integer, ChooserActivity.ServiceResultInfo[]>>
-            directShareTargets;
+    public Function2<UserHandle, Consumer<ShortcutLoader.Result>, ShortcutLoader>
+            shortcutLoaderFactory = (userHandle, callback) -> null;
     public ResolverListController resolverListController;
     public ResolverListController workResolverListController;
     public Boolean isVoiceInteraction;
@@ -69,15 +69,11 @@ public class ChooserActivityOverrideData {
     public UserHandle workProfileUserHandle;
     public boolean hasCrossProfileIntents;
     public boolean isQuietModeEnabled;
-    public boolean isWorkProfileUserRunning;
-    public boolean isWorkProfileUserUnlocked;
     public AbstractMultiProfilePagerAdapter.Injector multiPagerAdapterInjector;
     public PackageManager packageManager;
 
     public void reset() {
         onSafelyStartCallback = null;
-        onQueryDirectShareTargets = null;
-        directShareTargets = null;
         isVoiceInteraction = null;
         createPackageManager = null;
         previewThumbnail = null;
@@ -93,8 +89,6 @@ public class ChooserActivityOverrideData {
         workProfileUserHandle = null;
         hasCrossProfileIntents = true;
         isQuietModeEnabled = false;
-        isWorkProfileUserRunning = true;
-        isWorkProfileUserUnlocked = true;
         packageManager = null;
         multiPagerAdapterInjector = new AbstractMultiProfilePagerAdapter.Injector() {
             @Override
@@ -114,6 +108,7 @@ public class ChooserActivityOverrideData {
                 isQuietModeEnabled = enabled;
             }
         };
+        shortcutLoaderFactory = ((userHandle, resultConsumer) -> null);
     }
 
     private ChooserActivityOverrideData() {}
diff --git a/java/tests/src/com/android/intentresolver/ChooserWrapperActivity.java b/java/tests/src/com/android/intentresolver/ChooserWrapperActivity.java
index 8c7c28bb..6b74fcd4 100644
--- a/java/tests/src/com/android/intentresolver/ChooserWrapperActivity.java
+++ b/java/tests/src/com/android/intentresolver/ChooserWrapperActivity.java
@@ -19,11 +19,13 @@ package com.android.intentresolver;
 import static org.mockito.Mockito.when;
 
 import android.annotation.Nullable;
+import android.app.prediction.AppPredictor;
 import android.app.usage.UsageStatsManager;
 import android.content.ComponentName;
 import android.content.ContentResolver;
 import android.content.Context;
 import android.content.Intent;
+import android.content.IntentFilter;
 import android.content.pm.PackageManager;
 import android.content.pm.ResolveInfo;
 import android.content.res.Resources;
@@ -31,7 +33,6 @@ import android.database.Cursor;
 import android.graphics.Bitmap;
 import android.net.Uri;
 import android.os.UserHandle;
-import android.util.Pair;
 import android.util.Size;
 
 import com.android.intentresolver.AbstractMultiProfilePagerAdapter;
@@ -44,11 +45,12 @@ import com.android.intentresolver.ResolverListController;
 import com.android.intentresolver.chooser.DisplayResolveInfo;
 import com.android.intentresolver.chooser.NotSelectableTargetInfo;
 import com.android.intentresolver.chooser.TargetInfo;
+import com.android.intentresolver.shortcuts.ShortcutLoader;
 import com.android.internal.logging.MetricsLogger;
 import com.android.internal.logging.nano.MetricsProto.MetricsEvent;
 
-import java.util.Arrays;
 import java.util.List;
+import java.util.function.Consumer;
 
 /**
  * Simple wrapper around chooser activity to be able to initiate it under test. For more
@@ -256,41 +258,18 @@ public class ChooserWrapperActivity
     }
 
     @Override
-    protected void queryDirectShareTargets(
-            ChooserListAdapter adapter, boolean skipAppPredictionService) {
-        if (sOverrides.directShareTargets != null) {
-            Pair<Integer, ServiceResultInfo[]> result =
-                    sOverrides.directShareTargets.apply(this, adapter);
-            // Imitate asynchronous shortcut loading
-            getMainExecutor().execute(
-                    () -> onShortcutsLoaded(
-                            adapter, result.first, Arrays.asList(result.second)));
-            return;
-        }
-        if (sOverrides.onQueryDirectShareTargets != null) {
-            sOverrides.onQueryDirectShareTargets.apply(adapter);
-        }
-        super.queryDirectShareTargets(adapter, skipAppPredictionService);
-    }
-
-    @Override
-    protected boolean isQuietModeEnabled(UserHandle userHandle) {
-        return sOverrides.isQuietModeEnabled;
-    }
-
-    @Override
-    protected boolean isUserRunning(UserHandle userHandle) {
-        if (userHandle.equals(UserHandle.SYSTEM)) {
-            return super.isUserRunning(userHandle);
-        }
-        return sOverrides.isWorkProfileUserRunning;
-    }
-
-    @Override
-    protected boolean isUserUnlocked(UserHandle userHandle) {
-        if (userHandle.equals(UserHandle.SYSTEM)) {
-            return super.isUserUnlocked(userHandle);
+    protected ShortcutLoader createShortcutLoader(
+            Context context,
+            AppPredictor appPredictor,
+            UserHandle userHandle,
+            IntentFilter targetIntentFilter,
+            Consumer<ShortcutLoader.Result> callback) {
+        ShortcutLoader shortcutLoader =
+                sOverrides.shortcutLoaderFactory.invoke(userHandle, callback);
+        if (shortcutLoader != null) {
+            return shortcutLoader;
         }
-        return sOverrides.isWorkProfileUserUnlocked;
+        return super.createShortcutLoader(
+                context, appPredictor, userHandle, targetIntentFilter, callback);
     }
 }
diff --git a/java/tests/src/com/android/intentresolver/IChooserWrapper.java b/java/tests/src/com/android/intentresolver/IChooserWrapper.java
index f81cd023..0d44e147 100644
--- a/java/tests/src/com/android/intentresolver/IChooserWrapper.java
+++ b/java/tests/src/com/android/intentresolver/IChooserWrapper.java
@@ -25,6 +25,8 @@ import android.os.UserHandle;
 import com.android.intentresolver.ResolverListAdapter.ResolveInfoPresentationGetter;
 import com.android.intentresolver.chooser.DisplayResolveInfo;
 
+import java.util.concurrent.Executor;
+
 /**
  * Test-only extended API capabilities that an instrumented ChooserActivity subclass provides in
  * order to expose the internals for override/inspection. Implementations should apply the overrides
@@ -41,4 +43,5 @@ public interface IChooserWrapper {
             @Nullable ResolveInfoPresentationGetter resolveInfoPresentationGetter);
     UserHandle getCurrentUserHandle();
     ChooserActivityLogger getChooserActivityLogger();
+    Executor getMainExecutor();
 }
diff --git a/java/tests/src/com/android/intentresolver/TestApplication.kt b/java/tests/src/com/android/intentresolver/TestApplication.kt
new file mode 100644
index 00000000..849cfbab
--- /dev/null
+++ b/java/tests/src/com/android/intentresolver/TestApplication.kt
@@ -0,0 +1,27 @@
+/*
+ * Copyright (C) 2022 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.intentresolver
+
+import android.app.Application
+import android.content.Context
+import android.os.UserHandle
+
+class TestApplication : Application() {
+
+    // return the current context as a work profile doesn't really exist in these tests
+    override fun createContextAsUser(user: UserHandle, flags: Int): Context = this
+}
\ No newline at end of file
diff --git a/java/tests/src/com/android/intentresolver/UnbundledChooserActivityTest.java b/java/tests/src/com/android/intentresolver/UnbundledChooserActivityTest.java
index 7c304284..da72a749 100644
--- a/java/tests/src/com/android/intentresolver/UnbundledChooserActivityTest.java
+++ b/java/tests/src/com/android/intentresolver/UnbundledChooserActivityTest.java
@@ -38,7 +38,6 @@ import static com.android.intentresolver.MatcherUtils.first;
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static junit.framework.Assert.assertFalse;
 import static junit.framework.Assert.assertNull;
 
 import static org.hamcrest.CoreMatchers.allOf;
@@ -83,6 +82,7 @@ import android.os.UserHandle;
 import android.provider.DeviceConfig;
 import android.service.chooser.ChooserTarget;
 import android.util.Pair;
+import android.util.SparseArray;
 import android.view.View;
 
 import androidx.annotation.CallSuper;
@@ -93,9 +93,9 @@ import androidx.test.espresso.matcher.BoundedDiagnosingMatcher;
 import androidx.test.platform.app.InstrumentationRegistry;
 import androidx.test.rule.ActivityTestRule;
 
-import com.android.intentresolver.ChooserActivity.ServiceResultInfo;
 import com.android.intentresolver.ResolverActivity.ResolvedComponentInfo;
 import com.android.intentresolver.chooser.DisplayResolveInfo;
+import com.android.intentresolver.shortcuts.ShortcutLoader;
 import com.android.internal.config.sysui.SystemUiDeviceConfigFlags;
 import com.android.internal.logging.MetricsLogger;
 import com.android.internal.logging.nano.MetricsProto.MetricsEvent;
@@ -118,6 +118,7 @@ import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Consumer;
 import java.util.function.Function;
 
 /**
@@ -1279,7 +1280,7 @@ public class UnbundledChooserActivityTest {
     }
 
     // This test is too long and too slow and should not be taken as an example for future tests.
-    @Test @Ignore
+    @Test
     public void testDirectTargetSelectionLogging() {
         Intent sendIntent = createSendTextIntent();
         // We need app targets for direct targets to get displayed
@@ -1298,37 +1299,55 @@ public class UnbundledChooserActivityTest {
         // Set up resources
         MetricsLogger mockLogger = ChooserActivityOverrideData.getInstance().metricsLogger;
         ArgumentCaptor<LogMaker> logMakerCaptor = ArgumentCaptor.forClass(LogMaker.class);
-        // Create direct share target
-        List<ChooserTarget> serviceTargets = createDirectShareTargets(1, "");
-        ResolveInfo ri = ResolverDataProvider.createResolveInfo(3, 0);
+
+        // create test shortcut loader factory, remember loaders and their callbacks
+        SparseArray<Pair<ShortcutLoader, Consumer<ShortcutLoader.Result>>> shortcutLoaders =
+                createShortcutLoaderFactory();
 
         // Start activity
         final IChooserWrapper activity = (IChooserWrapper)
                 mActivityRule.launchActivity(Intent.createChooser(sendIntent, null));
+        waitForIdle();
 
-        // Insert the direct share target
-        Map<ChooserTarget, ShortcutInfo> directShareToShortcutInfos = new HashMap<>();
-        directShareToShortcutInfos.put(serviceTargets.get(0), null);
-        InstrumentationRegistry.getInstrumentation().runOnMainSync(
-                () -> activity.getAdapter().addServiceResults(
-                        activity.createTestDisplayResolveInfo(sendIntent,
-                                ri,
-                                "testLabel",
-                                "testInfo",
-                                sendIntent,
-                                /* resolveInfoPresentationGetter */ null),
-                        serviceTargets,
-                        TARGET_TYPE_CHOOSER_TARGET,
-                        directShareToShortcutInfos,
-                        /* directShareToAppTargets */ null)
+        // verify that ShortcutLoader was queried
+        ArgumentCaptor<DisplayResolveInfo[]> appTargets =
+                ArgumentCaptor.forClass(DisplayResolveInfo[].class);
+        verify(shortcutLoaders.get(0).first, times(1)).queryShortcuts(appTargets.capture());
+
+        // send shortcuts
+        assertThat(
+                "Wrong number of app targets",
+                appTargets.getValue().length,
+                is(resolvedComponentInfos.size()));
+        List<ChooserTarget> serviceTargets = createDirectShareTargets(1, "");
+        ShortcutLoader.Result result = new ShortcutLoader.Result(
+                true,
+                appTargets.getValue(),
+                new ShortcutLoader.ShortcutResultInfo[] {
+                        new ShortcutLoader.ShortcutResultInfo(
+                                appTargets.getValue()[0],
+                                serviceTargets
+                        )
+                },
+                new HashMap<>(),
+                new HashMap<>()
         );
+        activity.getMainExecutor().execute(() -> shortcutLoaders.get(0).second.accept(result));
+        waitForIdle();
 
-        assertThat("Chooser should have 3 targets (2 apps, 1 direct)",
-                activity.getAdapter().getCount(), is(3));
-        assertThat("Chooser should have exactly one selectable direct target",
-                activity.getAdapter().getSelectableServiceTargetCount(), is(1));
-        assertThat("The resolver info must match the resolver info used to create the target",
-                activity.getAdapter().getItem(0).getResolveInfo(), is(ri));
+        final ChooserListAdapter activeAdapter = activity.getAdapter();
+        assertThat(
+                "Chooser should have 3 targets (2 apps, 1 direct)",
+                activeAdapter.getCount(),
+                is(3));
+        assertThat(
+                "Chooser should have exactly one selectable direct target",
+                activeAdapter.getSelectableServiceTargetCount(),
+                is(1));
+        assertThat(
+                "The resolver info must match the resolver info used to create the target",
+                activeAdapter.getItem(0).getResolveInfo(),
+                is(resolvedComponentInfos.get(0).getResolveInfoAt(0)));
 
         // Click on the direct target
         String name = serviceTargets.get(0).getTitle().toString();
@@ -1336,23 +1355,29 @@ public class UnbundledChooserActivityTest {
                 .perform(click());
         waitForIdle();
 
-        // Currently we're seeing 3 invocations
-        //      1. ChooserActivity.onCreate()
-        //      2. ChooserActivity$ChooserRowAdapter.createContentPreviewView()
-        //      3. ChooserActivity.startSelected -- which is the one we're after
-        verify(mockLogger, Mockito.times(3)).write(logMakerCaptor.capture());
-        assertThat(logMakerCaptor.getAllValues().get(2).getCategory(),
+        // Currently we're seeing 4 invocations
+        //      1. ChooserActivity.logActionShareWithPreview()
+        //      2. ChooserActivity.onCreate()
+        //      3. ChooserActivity.logDirectShareTargetReceived()
+        //      4. ChooserActivity.startSelected -- which is the one we're after
+        verify(mockLogger, Mockito.times(4)).write(logMakerCaptor.capture());
+        LogMaker selectionLog = logMakerCaptor.getAllValues().get(3);
+        assertThat(
+                selectionLog.getCategory(),
                 is(MetricsEvent.ACTION_ACTIVITY_CHOOSER_PICKED_SERVICE_TARGET));
-        String hashedName = (String) logMakerCaptor
-                .getAllValues().get(2).getTaggedData(MetricsEvent.FIELD_HASHED_TARGET_NAME);
-        assertThat("Hash is not predictable but must be obfuscated",
+        String hashedName = (String) selectionLog.getTaggedData(
+                MetricsEvent.FIELD_HASHED_TARGET_NAME);
+        assertThat(
+                "Hash is not predictable but must be obfuscated",
                 hashedName, is(not(name)));
-        assertThat("The packages shouldn't match for app target and direct target", logMakerCaptor
-                .getAllValues().get(2).getTaggedData(MetricsEvent.FIELD_RANKED_POSITION), is(-1));
+        assertThat(
+                "The packages shouldn't match for app target and direct target",
+                selectionLog.getTaggedData(MetricsEvent.FIELD_RANKED_POSITION),
+                is(-1));
     }
 
     // This test is too long and too slow and should not be taken as an example for future tests.
-    @Test @Ignore
+    @Test
     public void testDirectTargetLoggingWithRankedAppTarget() {
         Intent sendIntent = createSendTextIntent();
         // We need app targets for direct targets to get displayed
@@ -1371,38 +1396,57 @@ public class UnbundledChooserActivityTest {
         // Set up resources
         MetricsLogger mockLogger = ChooserActivityOverrideData.getInstance().metricsLogger;
         ArgumentCaptor<LogMaker> logMakerCaptor = ArgumentCaptor.forClass(LogMaker.class);
-        // Create direct share target
-        List<ChooserTarget> serviceTargets = createDirectShareTargets(1,
-                resolvedComponentInfos.get(0).getResolveInfoAt(0).activityInfo.packageName);
-        ResolveInfo ri = ResolverDataProvider.createResolveInfo(3, 0);
+
+        // create test shortcut loader factory, remember loaders and their callbacks
+        SparseArray<Pair<ShortcutLoader, Consumer<ShortcutLoader.Result>>> shortcutLoaders =
+                createShortcutLoaderFactory();
 
         // Start activity
         final IChooserWrapper activity = (IChooserWrapper)
                 mActivityRule.launchActivity(Intent.createChooser(sendIntent, null));
+        waitForIdle();
 
-        // Insert the direct share target
-        Map<ChooserTarget, ShortcutInfo> directShareToShortcutInfos = new HashMap<>();
-        directShareToShortcutInfos.put(serviceTargets.get(0), null);
-        InstrumentationRegistry.getInstrumentation().runOnMainSync(
-                () -> activity.getAdapter().addServiceResults(
-                        activity.createTestDisplayResolveInfo(sendIntent,
-                                ri,
-                                "testLabel",
-                                "testInfo",
-                                sendIntent,
-                                /* resolveInfoPresentationGetter */ null),
-                        serviceTargets,
-                        TARGET_TYPE_CHOOSER_TARGET,
-                        directShareToShortcutInfos,
-                        /* directShareToAppTargets */ null)
+        // verify that ShortcutLoader was queried
+        ArgumentCaptor<DisplayResolveInfo[]> appTargets =
+                ArgumentCaptor.forClass(DisplayResolveInfo[].class);
+        verify(shortcutLoaders.get(0).first, times(1)).queryShortcuts(appTargets.capture());
+
+        // send shortcuts
+        assertThat(
+                "Wrong number of app targets",
+                appTargets.getValue().length,
+                is(resolvedComponentInfos.size()));
+        List<ChooserTarget> serviceTargets = createDirectShareTargets(
+                1,
+                resolvedComponentInfos.get(0).getResolveInfoAt(0).activityInfo.packageName);
+        ShortcutLoader.Result result = new ShortcutLoader.Result(
+                true,
+                appTargets.getValue(),
+                new ShortcutLoader.ShortcutResultInfo[] {
+                        new ShortcutLoader.ShortcutResultInfo(
+                                appTargets.getValue()[0],
+                                serviceTargets
+                        )
+                },
+                new HashMap<>(),
+                new HashMap<>()
         );
+        activity.getMainExecutor().execute(() -> shortcutLoaders.get(0).second.accept(result));
+        waitForIdle();
 
-        assertThat("Chooser should have 3 targets (2 apps, 1 direct)",
-                activity.getAdapter().getCount(), is(3));
-        assertThat("Chooser should have exactly one selectable direct target",
-                activity.getAdapter().getSelectableServiceTargetCount(), is(1));
-        assertThat("The resolver info must match the resolver info used to create the target",
-                activity.getAdapter().getItem(0).getResolveInfo(), is(ri));
+        final ChooserListAdapter activeAdapter = activity.getAdapter();
+        assertThat(
+                "Chooser should have 3 targets (2 apps, 1 direct)",
+                activeAdapter.getCount(),
+                is(3));
+        assertThat(
+                "Chooser should have exactly one selectable direct target",
+                activeAdapter.getSelectableServiceTargetCount(),
+                is(1));
+        assertThat(
+                "The resolver info must match the resolver info used to create the target",
+                activeAdapter.getItem(0).getResolveInfo(),
+                is(resolvedComponentInfos.get(0).getResolveInfoAt(0)));
 
         // Click on the direct target
         String name = serviceTargets.get(0).getTitle().toString();
@@ -1410,18 +1454,19 @@ public class UnbundledChooserActivityTest {
                 .perform(click());
         waitForIdle();
 
-        // Currently we're seeing 3 invocations
-        //      1. ChooserActivity.onCreate()
-        //      2. ChooserActivity$ChooserRowAdapter.createContentPreviewView()
-        //      3. ChooserActivity.startSelected -- which is the one we're after
-        verify(mockLogger, Mockito.times(3)).write(logMakerCaptor.capture());
-        assertThat(logMakerCaptor.getAllValues().get(2).getCategory(),
+        // Currently we're seeing 4 invocations
+        //      1. ChooserActivity.logActionShareWithPreview()
+        //      2. ChooserActivity.onCreate()
+        //      3. ChooserActivity.logDirectShareTargetReceived()
+        //      4. ChooserActivity.startSelected -- which is the one we're after
+        verify(mockLogger, Mockito.times(4)).write(logMakerCaptor.capture());
+        assertThat(logMakerCaptor.getAllValues().get(3).getCategory(),
                 is(MetricsEvent.ACTION_ACTIVITY_CHOOSER_PICKED_SERVICE_TARGET));
         assertThat("The packages should match for app target and direct target", logMakerCaptor
-                .getAllValues().get(2).getTaggedData(MetricsEvent.FIELD_RANKED_POSITION), is(0));
+                .getAllValues().get(3).getTaggedData(MetricsEvent.FIELD_RANKED_POSITION), is(0));
     }
 
-    @Test @Ignore
+    @Test
     public void testShortcutTargetWithApplyAppLimits() {
         // Set up resources
         ChooserActivityOverrideData.getInstance().resources = Mockito.spy(
@@ -1445,48 +1490,64 @@ public class UnbundledChooserActivityTest {
                                 Mockito.anyBoolean(),
                                 Mockito.isA(List.class)))
                 .thenReturn(resolvedComponentInfos);
-        // Create direct share target
-        List<ChooserTarget> serviceTargets = createDirectShareTargets(2,
-                resolvedComponentInfos.get(0).getResolveInfoAt(0).activityInfo.packageName);
-        ResolveInfo ri = ResolverDataProvider.createResolveInfo(3, 0);
+
+        // create test shortcut loader factory, remember loaders and their callbacks
+        SparseArray<Pair<ShortcutLoader, Consumer<ShortcutLoader.Result>>> shortcutLoaders =
+                createShortcutLoaderFactory();
 
         // Start activity
-        final ChooserActivity activity =
-                mActivityRule.launchActivity(Intent.createChooser(sendIntent, null));
-        final IChooserWrapper wrapper = (IChooserWrapper) activity;
+        final IChooserWrapper activity = (IChooserWrapper) mActivityRule
+                .launchActivity(Intent.createChooser(sendIntent, null));
+        waitForIdle();
 
-        // Insert the direct share target
-        Map<ChooserTarget, ShortcutInfo> directShareToShortcutInfos = new HashMap<>();
-        List<ShareShortcutInfo> shortcutInfos = createShortcuts(activity);
-        directShareToShortcutInfos.put(serviceTargets.get(0),
-                shortcutInfos.get(0).getShortcutInfo());
-        directShareToShortcutInfos.put(serviceTargets.get(1),
-                shortcutInfos.get(1).getShortcutInfo());
-        InstrumentationRegistry.getInstrumentation().runOnMainSync(
-                () -> wrapper.getAdapter().addServiceResults(
-                        wrapper.createTestDisplayResolveInfo(sendIntent,
-                                ri,
-                                "testLabel",
-                                "testInfo",
-                                sendIntent,
-                                /* resolveInfoPresentationGetter */ null),
-                        serviceTargets,
-                        TARGET_TYPE_SHORTCUTS_FROM_PREDICTION_SERVICE,
-                        directShareToShortcutInfos,
-                        /* directShareToAppTargets */ null)
+        // verify that ShortcutLoader was queried
+        ArgumentCaptor<DisplayResolveInfo[]> appTargets =
+                ArgumentCaptor.forClass(DisplayResolveInfo[].class);
+        verify(shortcutLoaders.get(0).first, times(1)).queryShortcuts(appTargets.capture());
+
+        // send shortcuts
+        assertThat(
+                "Wrong number of app targets",
+                appTargets.getValue().length,
+                is(resolvedComponentInfos.size()));
+        List<ChooserTarget> serviceTargets = createDirectShareTargets(
+                2,
+                resolvedComponentInfos.get(0).getResolveInfoAt(0).activityInfo.packageName);
+        ShortcutLoader.Result result = new ShortcutLoader.Result(
+                true,
+                appTargets.getValue(),
+                new ShortcutLoader.ShortcutResultInfo[] {
+                        new ShortcutLoader.ShortcutResultInfo(
+                                appTargets.getValue()[0],
+                                serviceTargets
+                        )
+                },
+                new HashMap<>(),
+                new HashMap<>()
         );
+        activity.getMainExecutor().execute(() -> shortcutLoaders.get(0).second.accept(result));
+        waitForIdle();
 
-        assertThat("Chooser should have 3 targets (2 apps, 1 direct)",
-                wrapper.getAdapter().getCount(), is(3));
-        assertThat("Chooser should have exactly one selectable direct target",
-                wrapper.getAdapter().getSelectableServiceTargetCount(), is(1));
-        assertThat("The resolver info must match the resolver info used to create the target",
-                wrapper.getAdapter().getItem(0).getResolveInfo(), is(ri));
-        assertThat("The display label must match",
-                wrapper.getAdapter().getItem(0).getDisplayLabel(), is("testTitle0"));
+        final ChooserListAdapter activeAdapter = activity.getAdapter();
+        assertThat(
+                "Chooser should have 3 targets (2 apps, 1 direct)",
+                activeAdapter.getCount(),
+                is(3));
+        assertThat(
+                "Chooser should have exactly one selectable direct target",
+                activeAdapter.getSelectableServiceTargetCount(),
+                is(1));
+        assertThat(
+                "The resolver info must match the resolver info used to create the target",
+                activeAdapter.getItem(0).getResolveInfo(),
+                is(resolvedComponentInfos.get(0).getResolveInfoAt(0)));
+        assertThat(
+                "The display label must match",
+                activeAdapter.getItem(0).getDisplayLabel(),
+                is("testTitle0"));
     }
 
-    @Test @Ignore
+    @Test
     public void testShortcutTargetWithoutApplyAppLimits() {
         setDeviceConfigProperty(
                 SystemUiDeviceConfigFlags.APPLY_SHARING_APP_LIMITS_IN_SYSUI,
@@ -1513,47 +1574,65 @@ public class UnbundledChooserActivityTest {
                                 Mockito.anyBoolean(),
                                 Mockito.isA(List.class)))
                 .thenReturn(resolvedComponentInfos);
-        // Create direct share target
-        List<ChooserTarget> serviceTargets = createDirectShareTargets(2,
-                resolvedComponentInfos.get(0).getResolveInfoAt(0).activityInfo.packageName);
-        ResolveInfo ri = ResolverDataProvider.createResolveInfo(3, 0);
+
+        // create test shortcut loader factory, remember loaders and their callbacks
+        SparseArray<Pair<ShortcutLoader, Consumer<ShortcutLoader.Result>>> shortcutLoaders =
+                createShortcutLoaderFactory();
 
         // Start activity
-        final ChooserActivity activity =
+        final IChooserWrapper activity = (IChooserWrapper)
                 mActivityRule.launchActivity(Intent.createChooser(sendIntent, null));
-        final IChooserWrapper wrapper = (IChooserWrapper) activity;
+        waitForIdle();
 
-        // Insert the direct share target
-        Map<ChooserTarget, ShortcutInfo> directShareToShortcutInfos = new HashMap<>();
-        List<ShareShortcutInfo> shortcutInfos = createShortcuts(activity);
-        directShareToShortcutInfos.put(serviceTargets.get(0),
-                shortcutInfos.get(0).getShortcutInfo());
-        directShareToShortcutInfos.put(serviceTargets.get(1),
-                shortcutInfos.get(1).getShortcutInfo());
-        InstrumentationRegistry.getInstrumentation().runOnMainSync(
-                () -> wrapper.getAdapter().addServiceResults(
-                        wrapper.createTestDisplayResolveInfo(sendIntent,
-                                ri,
-                                "testLabel",
-                                "testInfo",
-                                sendIntent,
-                                /* resolveInfoPresentationGetter */ null),
-                        serviceTargets,
-                        TARGET_TYPE_SHORTCUTS_FROM_PREDICTION_SERVICE,
-                        directShareToShortcutInfos,
-                        /* directShareToAppTargets */ null)
+        // verify that ShortcutLoader was queried
+        ArgumentCaptor<DisplayResolveInfo[]> appTargets =
+                ArgumentCaptor.forClass(DisplayResolveInfo[].class);
+        verify(shortcutLoaders.get(0).first, times(1)).queryShortcuts(appTargets.capture());
+
+        // send shortcuts
+        assertThat(
+                "Wrong number of app targets",
+                appTargets.getValue().length,
+                is(resolvedComponentInfos.size()));
+        List<ChooserTarget> serviceTargets = createDirectShareTargets(
+                2,
+                resolvedComponentInfos.get(0).getResolveInfoAt(0).activityInfo.packageName);
+        ShortcutLoader.Result result = new ShortcutLoader.Result(
+                true,
+                appTargets.getValue(),
+                new ShortcutLoader.ShortcutResultInfo[] {
+                        new ShortcutLoader.ShortcutResultInfo(
+                                appTargets.getValue()[0],
+                                serviceTargets
+                        )
+                },
+                new HashMap<>(),
+                new HashMap<>()
         );
+        activity.getMainExecutor().execute(() -> shortcutLoaders.get(0).second.accept(result));
+        waitForIdle();
 
-        assertThat("Chooser should have 4 targets (2 apps, 2 direct)",
-                wrapper.getAdapter().getCount(), is(4));
-        assertThat("Chooser should have exactly two selectable direct target",
-                wrapper.getAdapter().getSelectableServiceTargetCount(), is(2));
-        assertThat("The resolver info must match the resolver info used to create the target",
-                wrapper.getAdapter().getItem(0).getResolveInfo(), is(ri));
-        assertThat("The display label must match",
-                wrapper.getAdapter().getItem(0).getDisplayLabel(), is("testTitle0"));
-        assertThat("The display label must match",
-                wrapper.getAdapter().getItem(1).getDisplayLabel(), is("testTitle1"));
+        final ChooserListAdapter activeAdapter = activity.getAdapter();
+        assertThat(
+                "Chooser should have 4 targets (2 apps, 2 direct)",
+                activeAdapter.getCount(),
+                is(4));
+        assertThat(
+                "Chooser should have exactly two selectable direct target",
+                activeAdapter.getSelectableServiceTargetCount(),
+                is(2));
+        assertThat(
+                "The resolver info must match the resolver info used to create the target",
+                activeAdapter.getItem(0).getResolveInfo(),
+                is(resolvedComponentInfos.get(0).getResolveInfoAt(0)));
+        assertThat(
+                "The display label must match",
+                activeAdapter.getItem(0).getDisplayLabel(),
+                is("testTitle0"));
+        assertThat(
+                "The display label must match",
+                activeAdapter.getItem(1).getDisplayLabel(),
+                is("testTitle1"));
     }
 
     @Test
@@ -1948,43 +2027,59 @@ public class UnbundledChooserActivityTest {
                                 Mockito.isA(List.class)))
                 .thenReturn(resolvedComponentInfos);
 
-        // Create direct share target
-        List<ChooserTarget> serviceTargets = createDirectShareTargets(1,
-                resolvedComponentInfos.get(0).getResolveInfoAt(0).activityInfo.packageName);
-        ResolveInfo ri = ResolverDataProvider.createResolveInfo(3, 0);
-
-        ChooserActivityOverrideData
-                .getInstance()
-                .directShareTargets = (activity, adapter) -> {
-                    DisplayResolveInfo displayInfo = activity.createTestDisplayResolveInfo(
-                            sendIntent,
-                            ri,
-                             "testLabel",
-                             "testInfo",
-                            sendIntent,
-                            /* resolveInfoPresentationGetter */ null);
-                    ServiceResultInfo[] results = {
-                            new ServiceResultInfo(displayInfo, serviceTargets) };
-                    // TODO: consider covering the other type.
-                    //  Only 2 types are expected out of the shortcut loading logic:
-                    //  - TARGET_TYPE_SHORTCUTS_FROM_SHORTCUT_MANAGER, if shortcuts were loaded from
-                    //    the ShortcutManager, and;
-                    //  - TARGET_TYPE_SHORTCUTS_FROM_PREDICTION_SERVICE, if shortcuts were loaded
-                    //    from AppPredictor.
-                    //  Ideally, our tests should cover all of them.
-                    return new Pair<>(TARGET_TYPE_SHORTCUTS_FROM_SHORTCUT_MANAGER, results);
+        // create test shortcut loader factory, remember loaders and their callbacks
+        SparseArray<Pair<ShortcutLoader, Consumer<ShortcutLoader.Result>>> shortcutLoaders =
+                new SparseArray<>();
+        ChooserActivityOverrideData.getInstance().shortcutLoaderFactory =
+                (userHandle, callback) -> {
+                    Pair<ShortcutLoader, Consumer<ShortcutLoader.Result>> pair =
+                            new Pair<>(mock(ShortcutLoader.class), callback);
+                    shortcutLoaders.put(userHandle.getIdentifier(), pair);
+                    return pair.first;
                 };
 
         // Start activity
         final IChooserWrapper activity = (IChooserWrapper)
                 mActivityRule.launchActivity(Intent.createChooser(sendIntent, null));
+        waitForIdle();
+
+        // verify that ShortcutLoader was queried
+        ArgumentCaptor<DisplayResolveInfo[]> appTargets =
+                ArgumentCaptor.forClass(DisplayResolveInfo[].class);
+        verify(shortcutLoaders.get(0).first, times(1))
+                .queryShortcuts(appTargets.capture());
+
+        // send shortcuts
+        assertThat(
+                "Wrong number of app targets",
+                appTargets.getValue().length,
+                is(resolvedComponentInfos.size()));
+        List<ChooserTarget> serviceTargets = createDirectShareTargets(1,
+                resolvedComponentInfos.get(0).getResolveInfoAt(0).activityInfo.packageName);
+        ShortcutLoader.Result result = new ShortcutLoader.Result(
+                // TODO: test another value as well
+                false,
+                appTargets.getValue(),
+                new ShortcutLoader.ShortcutResultInfo[] {
+                        new ShortcutLoader.ShortcutResultInfo(
+                                appTargets.getValue()[0],
+                                serviceTargets
+                        )
+                },
+                new HashMap<>(),
+                new HashMap<>()
+        );
+        activity.getMainExecutor().execute(() -> shortcutLoaders.get(0).second.accept(result));
+        waitForIdle();
 
         assertThat("Chooser should have 3 targets (2 apps, 1 direct)",
                 activity.getAdapter().getCount(), is(3));
         assertThat("Chooser should have exactly one selectable direct target",
                 activity.getAdapter().getSelectableServiceTargetCount(), is(1));
-        assertThat("The resolver info must match the resolver info used to create the target",
-                activity.getAdapter().getItem(0).getResolveInfo(), is(ri));
+        assertThat(
+                "The resolver info must match the resolver info used to create the target",
+                activity.getAdapter().getItem(0).getResolveInfo(),
+                is(resolvedComponentInfos.get(0).getResolveInfoAt(0)));
 
         // Click on the direct target
         String name = serviceTargets.get(0).getTitle().toString();
@@ -2098,7 +2193,7 @@ public class UnbundledChooserActivityTest {
             return true;
         };
 
-        mActivityRule.launchActivity(sendIntent);
+        mActivityRule.launchActivity(Intent.createChooser(sendIntent, "Test"));
         waitForIdle();
 
         assertThat(chosen[0], is(personalResolvedComponentInfos.get(1).getResolveInfoAt(0)));
@@ -2273,21 +2368,20 @@ public class UnbundledChooserActivityTest {
     }
 
     @Test
-    public void testWorkTab_selectingWorkTabWithPausedWorkProfile_directShareTargetsNotQueried() {
+    public void test_query_shortcut_loader_for_the_selected_tab() {
         markWorkProfileUserAvailable();
         List<ResolvedComponentInfo> personalResolvedComponentInfos =
                 createResolvedComponentsForTestWithOtherProfile(3, /* userId */ 10);
         List<ResolvedComponentInfo> workResolvedComponentInfos =
                 createResolvedComponentsForTest(3);
         setupResolverControllers(personalResolvedComponentInfos, workResolvedComponentInfos);
-        ChooserActivityOverrideData.getInstance().isQuietModeEnabled = true;
-        boolean[] isQueryDirectShareCalledOnWorkProfile = new boolean[] { false };
-        ChooserActivityOverrideData.getInstance().onQueryDirectShareTargets =
-                chooserListAdapter -> {
-                    isQueryDirectShareCalledOnWorkProfile[0] =
-                            (chooserListAdapter.getUserHandle().getIdentifier() == 10);
-                    return null;
-                };
+        ShortcutLoader personalProfileShortcutLoader = mock(ShortcutLoader.class);
+        ShortcutLoader workProfileShortcutLoader = mock(ShortcutLoader.class);
+        final SparseArray<ShortcutLoader> shortcutLoaders = new SparseArray<>();
+        shortcutLoaders.put(0, personalProfileShortcutLoader);
+        shortcutLoaders.put(10, workProfileShortcutLoader);
+        ChooserActivityOverrideData.getInstance().shortcutLoaderFactory =
+                (userHandle, callback) -> shortcutLoaders.get(userHandle.getIdentifier(), null);
         Intent sendIntent = createSendTextIntent();
         sendIntent.setType(TEST_MIME_TYPE);
 
@@ -2295,118 +2389,14 @@ public class UnbundledChooserActivityTest {
         waitForIdle();
         onView(withId(com.android.internal.R.id.contentPanel))
                 .perform(swipeUp());
-        onView(withText(R.string.resolver_work_tab)).perform(click());
         waitForIdle();
 
-        assertFalse("Direct share targets were queried on a paused work profile",
-                isQueryDirectShareCalledOnWorkProfile[0]);
-    }
-
-    @Test
-    public void testWorkTab_selectingWorkTabWithNotRunningWorkUser_directShareTargetsNotQueried() {
-        markWorkProfileUserAvailable();
-        List<ResolvedComponentInfo> personalResolvedComponentInfos =
-                createResolvedComponentsForTestWithOtherProfile(3, /* userId */ 10);
-        List<ResolvedComponentInfo> workResolvedComponentInfos =
-                createResolvedComponentsForTest(3);
-        setupResolverControllers(personalResolvedComponentInfos, workResolvedComponentInfos);
-        ChooserActivityOverrideData.getInstance().isWorkProfileUserRunning = false;
-        boolean[] isQueryDirectShareCalledOnWorkProfile = new boolean[] { false };
-        ChooserActivityOverrideData.getInstance().onQueryDirectShareTargets =
-                chooserListAdapter -> {
-                    isQueryDirectShareCalledOnWorkProfile[0] =
-                            (chooserListAdapter.getUserHandle().getIdentifier() == 10);
-                    return null;
-                };
-        Intent sendIntent = createSendTextIntent();
-        sendIntent.setType(TEST_MIME_TYPE);
+        verify(personalProfileShortcutLoader, times(1)).queryShortcuts(any());
 
-        mActivityRule.launchActivity(Intent.createChooser(sendIntent, "work tab test"));
-        waitForIdle();
-        onView(withId(com.android.internal.R.id.contentPanel))
-                .perform(swipeUp());
         onView(withText(R.string.resolver_work_tab)).perform(click());
         waitForIdle();
 
-        assertFalse("Direct share targets were queried on a locked work profile user",
-                isQueryDirectShareCalledOnWorkProfile[0]);
-    }
-
-    @Test
-    public void testWorkTab_workUserNotRunning_workTargetsShown() {
-        markWorkProfileUserAvailable();
-        List<ResolvedComponentInfo> personalResolvedComponentInfos =
-                createResolvedComponentsForTestWithOtherProfile(3, /* userId */ 10);
-        List<ResolvedComponentInfo> workResolvedComponentInfos =
-                createResolvedComponentsForTest(3);
-        setupResolverControllers(personalResolvedComponentInfos, workResolvedComponentInfos);
-        Intent sendIntent = createSendTextIntent();
-        sendIntent.setType(TEST_MIME_TYPE);
-        ChooserActivityOverrideData.getInstance().isWorkProfileUserRunning = false;
-
-        final ChooserActivity activity =
-                mActivityRule.launchActivity(Intent.createChooser(sendIntent, "work tab test"));
-        final IChooserWrapper wrapper = (IChooserWrapper) activity;
-        waitForIdle();
-        onView(withId(com.android.internal.R.id.contentPanel)).perform(swipeUp());
-        onView(withText(R.string.resolver_work_tab)).perform(click());
-        waitForIdle();
-
-        assertEquals(3, wrapper.getWorkListAdapter().getCount());
-    }
-
-    @Test
-    public void testWorkTab_selectingWorkTabWithLockedWorkUser_directShareTargetsNotQueried() {
-        markWorkProfileUserAvailable();
-        List<ResolvedComponentInfo> personalResolvedComponentInfos =
-                createResolvedComponentsForTestWithOtherProfile(3, /* userId */ 10);
-        List<ResolvedComponentInfo> workResolvedComponentInfos =
-                createResolvedComponentsForTest(3);
-        setupResolverControllers(personalResolvedComponentInfos, workResolvedComponentInfos);
-        ChooserActivityOverrideData.getInstance().isWorkProfileUserUnlocked = false;
-        boolean[] isQueryDirectShareCalledOnWorkProfile = new boolean[] { false };
-        ChooserActivityOverrideData.getInstance().onQueryDirectShareTargets =
-                chooserListAdapter -> {
-                    isQueryDirectShareCalledOnWorkProfile[0] =
-                            (chooserListAdapter.getUserHandle().getIdentifier() == 10);
-                    return null;
-                };
-        Intent sendIntent = createSendTextIntent();
-        sendIntent.setType(TEST_MIME_TYPE);
-
-        mActivityRule.launchActivity(Intent.createChooser(sendIntent, "work tab test"));
-        waitForIdle();
-        onView(withId(com.android.internal.R.id.contentPanel))
-                .perform(swipeUp());
-        onView(withText(R.string.resolver_work_tab)).perform(click());
-        waitForIdle();
-
-        assertFalse("Direct share targets were queried on a locked work profile user",
-                isQueryDirectShareCalledOnWorkProfile[0]);
-    }
-
-    @Test
-    public void testWorkTab_workUserLocked_workTargetsShown() {
-        markWorkProfileUserAvailable();
-        List<ResolvedComponentInfo> personalResolvedComponentInfos =
-                createResolvedComponentsForTestWithOtherProfile(3, /* userId */ 10);
-        List<ResolvedComponentInfo> workResolvedComponentInfos =
-                createResolvedComponentsForTest(3);
-        setupResolverControllers(personalResolvedComponentInfos, workResolvedComponentInfos);
-        Intent sendIntent = createSendTextIntent();
-        sendIntent.setType(TEST_MIME_TYPE);
-        ChooserActivityOverrideData.getInstance().isWorkProfileUserUnlocked = false;
-
-        final ChooserActivity activity =
-                mActivityRule.launchActivity(Intent.createChooser(sendIntent, "work tab test"));
-        final IChooserWrapper wrapper = (IChooserWrapper) activity;
-        waitForIdle();
-        onView(withId(com.android.internal.R.id.contentPanel))
-                .perform(swipeUp());
-        onView(withText(R.string.resolver_work_tab)).perform(click());
-        waitForIdle();
-
-        assertEquals(3, wrapper.getWorkListAdapter().getCount());
+        verify(workProfileShortcutLoader, times(1)).queryShortcuts(any());
     }
 
     private Intent createChooserIntent(Intent intent, Intent[] initialIntents) {
@@ -2713,4 +2703,18 @@ public class UnbundledChooserActivityTest {
                         .getInteger(R.integer.config_chooser_max_targets_per_row))
                 .thenReturn(targetsPerRow);
     }
+
+    private SparseArray<Pair<ShortcutLoader, Consumer<ShortcutLoader.Result>>>
+            createShortcutLoaderFactory() {
+        SparseArray<Pair<ShortcutLoader, Consumer<ShortcutLoader.Result>>> shortcutLoaders =
+                new SparseArray<>();
+        ChooserActivityOverrideData.getInstance().shortcutLoaderFactory =
+                (userHandle, callback) -> {
+                    Pair<ShortcutLoader, Consumer<ShortcutLoader.Result>> pair =
+                            new Pair<>(mock(ShortcutLoader.class), callback);
+                    shortcutLoaders.put(userHandle.getIdentifier(), pair);
+                    return pair.first;
+                };
+        return shortcutLoaders;
+    }
 }
diff --git a/java/tests/src/com/android/intentresolver/shortcuts/ShortcutLoaderTest.kt b/java/tests/src/com/android/intentresolver/shortcuts/ShortcutLoaderTest.kt
new file mode 100644
index 00000000..5756a0cd
--- /dev/null
+++ b/java/tests/src/com/android/intentresolver/shortcuts/ShortcutLoaderTest.kt
@@ -0,0 +1,329 @@
+/*
+ * Copyright (C) 2022 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.intentresolver.shortcuts
+
+import android.app.prediction.AppPredictor
+import android.content.ComponentName
+import android.content.Context
+import android.content.IntentFilter
+import android.content.pm.ApplicationInfo
+import android.content.pm.PackageManager
+import android.content.pm.PackageManager.ApplicationInfoFlags
+import android.content.pm.ShortcutManager
+import android.os.UserHandle
+import android.os.UserManager
+import androidx.test.filters.SmallTest
+import com.android.intentresolver.any
+import com.android.intentresolver.chooser.DisplayResolveInfo
+import com.android.intentresolver.createAppTarget
+import com.android.intentresolver.createShareShortcutInfo
+import com.android.intentresolver.createShortcutInfo
+import com.android.intentresolver.mock
+import com.android.intentresolver.whenever
+import org.junit.Assert.assertArrayEquals
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertFalse
+import org.junit.Assert.assertTrue
+import org.junit.Test
+import org.mockito.ArgumentCaptor
+import org.mockito.Mockito.anyInt
+import org.mockito.Mockito.never
+import org.mockito.Mockito.times
+import org.mockito.Mockito.verify
+import java.util.concurrent.Executor
+import java.util.function.Consumer
+
+@SmallTest
+class ShortcutLoaderTest {
+    private val appInfo = ApplicationInfo().apply {
+        enabled = true
+        flags = 0
+    }
+    private val pm = mock<PackageManager> {
+        whenever(getApplicationInfo(any(), any<ApplicationInfoFlags>())).thenReturn(appInfo)
+    }
+    private val context = mock<Context> {
+        whenever(packageManager).thenReturn(pm)
+        whenever(createContextAsUser(any(), anyInt())).thenReturn(this)
+    }
+    private val executor = ImmediateExecutor()
+    private val intentFilter = mock<IntentFilter>()
+    private val appPredictor = mock<ShortcutLoader.AppPredictorProxy>()
+    private val callback = mock<Consumer<ShortcutLoader.Result>>()
+
+    @Test
+    fun test_app_predictor_result() {
+        val componentName = ComponentName("pkg", "Class")
+        val appTarget = mock<DisplayResolveInfo> {
+            whenever(resolvedComponentName).thenReturn(componentName)
+        }
+        val appTargets = arrayOf(appTarget)
+        val testSubject = ShortcutLoader(
+            context,
+            appPredictor,
+            UserHandle.of(0),
+            true,
+            intentFilter,
+            executor,
+            executor,
+            callback
+        )
+
+        testSubject.queryShortcuts(appTargets)
+
+        verify(appPredictor, times(1)).requestPredictionUpdate()
+        val appPredictorCallbackCaptor = ArgumentCaptor.forClass(AppPredictor.Callback::class.java)
+        verify(appPredictor, times(1))
+            .registerPredictionUpdates(any(), appPredictorCallbackCaptor.capture())
+
+        val matchingShortcutInfo = createShortcutInfo("id-0", componentName, 1)
+        val matchingAppTarget = createAppTarget(matchingShortcutInfo)
+        val shortcuts = listOf(
+            matchingAppTarget,
+            // mismatching shortcut
+            createAppTarget(
+                createShortcutInfo("id-1", ComponentName("mismatching.pkg", "Class"), 1)
+            )
+        )
+        appPredictorCallbackCaptor.value.onTargetsAvailable(shortcuts)
+
+        val resultCaptor = ArgumentCaptor.forClass(ShortcutLoader.Result::class.java)
+        verify(callback, times(1)).accept(resultCaptor.capture())
+
+        val result = resultCaptor.value
+        assertTrue("An app predictor result is expected", result.isFromAppPredictor)
+        assertArrayEquals("Wrong input app targets in the result", appTargets, result.appTargets)
+        assertEquals("Wrong shortcut count", 1, result.shortcutsByApp.size)
+        assertEquals("Wrong app target", appTarget, result.shortcutsByApp[0].appTarget)
+        for (shortcut in result.shortcutsByApp[0].shortcuts) {
+            assertEquals(
+                "Wrong AppTarget in the cache",
+                matchingAppTarget,
+                result.directShareAppTargetCache[shortcut]
+            )
+            assertEquals(
+                "Wrong ShortcutInfo in the cache",
+                matchingShortcutInfo,
+                result.directShareShortcutInfoCache[shortcut]
+            )
+        }
+    }
+
+    @Test
+    fun test_shortcut_manager_result() {
+        val componentName = ComponentName("pkg", "Class")
+        val appTarget = mock<DisplayResolveInfo> {
+            whenever(resolvedComponentName).thenReturn(componentName)
+        }
+        val appTargets = arrayOf(appTarget)
+        val matchingShortcutInfo = createShortcutInfo("id-0", componentName, 1)
+        val shortcutManagerResult = listOf(
+            ShortcutManager.ShareShortcutInfo(matchingShortcutInfo, componentName),
+            // mismatching shortcut
+            createShareShortcutInfo("id-1", ComponentName("mismatching.pkg", "Class"), 1)
+        )
+        val shortcutManager = mock<ShortcutManager> {
+            whenever(getShareTargets(intentFilter)).thenReturn(shortcutManagerResult)
+        }
+        whenever(context.getSystemService(Context.SHORTCUT_SERVICE)).thenReturn(shortcutManager)
+        val testSubject = ShortcutLoader(
+            context,
+            null,
+            UserHandle.of(0),
+            true,
+            intentFilter,
+            executor,
+            executor,
+            callback
+        )
+
+        testSubject.queryShortcuts(appTargets)
+
+        val resultCaptor = ArgumentCaptor.forClass(ShortcutLoader.Result::class.java)
+        verify(callback, times(1)).accept(resultCaptor.capture())
+
+        val result = resultCaptor.value
+        assertFalse("An ShortcutManager result is expected", result.isFromAppPredictor)
+        assertArrayEquals("Wrong input app targets in the result", appTargets, result.appTargets)
+        assertEquals("Wrong shortcut count", 1, result.shortcutsByApp.size)
+        assertEquals("Wrong app target", appTarget, result.shortcutsByApp[0].appTarget)
+        for (shortcut in result.shortcutsByApp[0].shortcuts) {
+            assertTrue(
+                "AppTargets are not expected the cache of a ShortcutManager result",
+                result.directShareAppTargetCache.isEmpty()
+            )
+            assertEquals(
+                "Wrong ShortcutInfo in the cache",
+                matchingShortcutInfo,
+                result.directShareShortcutInfoCache[shortcut]
+            )
+        }
+    }
+
+    @Test
+    fun test_fallback_to_shortcut_manager() {
+        val componentName = ComponentName("pkg", "Class")
+        val appTarget = mock<DisplayResolveInfo> {
+            whenever(resolvedComponentName).thenReturn(componentName)
+        }
+        val appTargets = arrayOf(appTarget)
+        val matchingShortcutInfo = createShortcutInfo("id-0", componentName, 1)
+        val shortcutManagerResult = listOf(
+            ShortcutManager.ShareShortcutInfo(matchingShortcutInfo, componentName),
+            // mismatching shortcut
+            createShareShortcutInfo("id-1", ComponentName("mismatching.pkg", "Class"), 1)
+        )
+        val shortcutManager = mock<ShortcutManager> {
+            whenever(getShareTargets(intentFilter)).thenReturn(shortcutManagerResult)
+        }
+        whenever(context.getSystemService(Context.SHORTCUT_SERVICE)).thenReturn(shortcutManager)
+        val testSubject = ShortcutLoader(
+            context,
+            appPredictor,
+            UserHandle.of(0),
+            true,
+            intentFilter,
+            executor,
+            executor,
+            callback
+        )
+
+        testSubject.queryShortcuts(appTargets)
+
+        verify(appPredictor, times(1)).requestPredictionUpdate()
+        val appPredictorCallbackCaptor = ArgumentCaptor.forClass(AppPredictor.Callback::class.java)
+        verify(appPredictor, times(1))
+            .registerPredictionUpdates(any(), appPredictorCallbackCaptor.capture())
+        appPredictorCallbackCaptor.value.onTargetsAvailable(emptyList())
+
+        val resultCaptor = ArgumentCaptor.forClass(ShortcutLoader.Result::class.java)
+        verify(callback, times(1)).accept(resultCaptor.capture())
+
+        val result = resultCaptor.value
+        assertFalse("An ShortcutManager result is expected", result.isFromAppPredictor)
+        assertArrayEquals("Wrong input app targets in the result", appTargets, result.appTargets)
+        assertEquals("Wrong shortcut count", 1, result.shortcutsByApp.size)
+        assertEquals("Wrong app target", appTarget, result.shortcutsByApp[0].appTarget)
+        for (shortcut in result.shortcutsByApp[0].shortcuts) {
+            assertTrue(
+                "AppTargets are not expected the cache of a ShortcutManager result",
+                result.directShareAppTargetCache.isEmpty()
+            )
+            assertEquals(
+                "Wrong ShortcutInfo in the cache",
+                matchingShortcutInfo,
+                result.directShareShortcutInfoCache[shortcut]
+            )
+        }
+    }
+
+    @Test
+    fun test_do_not_call_services_for_not_running_work_profile() {
+        testDisabledWorkProfileDoNotCallSystem(isUserRunning = false)
+    }
+
+    @Test
+    fun test_do_not_call_services_for_locked_work_profile() {
+        testDisabledWorkProfileDoNotCallSystem(isUserUnlocked = false)
+    }
+
+    @Test
+    fun test_do_not_call_services_if_quite_mode_is_enabled_for_work_profile() {
+        testDisabledWorkProfileDoNotCallSystem(isQuietModeEnabled = true)
+    }
+
+    @Test
+    fun test_call_services_for_not_running_main_profile() {
+        testAlwaysCallSystemForMainProfile(isUserRunning = false)
+    }
+
+    @Test
+    fun test_call_services_for_locked_main_profile() {
+        testAlwaysCallSystemForMainProfile(isUserUnlocked = false)
+    }
+
+    @Test
+    fun test_call_services_if_quite_mode_is_enabled_for_main_profile() {
+        testAlwaysCallSystemForMainProfile(isQuietModeEnabled = true)
+    }
+
+    private fun testDisabledWorkProfileDoNotCallSystem(
+        isUserRunning: Boolean = true,
+        isUserUnlocked: Boolean = true,
+        isQuietModeEnabled: Boolean = false
+    ) {
+        val userHandle = UserHandle.of(10)
+        val userManager = mock<UserManager> {
+            whenever(isUserRunning(userHandle)).thenReturn(isUserRunning)
+            whenever(isUserUnlocked(userHandle)).thenReturn(isUserUnlocked)
+            whenever(isQuietModeEnabled(userHandle)).thenReturn(isQuietModeEnabled)
+        }
+        whenever(context.getSystemService(Context.USER_SERVICE)).thenReturn(userManager);
+        val appPredictor = mock<ShortcutLoader.AppPredictorProxy>()
+        val callback = mock<Consumer<ShortcutLoader.Result>>()
+        val testSubject = ShortcutLoader(
+            context,
+            appPredictor,
+            userHandle,
+            false,
+            intentFilter,
+            executor,
+            executor,
+            callback
+        )
+
+        testSubject.queryShortcuts(arrayOf<DisplayResolveInfo>(mock()))
+
+        verify(appPredictor, never()).requestPredictionUpdate()
+    }
+
+    private fun testAlwaysCallSystemForMainProfile(
+        isUserRunning: Boolean = true,
+        isUserUnlocked: Boolean = true,
+        isQuietModeEnabled: Boolean = false
+    ) {
+        val userHandle = UserHandle.of(10)
+        val userManager = mock<UserManager> {
+            whenever(isUserRunning(userHandle)).thenReturn(isUserRunning)
+            whenever(isUserUnlocked(userHandle)).thenReturn(isUserUnlocked)
+            whenever(isQuietModeEnabled(userHandle)).thenReturn(isQuietModeEnabled)
+        }
+        whenever(context.getSystemService(Context.USER_SERVICE)).thenReturn(userManager);
+        val appPredictor = mock<ShortcutLoader.AppPredictorProxy>()
+        val callback = mock<Consumer<ShortcutLoader.Result>>()
+        val testSubject = ShortcutLoader(
+            context,
+            appPredictor,
+            userHandle,
+            true,
+            intentFilter,
+            executor,
+            executor,
+            callback
+        )
+
+        testSubject.queryShortcuts(arrayOf<DisplayResolveInfo>(mock()))
+
+        verify(appPredictor, times(1)).requestPredictionUpdate()
+    }
+}
+
+private class ImmediateExecutor : Executor {
+    override fun execute(r: Runnable) {
+        r.run()
+    }
+}
diff --git a/java/tests/src/com/android/intentresolver/ShortcutToChooserTargetConverterTest.kt b/java/tests/src/com/android/intentresolver/shortcuts/ShortcutToChooserTargetConverterTest.kt
index 5529e714..e0de005d 100644
--- a/java/tests/src/com/android/intentresolver/ShortcutToChooserTargetConverterTest.kt
+++ b/java/tests/src/com/android/intentresolver/shortcuts/ShortcutToChooserTargetConverterTest.kt
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.intentresolver
+package com.android.intentresolver.shortcuts
 
 import android.app.prediction.AppTarget
 import android.content.ComponentName
@@ -22,6 +22,8 @@ import android.content.Intent
 import android.content.pm.ShortcutInfo
 import android.content.pm.ShortcutManager.ShareShortcutInfo
 import android.service.chooser.ChooserTarget
+import com.android.intentresolver.createAppTarget
+import com.android.intentresolver.createShareShortcutInfo
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertNotNull
 import org.junit.Test