Moves targetIntent logic and WorkProfileAvailability into ActivityLogic

Test: atest com.android.intentresolver.v2
BUG: 302113519
Change-Id: I590c10264220e5d328cad057215e5c4d8b120a4a

9 files changed, 113 insertions, 89 deletions

diff --git a/java/src/com/android/intentresolver/v2/ActivityLogic.kt b/java/src/com/android/intentresolver/v2/ActivityLogic.kt
index e5b89dfa..95f2bed4 100644
--- a/java/src/com/android/intentresolver/v2/ActivityLogic.kt
+++ b/java/src/com/android/intentresolver/v2/ActivityLogic.kt
@@ -13,6 +13,7 @@ import android.util.Log
 import androidx.activity.ComponentActivity
 import com.android.intentresolver.AnnotatedUserHandles
 import com.android.intentresolver.R
+import com.android.intentresolver.WorkProfileAvailabilityManager
 import com.android.intentresolver.icons.TargetDataLoader
 
 /**
@@ -22,12 +23,10 @@ import com.android.intentresolver.icons.TargetDataLoader
  * CommonActivityLogic implementation.
  */
 interface ActivityLogic : CommonActivityLogic {
-    /** The intent for the target. This will always come before [additionalTargets], if any. */
+    /** The intent for the target. This will always come before additional targets, if any. */
     val targetIntent: Intent
     /** Whether the intent is for home. */
     val resolvingHome: Boolean
-    /** Intents for additional targets. These will always come after [targetIntent]. */
-    val additionalTargets: List<Intent>?
     /** Custom title to display. */
     val title: CharSequence?
     /** Resource ID for the title to display when there is no custom title. */
@@ -44,6 +43,8 @@ interface ActivityLogic : CommonActivityLogic {
      * Message showing that intent is forwarded from managed profile to owner or other way around.
      */
     val profileSwitchMessage: String?
+    /** The intents for potential actual targets. [targetIntent] must be first. */
+    val payloadIntents: List<Intent>
 
     /**
      * Called after Activity superclass creation, but before any other onCreate logic is performed.
@@ -71,6 +72,8 @@ interface CommonActivityLogic {
     val devicePolicyManager: DevicePolicyManager
     /** Current [UserHandle]s retrievable by type. */
     val annotatedUserHandles: AnnotatedUserHandles?
+    /** Monitors for changes to work profile availability. */
+    val workProfileAvailabilityManager: WorkProfileAvailabilityManager
 
     /** Returns display message indicating intent forwarding or null if not intent forwarding. */
     fun forwardMessageFor(intent: Intent): String?
@@ -91,6 +94,7 @@ interface CommonActivityLogic {
 class CommonActivityLogicImpl(
     override val tag: String,
     activityProvider: () -> ComponentActivity,
+    onWorkProfileStatusUpdated: () -> Unit,
 ) : CommonActivityLogic {
 
     override val activity: ComponentActivity by lazy { activityProvider() }
@@ -122,6 +126,14 @@ class CommonActivityLogicImpl(
         }
     }
 
+    override val workProfileAvailabilityManager: WorkProfileAvailabilityManager by lazy {
+        WorkProfileAvailabilityManager(
+            userManager,
+            annotatedUserHandles?.workProfileUserHandle,
+            onWorkProfileStatusUpdated,
+        )
+    }
+
     private val forwardToPersonalMessage: String? by lazy {
         devicePolicyManager.resources.getString(FORWARD_INTENT_TO_PERSONAL) {
             activity.context.getString(R.string.forward_intent_to_owner)
diff --git a/java/src/com/android/intentresolver/v2/ChooserActivity.java b/java/src/com/android/intentresolver/v2/ChooserActivity.java
index 446577ef..84258850 100644
--- a/java/src/com/android/intentresolver/v2/ChooserActivity.java
+++ b/java/src/com/android/intentresolver/v2/ChooserActivity.java
@@ -241,6 +241,7 @@ public class ChooserActivity extends Hilt_ChooserActivity implements
         mLogic = new ChooserActivityLogic(
                 TAG,
                 () -> this,
+                this::onWorkProfileStatusUpdated,
                 () -> mTargetDataLoader,
                 this::onPreinitialization
         );
@@ -496,7 +497,7 @@ public class ChooserActivity extends Hilt_ChooserActivity implements
             TargetDataLoader targetDataLoader) {
         ChooserGridAdapter adapter = createChooserGridAdapter(
                 /* context */ this,
-                /* payloadIntents */ mIntents,
+                mLogic.getPayloadIntents(),
                 initialIntents,
                 rList,
                 filterLastUsed,
@@ -521,7 +522,7 @@ public class ChooserActivity extends Hilt_ChooserActivity implements
         int selectedProfile = findSelectedProfile();
         ChooserGridAdapter personalAdapter = createChooserGridAdapter(
                 /* context */ this,
-                /* payloadIntents */ mIntents,
+                mLogic.getPayloadIntents(),
                 selectedProfile == PROFILE_PERSONAL ? initialIntents : null,
                 rList,
                 filterLastUsed,
@@ -529,7 +530,7 @@ public class ChooserActivity extends Hilt_ChooserActivity implements
                 targetDataLoader);
         ChooserGridAdapter workAdapter = createChooserGridAdapter(
                 /* context */ this,
-                /* payloadIntents */ mIntents,
+                mLogic.getPayloadIntents(),
                 selectedProfile == PROFILE_WORK ? initialIntents : null,
                 rList,
                 filterLastUsed,
@@ -540,7 +541,7 @@ public class ChooserActivity extends Hilt_ChooserActivity implements
                 personalAdapter,
                 workAdapter,
                 createEmptyStateProvider(requireAnnotatedUserHandles().workProfileUserHandle),
-                () -> mWorkProfileAvailability.isQuietModeEnabled(),
+                () -> mLogic.getWorkProfileAvailabilityManager().isQuietModeEnabled(),
                 selectedProfile,
                 requireAnnotatedUserHandles().workProfileUserHandle,
                 requireAnnotatedUserHandles().cloneProfileUserHandle,
@@ -1015,7 +1016,7 @@ public class ChooserActivity extends Hilt_ChooserActivity implements
         if (info != null) {
             sendClickToAppPredictor(info);
             final ResolveInfo ri = info.getResolveInfo();
-            Intent targetIntent = getTargetIntent();
+            Intent targetIntent = mLogic.getTargetIntent();
             if (ri != null && ri.activityInfo != null && targetIntent != null) {
                 ChooserListAdapter currentListAdapter =
                         mChooserMultiProfilePagerAdapter.getActiveListAdapter();
@@ -1189,7 +1190,7 @@ public class ChooserActivity extends Hilt_ChooserActivity implements
                 filterLastUsed,
                 createListController(userHandle),
                 userHandle,
-                getTargetIntent(),
+                mLogic.getTargetIntent(),
                 requireChooserRequest(),
                 mMaxTargetsPerRow,
                 targetDataLoader);
@@ -1274,13 +1275,13 @@ public class ChooserActivity extends Hilt_ChooserActivity implements
     }
 
     @Override
-    protected void onWorkProfileStatusUpdated() {
+    protected Unit onWorkProfileStatusUpdated() {
         UserHandle workUser = requireAnnotatedUserHandles().workProfileUserHandle;
         ProfileRecord record = workUser == null ? null : getProfileRecord(workUser);
         if (record != null && record.shortcutLoader != null) {
             record.shortcutLoader.reset();
         }
-        super.onWorkProfileStatusUpdated();
+        return super.onWorkProfileStatusUpdated();
     }
 
     @Override
@@ -1289,14 +1290,20 @@ public class ChooserActivity extends Hilt_ChooserActivity implements
         AppPredictor appPredictor = getAppPredictor(userHandle);
         AbstractResolverComparator resolverComparator;
         if (appPredictor != null) {
-            resolverComparator = new AppPredictionServiceResolverComparator(this, getTargetIntent(),
-                    mLogic.getReferrerPackageName(), appPredictor, userHandle, getEventLog(),
-                    mNearbyShare.orElse(null));
+            resolverComparator = new AppPredictionServiceResolverComparator(
+                    this,
+                    mLogic.getTargetIntent(),
+                    mLogic.getReferrerPackageName(),
+                    appPredictor,
+                    userHandle,
+                    getEventLog(),
+                    mNearbyShare.orElse(null)
+            );
         } else {
             resolverComparator =
                     new ResolverRankerServiceResolverComparator(
                             this,
-                            getTargetIntent(),
+                            mLogic.getTargetIntent(),
                             mLogic.getReferrerPackageName(),
                             null,
                             getEventLog(),
@@ -1307,7 +1314,7 @@ public class ChooserActivity extends Hilt_ChooserActivity implements
         return new ChooserListController(
                 this,
                 mPm,
-                getTargetIntent(),
+                mLogic.getTargetIntent(),
                 mLogic.getReferrerPackageName(),
                 requireAnnotatedUserHandles().userIdOfCallingApp,
                 resolverComparator,
diff --git a/java/src/com/android/intentresolver/v2/ChooserActivityLogic.kt b/java/src/com/android/intentresolver/v2/ChooserActivityLogic.kt
index 838c39e2..5303a7e7 100644
--- a/java/src/com/android/intentresolver/v2/ChooserActivityLogic.kt
+++ b/java/src/com/android/intentresolver/v2/ChooserActivityLogic.kt
@@ -21,18 +21,21 @@ import com.android.intentresolver.v2.util.mutableLazy
 open class ChooserActivityLogic(
     tag: String,
     activityProvider: () -> ComponentActivity,
+    onWorkProfileStatusUpdated: () -> Unit,
     targetDataLoaderProvider: () -> TargetDataLoader,
     private val onPreInitialization: () -> Unit,
-) : ActivityLogic, CommonActivityLogic by CommonActivityLogicImpl(tag, activityProvider) {
+) :
+    ActivityLogic,
+    CommonActivityLogic by CommonActivityLogicImpl(
+        tag,
+        activityProvider,
+        onWorkProfileStatusUpdated,
+    ) {
 
     override val targetIntent: Intent by lazy { chooserRequestParameters?.targetIntent ?: Intent() }
 
     override val resolvingHome: Boolean = false
 
-    override val additionalTargets: List<Intent>? by lazy {
-        chooserRequestParameters?.additionalTargets?.toList()
-    }
-
     override val title: CharSequence? by lazy { chooserRequestParameters?.title }
 
     override val defaultTitleResId: Int by lazy {
@@ -52,6 +55,13 @@ open class ChooserActivityLogic(
     private val _profileSwitchMessage = mutableLazy { forwardMessageFor(targetIntent) }
     override val profileSwitchMessage: String? by _profileSwitchMessage
 
+    override val payloadIntents: List<Intent> by lazy {
+        buildList {
+            add(targetIntent)
+            chooserRequestParameters?.additionalTargets?.let { addAll(it) }
+        }
+    }
+
     val chooserRequestParameters: ChooserRequestParameters? by lazy {
         try {
             ChooserRequestParameters(
diff --git a/java/src/com/android/intentresolver/v2/ResolverActivity.java b/java/src/com/android/intentresolver/v2/ResolverActivity.java
index a97ef962..a7f2047d 100644
--- a/java/src/com/android/intentresolver/v2/ResolverActivity.java
+++ b/java/src/com/android/intentresolver/v2/ResolverActivity.java
@@ -124,9 +124,10 @@ import com.android.internal.logging.MetricsLogger;
 import com.android.internal.logging.nano.MetricsProto;
 import com.android.internal.util.LatencyTracker;
 
+import kotlin.Unit;
+
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
@@ -142,7 +143,11 @@ import java.util.Set;
 public class ResolverActivity extends FragmentActivity implements
         ResolverListAdapter.ResolverListCommunicator {
 
-    protected ActivityLogic mLogic = new ResolverActivityLogic(TAG, () -> this);
+    protected ActivityLogic mLogic = new ResolverActivityLogic(
+            TAG,
+            () -> this,
+            this::onWorkProfileStatusUpdated
+    );
 
     public ResolverActivity() {
         mIsIntentPicker = getClass().equals(ResolverActivity.class);
@@ -157,8 +162,6 @@ public class ResolverActivity extends FragmentActivity implements
     protected View mProfileView;
     private int mLastSelected = AbsListView.INVALID_POSITION;
     private int mLayoutId;
-    @VisibleForTesting
-    protected final ArrayList<Intent> mIntents = new ArrayList<>();
     private PickTargetOptionRequest mPickOptionRequest;
     // Expected to be true if this object is ResolverActivity or is ResolverWrapperActivity.
     private final boolean mIsIntentPicker;
@@ -192,7 +195,6 @@ public class ResolverActivity extends FragmentActivity implements
     @VisibleForTesting
     protected MultiProfilePagerAdapter mMultiProfilePagerAdapter;
 
-    protected WorkProfileAvailabilityManager mWorkProfileAvailability;
 
     // Intent extra for connected audio devices
     public static final String EXTRA_IS_AUDIO_CAPTURE_DEVICE = "is_audio_capture_device";
@@ -318,8 +320,6 @@ public class ResolverActivity extends FragmentActivity implements
         mLogic.preInitialization();
         init(
                 mLogic.getTargetIntent(),
-                mLogic.getAdditionalTargets() == null
-                        ? null : mLogic.getAdditionalTargets().toArray(new Intent[0]),
                 mLogic.getInitialIntents() == null
                         ? null : mLogic.getInitialIntents().toArray(new Intent[0]),
                 mLogic.getTargetDataLoader()
@@ -328,7 +328,6 @@ public class ResolverActivity extends FragmentActivity implements
 
     protected void init(
             Intent intent,
-            Intent[] additionalTargets,
             Intent[] initialIntents,
             TargetDataLoader targetDataLoader
     ) {
@@ -338,16 +337,8 @@ public class ResolverActivity extends FragmentActivity implements
             return;
         }
 
-        mWorkProfileAvailability = createWorkProfileAvailabilityManager();
-
         mPm = getPackageManager();
 
-        // The initial intent must come before any other targets that are to be added.
-        mIntents.add(0, new Intent(intent));
-        if (additionalTargets != null) {
-            Collections.addAll(mIntents, additionalTargets);
-        }
-
         // The last argument of createResolverListAdapter is whether to do special handling
         // of the last used choice to highlight it in the list.  We need to always
         // turn this off when running under voice interaction, since it results in
@@ -582,7 +573,7 @@ public class ResolverActivity extends FragmentActivity implements
             }
         }
         // TODO: should we clean up the work-profile manager before we potentially finish() above?
-        mWorkProfileAvailability.unregisterWorkProfileStateReceiver(this);
+        mLogic.getWorkProfileAvailabilityManager().unregisterWorkProfileStateReceiver(this);
     }
 
     @Override
@@ -857,7 +848,7 @@ public class ResolverActivity extends FragmentActivity implements
         ResolverRankerServiceResolverComparator resolverComparator =
                 new ResolverRankerServiceResolverComparator(
                         this,
-                        getTargetIntent(),
+                        mLogic.getTargetIntent(),
                         mLogic.getReferrerPackageName(),
                         null,
                         null,
@@ -866,7 +857,7 @@ public class ResolverActivity extends FragmentActivity implements
         return new ResolverListController(
                 this,
                 mPm,
-                getTargetIntent(),
+                mLogic.getTargetIntent(),
                 mLogic.getReferrerPackageName(),
                 requireAnnotatedUserHandles().userIdOfCallingApp,
                 resolverComparator,
@@ -958,7 +949,7 @@ public class ResolverActivity extends FragmentActivity implements
         if (listAdapter == mMultiProfilePagerAdapter.getActiveListAdapter()) {
             if (listAdapter.getUserHandle().equals(
                     requireAnnotatedUserHandles().workProfileUserHandle)
-                    && mWorkProfileAvailability.isWaitingToEnableWorkProfile()) {
+                    && mLogic.getWorkProfileAvailabilityManager().isWaitingToEnableWorkProfile()) {
                 // We have just turned on the work profile and entered the pass code to start it,
                 // now we are waiting to receive the ACTION_USER_UNLOCKED broadcast. There is no
                 // point in reloading the list now, since the work profile user is still
@@ -994,20 +985,14 @@ public class ResolverActivity extends FragmentActivity implements
         return new CrossProfileIntentsChecker(getContentResolver());
     }
 
-    protected WorkProfileAvailabilityManager createWorkProfileAvailabilityManager() {
-        return new WorkProfileAvailabilityManager(
-                getSystemService(UserManager.class),
-                requireAnnotatedUserHandles().workProfileUserHandle,
-                this::onWorkProfileStatusUpdated);
-    }
-
-    protected void onWorkProfileStatusUpdated() {
+    protected Unit onWorkProfileStatusUpdated() {
         if (mMultiProfilePagerAdapter.getCurrentUserHandle().equals(
                 requireAnnotatedUserHandles().workProfileUserHandle)) {
             mMultiProfilePagerAdapter.rebuildActiveTab(true);
         } else {
             mMultiProfilePagerAdapter.clearInactiveProfileCache();
         }
+        return Unit.INSTANCE;
     }
 
     // @NonFinalForTesting
@@ -1031,7 +1016,7 @@ public class ResolverActivity extends FragmentActivity implements
                 filterLastUsed,
                 createListController(userHandle),
                 userHandle,
-                getTargetIntent(),
+                mLogic.getTargetIntent(),
                 this,
                 initialIntentsUserSpace,
                 targetDataLoader);
@@ -1059,7 +1044,7 @@ public class ResolverActivity extends FragmentActivity implements
 
         final EmptyStateProvider workProfileOffEmptyStateProvider =
                 new WorkProfilePausedEmptyStateProvider(this, workProfileUserHandle,
-                        mWorkProfileAvailability,
+                        mLogic.getWorkProfileAvailabilityManager(),
                         /* onSwitchOnWorkSelectedListener= */
                         () -> {
                             if (mOnSwitchOnWorkSelectedListener != null) {
@@ -1092,7 +1077,7 @@ public class ResolverActivity extends FragmentActivity implements
                     TargetDataLoader targetDataLoader) {
         ResolverListAdapter adapter = createResolverListAdapter(
                 /* context */ this,
-                /* payloadIntents */ mIntents,
+                mLogic.getPayloadIntents(),
                 initialIntents,
                 resolutionList,
                 filterLastUsed,
@@ -1140,7 +1125,7 @@ public class ResolverActivity extends FragmentActivity implements
         // resolver list. So filterLastUsed should be false for the other profile.
         ResolverListAdapter personalAdapter = createResolverListAdapter(
                 /* context */ this,
-                /* payloadIntents */ mIntents,
+                mLogic.getPayloadIntents(),
                 selectedProfile == PROFILE_PERSONAL ? initialIntents : null,
                 resolutionList,
                 (filterLastUsed && UserHandle.myUserId()
@@ -1150,7 +1135,7 @@ public class ResolverActivity extends FragmentActivity implements
         UserHandle workProfileUserHandle = requireAnnotatedUserHandles().workProfileUserHandle;
         ResolverListAdapter workAdapter = createResolverListAdapter(
                 /* context */ this,
-                /* payloadIntents */ mIntents,
+                mLogic.getPayloadIntents(),
                 selectedProfile == PROFILE_WORK ? initialIntents : null,
                 resolutionList,
                 (filterLastUsed && UserHandle.myUserId()
@@ -1162,7 +1147,7 @@ public class ResolverActivity extends FragmentActivity implements
                 personalAdapter,
                 workAdapter,
                 createEmptyStateProvider(workProfileUserHandle),
-                () -> mWorkProfileAvailability.isQuietModeEnabled(),
+                () -> mLogic.getWorkProfileAvailabilityManager().isQuietModeEnabled(),
                 selectedProfile,
                 workProfileUserHandle,
                 requireAnnotatedUserHandles().cloneProfileUserHandle);
@@ -1286,10 +1271,6 @@ public class ResolverActivity extends FragmentActivity implements
         return new Option(getOrLoadDisplayLabel(target), index);
     }
 
-    public final Intent getTargetIntent() {
-        return mIntents.isEmpty() ? null : mIntents.get(0);
-    }
-
     @Override // ResolverListCommunicator
     public final void updateProfileViewButton() {
         if (mProfileView == null) {
@@ -1360,9 +1341,11 @@ public class ResolverActivity extends FragmentActivity implements
             }
             mRegistered = true;
         }
-        if (shouldShowTabs() && mWorkProfileAvailability.isWaitingToEnableWorkProfile()) {
-            if (mWorkProfileAvailability.isQuietModeEnabled()) {
-                mWorkProfileAvailability.markWorkProfileEnabledBroadcastReceived();
+        WorkProfileAvailabilityManager workProfileAvailabilityManager =
+                mLogic.getWorkProfileAvailabilityManager();
+        if (shouldShowTabs() && workProfileAvailabilityManager.isWaitingToEnableWorkProfile()) {
+            if (workProfileAvailabilityManager.isQuietModeEnabled()) {
+                workProfileAvailabilityManager.markWorkProfileEnabledBroadcastReceived();
             }
         }
         mMultiProfilePagerAdapter.getActiveListAdapter().handlePackagesChanged();
@@ -1375,7 +1358,7 @@ public class ResolverActivity extends FragmentActivity implements
 
         this.getWindow().addSystemFlags(SYSTEM_FLAG_HIDE_NON_SYSTEM_OVERLAY_WINDOWS);
         if (shouldShowTabs()) {
-            mWorkProfileAvailability.registerWorkProfileStateReceiver(this);
+            mLogic.getWorkProfileAvailabilityManager().registerWorkProfileStateReceiver(this);
         }
     }
 
@@ -2062,7 +2045,7 @@ public class ResolverActivity extends FragmentActivity implements
 
         CharSequence title = mLogic.getTitle() != null
                 ? mLogic.getTitle()
-                : getTitleForAction(getTargetIntent(), mLogic.getDefaultTitleResId());
+                : getTitleForAction(mLogic.getTargetIntent(), mLogic.getDefaultTitleResId());
 
         if (!TextUtils.isEmpty(title)) {
             final TextView titleView = findViewById(com.android.internal.R.id.title);
diff --git a/java/src/com/android/intentresolver/v2/ResolverActivityLogic.kt b/java/src/com/android/intentresolver/v2/ResolverActivityLogic.kt
index 1b936159..8cca5b85 100644
--- a/java/src/com/android/intentresolver/v2/ResolverActivityLogic.kt
+++ b/java/src/com/android/intentresolver/v2/ResolverActivityLogic.kt
@@ -2,16 +2,25 @@ package com.android.intentresolver.v2
 
 import android.content.Intent
 import androidx.activity.ComponentActivity
+import androidx.annotation.OpenForTesting
 import com.android.intentresolver.R
 import com.android.intentresolver.icons.DefaultTargetDataLoader
 import com.android.intentresolver.icons.TargetDataLoader
 import com.android.intentresolver.v2.util.mutableLazy
 
 /** Activity logic for [ResolverActivity]. */
-class ResolverActivityLogic(
+@OpenForTesting
+open class ResolverActivityLogic(
     tag: String,
     activityProvider: () -> ComponentActivity,
-) : ActivityLogic, CommonActivityLogic by CommonActivityLogicImpl(tag, activityProvider) {
+    onWorkProfileStatusUpdated: () -> Unit,
+) :
+    ActivityLogic,
+    CommonActivityLogic by CommonActivityLogicImpl(
+        tag,
+        activityProvider,
+        onWorkProfileStatusUpdated,
+    ) {
 
     override val targetIntent: Intent by lazy {
         val intent = Intent(activity.intent)
@@ -37,8 +46,6 @@ class ResolverActivityLogic(
             targetIntent.categories.contains(Intent.CATEGORY_HOME)
     }
 
-    override val additionalTargets: List<Intent>? = null
-
     override val title: CharSequence? = null
 
     override val defaultTitleResId: Int = 0
@@ -63,6 +70,8 @@ class ResolverActivityLogic(
     private val _profileSwitchMessage = mutableLazy { forwardMessageFor(targetIntent) }
     override val profileSwitchMessage: String? by _profileSwitchMessage
 
+    override val payloadIntents: List<Intent> by lazy { listOf(targetIntent) }
+
     override fun preInitialization() {
         // Do nothing
     }
diff --git a/java/tests/src/com/android/intentresolver/v2/ChooserWrapperActivity.java b/java/tests/src/com/android/intentresolver/v2/ChooserWrapperActivity.java
index 6fdba4c2..5572bb24 100644
--- a/java/tests/src/com/android/intentresolver/v2/ChooserWrapperActivity.java
+++ b/java/tests/src/com/android/intentresolver/v2/ChooserWrapperActivity.java
@@ -38,7 +38,6 @@ import com.android.intentresolver.ChooserRequestParameters;
 import com.android.intentresolver.IChooserWrapper;
 import com.android.intentresolver.ResolverListController;
 import com.android.intentresolver.TestContentPreviewViewModel;
-import com.android.intentresolver.WorkProfileAvailabilityManager;
 import com.android.intentresolver.chooser.DisplayResolveInfo;
 import com.android.intentresolver.chooser.TargetInfo;
 import com.android.intentresolver.emptystate.CrossProfileIntentsChecker;
@@ -63,8 +62,9 @@ public class ChooserWrapperActivity extends ChooserActivity implements IChooserW
         mLogic = new TestChooserActivityLogic(
                 "ChooserWrapper",
                 () -> this,
+                this::onWorkProfileStatusUpdated,
                 () -> mTargetDataLoader,
-                super::onPreinitialization,
+                this::onPreinitialization,
                 sOverrides
         );
     }
@@ -160,14 +160,6 @@ public class ChooserWrapperActivity extends ChooserActivity implements IChooserW
     }
 
     @Override
-    protected WorkProfileAvailabilityManager createWorkProfileAvailabilityManager() {
-        if (sOverrides.mWorkProfileAvailability != null) {
-            return sOverrides.mWorkProfileAvailability;
-        }
-        return super.createWorkProfileAvailabilityManager();
-    }
-
-    @Override
     public void safelyStartActivityInternal(TargetInfo cti, UserHandle user,
             @Nullable Bundle options) {
         if (sOverrides.onSafelyStartInternalCallback != null
diff --git a/java/tests/src/com/android/intentresolver/v2/ResolverWrapperActivity.java b/java/tests/src/com/android/intentresolver/v2/ResolverWrapperActivity.java
index e5617090..92b73d92 100644
--- a/java/tests/src/com/android/intentresolver/v2/ResolverWrapperActivity.java
+++ b/java/tests/src/com/android/intentresolver/v2/ResolverWrapperActivity.java
@@ -45,6 +45,8 @@ import com.android.intentresolver.emptystate.CrossProfileIntentsChecker;
 import com.android.intentresolver.icons.LabelInfo;
 import com.android.intentresolver.icons.TargetDataLoader;
 
+import kotlin.Unit;
+
 import java.util.List;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -63,6 +65,10 @@ public class ResolverWrapperActivity extends ResolverActivity {
         mLogic = new TestResolverActivityLogic(
                 "ResolverWrapper",
                 () -> this,
+                () -> {
+                    onWorkProfileStatusUpdated();
+                    return Unit.INSTANCE;
+                },
                 sOverrides
         );
     }
@@ -102,14 +108,6 @@ public class ResolverWrapperActivity extends ResolverActivity {
         return super.createCrossProfileIntentsChecker();
     }
 
-    @Override
-    protected WorkProfileAvailabilityManager createWorkProfileAvailabilityManager() {
-        if (sOverrides.mWorkProfileAvailability != null) {
-            return sOverrides.mWorkProfileAvailability;
-        }
-        return super.createWorkProfileAvailabilityManager();
-    }
-
     ResolverListAdapter getAdapter() {
         return mMultiProfilePagerAdapter.getActiveListAdapter();
     }
diff --git a/java/tests/src/com/android/intentresolver/v2/TestChooserActivityLogic.kt b/java/tests/src/com/android/intentresolver/v2/TestChooserActivityLogic.kt
index fb1eab6c..198b9236 100644
--- a/java/tests/src/com/android/intentresolver/v2/TestChooserActivityLogic.kt
+++ b/java/tests/src/com/android/intentresolver/v2/TestChooserActivityLogic.kt
@@ -2,19 +2,22 @@ package com.android.intentresolver.v2
 
 import androidx.activity.ComponentActivity
 import com.android.intentresolver.AnnotatedUserHandles
+import com.android.intentresolver.WorkProfileAvailabilityManager
 import com.android.intentresolver.icons.TargetDataLoader
 
 /** Activity logic for use when testing [ChooserActivity]. */
 class TestChooserActivityLogic(
     tag: String,
     activityProvider: () -> ComponentActivity,
+    onWorkProfileStatusUpdated: () -> Unit,
     targetDataLoaderProvider: () -> TargetDataLoader,
     onPreinitialization: () -> Unit,
-    overrideData: ChooserActivityOverrideData,
+    private val overrideData: ChooserActivityOverrideData,
 ) :
     ChooserActivityLogic(
         tag,
         activityProvider,
+        onWorkProfileStatusUpdated,
         targetDataLoaderProvider,
         onPreinitialization,
     ) {
@@ -22,4 +25,8 @@ class TestChooserActivityLogic(
     override val annotatedUserHandles: AnnotatedUserHandles? by lazy {
         overrideData.annotatedUserHandles
     }
+
+    override val workProfileAvailabilityManager: WorkProfileAvailabilityManager by lazy {
+        overrideData.mWorkProfileAvailability ?: super.workProfileAvailabilityManager
+    }
 }
diff --git a/java/tests/src/com/android/intentresolver/v2/TestResolverActivityLogic.kt b/java/tests/src/com/android/intentresolver/v2/TestResolverActivityLogic.kt
index 7f8e6f70..7581043e 100644
--- a/java/tests/src/com/android/intentresolver/v2/TestResolverActivityLogic.kt
+++ b/java/tests/src/com/android/intentresolver/v2/TestResolverActivityLogic.kt
@@ -2,15 +2,21 @@ package com.android.intentresolver.v2
 
 import androidx.activity.ComponentActivity
 import com.android.intentresolver.AnnotatedUserHandles
+import com.android.intentresolver.WorkProfileAvailabilityManager
 
 /** Activity logic for use when testing [ResolverActivity]. */
 class TestResolverActivityLogic(
     tag: String,
     activityProvider: () -> ComponentActivity,
-    overrideData: ResolverWrapperActivity.OverrideData,
-) : ActivityLogic by ResolverActivityLogic(tag, activityProvider) {
+    onWorkProfileStatusUpdated: () -> Unit,
+    private val overrideData: ResolverWrapperActivity.OverrideData,
+) : ResolverActivityLogic(tag, activityProvider, onWorkProfileStatusUpdated) {
 
     override val annotatedUserHandles: AnnotatedUserHandles? by lazy {
         overrideData.annotatedUserHandles
     }
+
+    override val workProfileAvailabilityManager: WorkProfileAvailabilityManager by lazy {
+        overrideData.mWorkProfileAvailability ?: super.workProfileAvailabilityManager
+    }
 }