Revert "Extract shortcuts loading logic from ChooserActivity"

This reverts commit 9f9dceaa1398c925b233bcc54fb47cf3531a88da.

Reason for revert: DroidMonitor-triggered revert due to breakage b/259256230.

Change-Id: I399625feb7ee3e6858db2863c8d66e62eccc98c4

7 files changed, 335 insertions, 674 deletions

diff --git a/java/tests/src/com/android/intentresolver/ChooserActivityOverrideData.java b/java/tests/src/com/android/intentresolver/ChooserActivityOverrideData.java
index dd78b69e..e474938b 100644
--- a/java/tests/src/com/android/intentresolver/ChooserActivityOverrideData.java
+++ b/java/tests/src/com/android/intentresolver/ChooserActivityOverrideData.java
@@ -24,17 +24,15 @@ 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.Consumer;
+import java.util.function.BiFunction;
 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
@@ -53,8 +51,10 @@ public class ChooserActivityOverrideData {
     @SuppressWarnings("Since15")
     public Function<PackageManager, PackageManager> createPackageManager;
     public Function<TargetInfo, Boolean> onSafelyStartCallback;
-    public Function2<UserHandle, Consumer<ShortcutLoader.Result>, ShortcutLoader>
-            shortcutLoaderFactory = (userHandle, callback) -> null;
+    public Function<ChooserListAdapter, Void> onQueryDirectShareTargets;
+    public BiFunction<
+            IChooserWrapper, ChooserListAdapter, Pair<Integer, ChooserActivity.ServiceResultInfo[]>>
+            directShareTargets;
     public ResolverListController resolverListController;
     public ResolverListController workResolverListController;
     public Boolean isVoiceInteraction;
@@ -69,11 +69,15 @@ 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;
@@ -89,6 +93,8 @@ public class ChooserActivityOverrideData {
         workProfileUserHandle = null;
         hasCrossProfileIntents = true;
         isQuietModeEnabled = false;
+        isWorkProfileUserRunning = true;
+        isWorkProfileUserUnlocked = true;
         packageManager = null;
         multiPagerAdapterInjector = new AbstractMultiProfilePagerAdapter.Injector() {
             @Override
@@ -108,7 +114,6 @@ 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 6b74fcd4..8c7c28bb 100644
--- a/java/tests/src/com/android/intentresolver/ChooserWrapperActivity.java
+++ b/java/tests/src/com/android/intentresolver/ChooserWrapperActivity.java
@@ -19,13 +19,11 @@ 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;
@@ -33,6 +31,7 @@ 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;
@@ -45,12 +44,11 @@ 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
@@ -258,18 +256,41 @@ public class ChooserWrapperActivity
     }
 
     @Override
-    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;
+    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);
         }
-        return super.createShortcutLoader(
-                context, appPredictor, userHandle, targetIntentFilter, callback);
+        return sOverrides.isWorkProfileUserUnlocked;
     }
 }
diff --git a/java/tests/src/com/android/intentresolver/IChooserWrapper.java b/java/tests/src/com/android/intentresolver/IChooserWrapper.java
index 0d44e147..f81cd023 100644
--- a/java/tests/src/com/android/intentresolver/IChooserWrapper.java
+++ b/java/tests/src/com/android/intentresolver/IChooserWrapper.java
@@ -25,8 +25,6 @@ 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
@@ -43,5 +41,4 @@ public interface IChooserWrapper {
             @Nullable ResolveInfoPresentationGetter resolveInfoPresentationGetter);
     UserHandle getCurrentUserHandle();
     ChooserActivityLogger getChooserActivityLogger();
-    Executor getMainExecutor();
 }
diff --git a/java/tests/src/com/android/intentresolver/shortcuts/ShortcutToChooserTargetConverterTest.kt b/java/tests/src/com/android/intentresolver/ShortcutToChooserTargetConverterTest.kt
index e0de005d..5529e714 100644
--- a/java/tests/src/com/android/intentresolver/shortcuts/ShortcutToChooserTargetConverterTest.kt
+++ b/java/tests/src/com/android/intentresolver/ShortcutToChooserTargetConverterTest.kt
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.intentresolver.shortcuts
+package com.android.intentresolver
 
 import android.app.prediction.AppTarget
 import android.content.ComponentName
@@ -22,8 +22,6 @@ 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
diff --git a/java/tests/src/com/android/intentresolver/TestApplication.kt b/java/tests/src/com/android/intentresolver/TestApplication.kt
deleted file mode 100644
index 849cfbab..00000000
--- a/java/tests/src/com/android/intentresolver/TestApplication.kt
+++ /dev/null
@@ -1,27 +0,0 @@
-/*
- * 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 da72a749..7c304284 100644
--- a/java/tests/src/com/android/intentresolver/UnbundledChooserActivityTest.java
+++ b/java/tests/src/com/android/intentresolver/UnbundledChooserActivityTest.java
@@ -38,6 +38,7 @@ 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;
@@ -82,7 +83,6 @@ 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,7 +118,6 @@ 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;
 
 /**
@@ -1280,7 +1279,7 @@ public class UnbundledChooserActivityTest {
     }
 
     // This test is too long and too slow and should not be taken as an example for future tests.
-    @Test
+    @Test @Ignore
     public void testDirectTargetSelectionLogging() {
         Intent sendIntent = createSendTextIntent();
         // We need app targets for direct targets to get displayed
@@ -1299,55 +1298,37 @@ public class UnbundledChooserActivityTest {
         // Set up resources
         MetricsLogger mockLogger = ChooserActivityOverrideData.getInstance().metricsLogger;
         ArgumentCaptor<LogMaker> logMakerCaptor = ArgumentCaptor.forClass(LogMaker.class);
-
-        // create test shortcut loader factory, remember loaders and their callbacks
-        SparseArray<Pair<ShortcutLoader, Consumer<ShortcutLoader.Result>>> shortcutLoaders =
-                createShortcutLoaderFactory();
+        // Create direct share target
+        List<ChooserTarget> serviceTargets = createDirectShareTargets(1, "");
+        ResolveInfo ri = ResolverDataProvider.createResolveInfo(3, 0);
 
         // 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, "");
-        ShortcutLoader.Result result = new ShortcutLoader.Result(
-                true,
-                appTargets.getValue(),
-                new ShortcutLoader.ShortcutResultInfo[] {
-                        new ShortcutLoader.ShortcutResultInfo(
-                                appTargets.getValue()[0],
-                                serviceTargets
-                        )
-                },
-                new HashMap<>(),
-                new HashMap<>()
+        // 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)
         );
-        activity.getMainExecutor().execute(() -> shortcutLoaders.get(0).second.accept(result));
-        waitForIdle();
 
-        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("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));
 
         // Click on the direct target
         String name = serviceTargets.get(0).getTitle().toString();
@@ -1355,29 +1336,23 @@ public class UnbundledChooserActivityTest {
                 .perform(click());
         waitForIdle();
 
-        // 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(),
+        // 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(),
                 is(MetricsEvent.ACTION_ACTIVITY_CHOOSER_PICKED_SERVICE_TARGET));
-        String hashedName = (String) selectionLog.getTaggedData(
-                MetricsEvent.FIELD_HASHED_TARGET_NAME);
-        assertThat(
-                "Hash is not predictable but must be obfuscated",
+        String hashedName = (String) logMakerCaptor
+                .getAllValues().get(2).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",
-                selectionLog.getTaggedData(MetricsEvent.FIELD_RANKED_POSITION),
-                is(-1));
+        assertThat("The packages shouldn't match for app target and direct target", logMakerCaptor
+                .getAllValues().get(2).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
+    @Test @Ignore
     public void testDirectTargetLoggingWithRankedAppTarget() {
         Intent sendIntent = createSendTextIntent();
         // We need app targets for direct targets to get displayed
@@ -1396,57 +1371,38 @@ public class UnbundledChooserActivityTest {
         // Set up resources
         MetricsLogger mockLogger = ChooserActivityOverrideData.getInstance().metricsLogger;
         ArgumentCaptor<LogMaker> logMakerCaptor = ArgumentCaptor.forClass(LogMaker.class);
-
-        // create test shortcut loader factory, remember loaders and their callbacks
-        SparseArray<Pair<ShortcutLoader, Consumer<ShortcutLoader.Result>>> shortcutLoaders =
-                createShortcutLoaderFactory();
+        // Create direct share target
+        List<ChooserTarget> serviceTargets = createDirectShareTargets(1,
+                resolvedComponentInfos.get(0).getResolveInfoAt(0).activityInfo.packageName);
+        ResolveInfo ri = ResolverDataProvider.createResolveInfo(3, 0);
 
         // 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(
-                true,
-                appTargets.getValue(),
-                new ShortcutLoader.ShortcutResultInfo[] {
-                        new ShortcutLoader.ShortcutResultInfo(
-                                appTargets.getValue()[0],
-                                serviceTargets
-                        )
-                },
-                new HashMap<>(),
-                new HashMap<>()
+        // 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)
         );
-        activity.getMainExecutor().execute(() -> shortcutLoaders.get(0).second.accept(result));
-        waitForIdle();
 
-        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("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));
 
         // Click on the direct target
         String name = serviceTargets.get(0).getTitle().toString();
@@ -1454,19 +1410,18 @@ public class UnbundledChooserActivityTest {
                 .perform(click());
         waitForIdle();
 
-        // 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(),
+        // 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(),
                 is(MetricsEvent.ACTION_ACTIVITY_CHOOSER_PICKED_SERVICE_TARGET));
         assertThat("The packages should match for app target and direct target", logMakerCaptor
-                .getAllValues().get(3).getTaggedData(MetricsEvent.FIELD_RANKED_POSITION), is(0));
+                .getAllValues().get(2).getTaggedData(MetricsEvent.FIELD_RANKED_POSITION), is(0));
     }
 
-    @Test
+    @Test @Ignore
     public void testShortcutTargetWithApplyAppLimits() {
         // Set up resources
         ChooserActivityOverrideData.getInstance().resources = Mockito.spy(
@@ -1490,64 +1445,48 @@ public class UnbundledChooserActivityTest {
                                 Mockito.anyBoolean(),
                                 Mockito.isA(List.class)))
                 .thenReturn(resolvedComponentInfos);
-
-        // create test shortcut loader factory, remember loaders and their callbacks
-        SparseArray<Pair<ShortcutLoader, Consumer<ShortcutLoader.Result>>> shortcutLoaders =
-                createShortcutLoaderFactory();
+        // Create direct share target
+        List<ChooserTarget> serviceTargets = createDirectShareTargets(2,
+                resolvedComponentInfos.get(0).getResolveInfoAt(0).activityInfo.packageName);
+        ResolveInfo ri = ResolverDataProvider.createResolveInfo(3, 0);
 
         // 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());
+        final ChooserActivity activity =
+                mActivityRule.launchActivity(Intent.createChooser(sendIntent, null));
+        final IChooserWrapper wrapper = (IChooserWrapper) activity;
 
-        // 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<>()
+        // 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)
         );
-        activity.getMainExecutor().execute(() -> shortcutLoaders.get(0).second.accept(result));
-        waitForIdle();
 
-        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"));
+        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"));
     }
 
-    @Test
+    @Test @Ignore
     public void testShortcutTargetWithoutApplyAppLimits() {
         setDeviceConfigProperty(
                 SystemUiDeviceConfigFlags.APPLY_SHARING_APP_LIMITS_IN_SYSUI,
@@ -1574,65 +1513,47 @@ public class UnbundledChooserActivityTest {
                                 Mockito.anyBoolean(),
                                 Mockito.isA(List.class)))
                 .thenReturn(resolvedComponentInfos);
-
-        // create test shortcut loader factory, remember loaders and their callbacks
-        SparseArray<Pair<ShortcutLoader, Consumer<ShortcutLoader.Result>>> shortcutLoaders =
-                createShortcutLoaderFactory();
+        // Create direct share target
+        List<ChooserTarget> serviceTargets = createDirectShareTargets(2,
+                resolvedComponentInfos.get(0).getResolveInfoAt(0).activityInfo.packageName);
+        ResolveInfo ri = ResolverDataProvider.createResolveInfo(3, 0);
 
         // Start activity
-        final IChooserWrapper activity = (IChooserWrapper)
+        final ChooserActivity activity =
                 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());
+        final IChooserWrapper wrapper = (IChooserWrapper) activity;
 
-        // 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<>()
+        // 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)
         );
-        activity.getMainExecutor().execute(() -> shortcutLoaders.get(0).second.accept(result));
-        waitForIdle();
 
-        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"));
+        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"));
     }
 
     @Test
@@ -2027,59 +1948,43 @@ public class UnbundledChooserActivityTest {
                                 Mockito.isA(List.class)))
                 .thenReturn(resolvedComponentInfos);
 
-        // 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;
+        // 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);
                 };
 
         // 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(resolvedComponentInfos.get(0).getResolveInfoAt(0)));
+        assertThat("The resolver info must match the resolver info used to create the target",
+                activity.getAdapter().getItem(0).getResolveInfo(), is(ri));
 
         // Click on the direct target
         String name = serviceTargets.get(0).getTitle().toString();
@@ -2193,7 +2098,7 @@ public class UnbundledChooserActivityTest {
             return true;
         };
 
-        mActivityRule.launchActivity(Intent.createChooser(sendIntent, "Test"));
+        mActivityRule.launchActivity(sendIntent);
         waitForIdle();
 
         assertThat(chosen[0], is(personalResolvedComponentInfos.get(1).getResolveInfoAt(0)));
@@ -2368,20 +2273,21 @@ public class UnbundledChooserActivityTest {
     }
 
     @Test
-    public void test_query_shortcut_loader_for_the_selected_tab() {
+    public void testWorkTab_selectingWorkTabWithPausedWorkProfile_directShareTargetsNotQueried() {
         markWorkProfileUserAvailable();
         List<ResolvedComponentInfo> personalResolvedComponentInfos =
                 createResolvedComponentsForTestWithOtherProfile(3, /* userId */ 10);
         List<ResolvedComponentInfo> workResolvedComponentInfos =
                 createResolvedComponentsForTest(3);
         setupResolverControllers(personalResolvedComponentInfos, workResolvedComponentInfos);
-        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);
+        ChooserActivityOverrideData.getInstance().isQuietModeEnabled = true;
+        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);
 
@@ -2389,14 +2295,118 @@ public class UnbundledChooserActivityTest {
         waitForIdle();
         onView(withId(com.android.internal.R.id.contentPanel))
                 .perform(swipeUp());
+        onView(withText(R.string.resolver_work_tab)).perform(click());
         waitForIdle();
 
-        verify(personalProfileShortcutLoader, times(1)).queryShortcuts(any());
+        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);
 
+        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();
 
-        verify(workProfileShortcutLoader, times(1)).queryShortcuts(any());
+        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());
     }
 
     private Intent createChooserIntent(Intent intent, Intent[] initialIntents) {
@@ -2703,18 +2713,4 @@ 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
deleted file mode 100644
index 5756a0cd..00000000
--- a/java/tests/src/com/android/intentresolver/shortcuts/ShortcutLoaderTest.kt
+++ /dev/null
@@ -1,329 +0,0 @@
-/*
- * 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()
-    }
-}