Use a ViewState to preserve refinement state

Refinement manager keeps its own state across config changes so that it
can receive callbacks on the ResultReceiver. ChooserActivity can watch
for the state of refinement completion and launch the refined TargetInfo
if refinement succeeds, otherwise just finish().

If onResume happens when refinement is in progress, ChooserActivity will
finish, as this is expected to be abandonment of the refinement process
without a result. But, onResume also happens during a config change
during refinement, and the activity must not finish in those cases. The
logic for this is contained within ChooserRefinementManager, with
ChooserActivity only needing to pass along a couple of lifecycle events.

Add a bunch of test coverage to new and old functionality.

Also fix some kotlin nullability errors in TargetInfo tests that sysui
studio was complaining about.

Bug: 279514914
Test: atest ChooserRefinementManagerTest
Test: manually test the refinement flow in Photos, both completing and
      cancelling the process, rotating at various points.
Change-Id: I2b22155894e29b062c78b4af20e8fe0683d40bea

5 files changed, 329 insertions, 103 deletions

diff --git a/java/src/com/android/intentresolver/ChooserActivity.java b/java/src/com/android/intentresolver/ChooserActivity.java
index 7f55f78f..8067876f 100644
--- a/java/src/com/android/intentresolver/ChooserActivity.java
+++ b/java/src/com/android/intentresolver/ChooserActivity.java
@@ -72,6 +72,7 @@ import android.view.animation.LinearInterpolator;
 import android.widget.TextView;
 
 import androidx.annotation.MainThread;
+import androidx.lifecycle.ViewModelProvider;
 import androidx.recyclerview.widget.GridLayoutManager;
 import androidx.recyclerview.widget.RecyclerView;
 import androidx.viewpager.widget.ViewPager;
@@ -250,20 +251,25 @@ public class ChooserActivity extends ResolverActivity implements
             return;
         }
 
-        mRefinementManager = new ChooserRefinementManager(
-                this,
-                mChooserRequest.getRefinementIntentSender(),
-                (validatedRefinedTarget) -> {
-                    maybeRemoveSharedText(validatedRefinedTarget);
+        mRefinementManager = new ViewModelProvider(this).get(ChooserRefinementManager.class);
+
+        mRefinementManager.getRefinementCompletion().observe(this, completion -> {
+            if (completion.consume()) {
+                TargetInfo targetInfo = completion.getTargetInfo();
+                // targetInfo is non-null if the refinement process was successful.
+                if (targetInfo != null) {
+                    maybeRemoveSharedText(targetInfo);
 
                     // We already block suspended targets from going to refinement, and we probably
                     // can't recover a Chooser session if that's the reason the refined target fails
                     // to launch now. Fire-and-forget the refined launch; ignore the return value
                     // and just make sure the Sharesheet session gets cleaned up regardless.
-                    super.onTargetSelected(validatedRefinedTarget, false);
-                    finish();
-                },
-                this::finish);
+                    ChooserActivity.super.onTargetSelected(targetInfo, false);
+                }
+
+                finish();
+            }
+        });
 
         mChooserContentPreviewUi = new ChooserContentPreviewUi(
                 mChooserRequest.getTargetIntent(),
@@ -611,14 +617,7 @@ public class ChooserActivity extends ResolverActivity implements
         Log.d(TAG, "onResume: " + getComponentName().flattenToShortString());
         maybeCancelFinishAnimation();
 
-        if (mRefinementManager.isAwaitingRefinementResult()) {
-            // This can happen if the refinement activity terminates without ever sending a response
-            // to our `ResultReceiver`. We're probably not prepared to return the user into a valid
-            // Chooser session, so we'll treat it as a cancellation instead.
-            Log.w(TAG, "Chooser resumed while awaiting refinement result; aborting");
-            mRefinementManager.destroy();
-            finish();
-        }
+        mRefinementManager.onActivityResume();
     }
 
     @Override
@@ -730,6 +729,8 @@ public class ChooserActivity extends ResolverActivity implements
     @Override
     protected void onStop() {
         super.onStop();
+        mRefinementManager.onActivityStop(isChangingConfigurations());
+
         if (maybeCancelFinishAnimation()) {
             finish();
         }
@@ -743,11 +744,6 @@ public class ChooserActivity extends ResolverActivity implements
             mLatencyTracker.onActionCancel(ACTION_LOAD_SHARE_SHEET);
         }
 
-        if (mRefinementManager != null) {  // TODO: null-checked in case of early-destroy, or skip?
-            mRefinementManager.destroy();
-            mRefinementManager = null;
-        }
-
         mBackgroundThreadPoolExecutor.shutdownNow();
 
         destroyProfileRecords();
@@ -873,7 +869,11 @@ public class ChooserActivity extends ResolverActivity implements
 
     @Override
     protected boolean onTargetSelected(TargetInfo target, boolean alwaysCheck) {
-        if (mRefinementManager.maybeHandleSelection(target)) {
+        if (mRefinementManager.maybeHandleSelection(
+                target,
+                mChooserRequest.getRefinementIntentSender(),
+                getApplication(),
+                getMainThreadHandler())) {
             return false;
         }
         updateModelAndChooserCounts(target);
diff --git a/java/src/com/android/intentresolver/ChooserRefinementManager.java b/java/src/com/android/intentresolver/ChooserRefinementManager.java
index 8d7b1aac..2ebe48a6 100644
--- a/java/src/com/android/intentresolver/ChooserRefinementManager.java
+++ b/java/src/com/android/intentresolver/ChooserRefinementManager.java
@@ -19,16 +19,19 @@ package com.android.intentresolver;
 import android.annotation.Nullable;
 import android.annotation.UiThread;
 import android.app.Activity;
-import android.content.Context;
+import android.app.Application;
 import android.content.Intent;
 import android.content.IntentSender;
-import android.content.IntentSender.SendIntentException;
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.Parcel;
 import android.os.ResultReceiver;
 import android.util.Log;
 
+import androidx.lifecycle.LiveData;
+import androidx.lifecycle.MutableLiveData;
+import androidx.lifecycle.ViewModel;
+
 import com.android.intentresolver.chooser.TargetInfo;
 
 import java.util.List;
@@ -42,38 +45,51 @@ import java.util.function.Consumer;
  * call).
  */
 @UiThread
-public final class ChooserRefinementManager {
+public final class ChooserRefinementManager extends ViewModel {
     private static final String TAG = "ChooserRefinement";
 
-    @Nullable
-    private final IntentSender mRefinementIntentSender;
-
-    private final Context mContext;
-    private final Consumer<TargetInfo> mOnSelectionRefined;
-    private final Runnable mOnRefinementCancelled;
-
     @Nullable    // Non-null only during an active refinement session.
     private RefinementResultReceiver mRefinementResultReceiver;
 
-    public ChooserRefinementManager(
-            Context context,
-            @Nullable IntentSender refinementIntentSender,
-            Consumer<TargetInfo> onSelectionRefined,
-            Runnable onRefinementCancelled) {
-        mContext = context;
-        mRefinementIntentSender = refinementIntentSender;
-        mOnSelectionRefined = onSelectionRefined;
-        mOnRefinementCancelled = onRefinementCancelled;
-    }
+    private boolean mConfigurationChangeInProgress = false;
 
     /**
-     * @return whether a refinement session has been initiated (i.e., an earlier call to
-     * {@link #maybeHandleSelection(TargetInfo)} returned true), and isn't yet complete. The session
-     * is complete if the refinement activity calls {@link ResultReceiver#onResultReceived()} (with
-     * any result), or if it's cancelled on our side by {@link ChooserRefinementManager#destroy()}.
+     * A token for the completion of a refinement process that can be consumed exactly once.
      */
-    public boolean isAwaitingRefinementResult() {
-        return (mRefinementResultReceiver != null);
+    public static class RefinementCompletion {
+        private TargetInfo mTargetInfo;
+        private boolean mConsumed;
+
+        RefinementCompletion(TargetInfo targetInfo) {
+            mTargetInfo = targetInfo;
+        }
+
+        /**
+         * @return The output of the completed refinement process. Null if the process was aborted
+         *         or failed.
+         */
+        public TargetInfo getTargetInfo() {
+            return mTargetInfo;
+        }
+
+        /**
+         * Mark this event as consumed if it wasn't already.
+         *
+         * @return true if this had not already been consumed.
+         */
+        public boolean consume() {
+            if (!mConsumed) {
+                mConsumed = true;
+                return true;
+            }
+            return false;
+        }
+    }
+
+    private MutableLiveData<RefinementCompletion> mRefinementCompletion = new MutableLiveData<>();
+
+    public LiveData<RefinementCompletion> getRefinementCompletion() {
+        return mRefinementCompletion;
     }
 
     /**
@@ -81,8 +97,9 @@ public final class ChooserRefinementManager {
      * @return true if the selection should wait for a now-started refinement flow, or false if it
      * can proceed by the default (non-refinement) logic.
      */
-    public boolean maybeHandleSelection(TargetInfo selectedTarget) {
-        if (mRefinementIntentSender == null) {
+    public boolean maybeHandleSelection(TargetInfo selectedTarget,
+            IntentSender refinementIntentSender, Application application, Handler mainHandler) {
+        if (refinementIntentSender == null) {
             return false;
         }
         if (selectedTarget.getAllSourceIntents().isEmpty()) {
@@ -100,33 +117,61 @@ public final class ChooserRefinementManager {
         mRefinementResultReceiver = new RefinementResultReceiver(
                 refinedIntent -> {
                     destroy();
+
                     TargetInfo refinedTarget =
                             selectedTarget.tryToCloneWithAppliedRefinement(refinedIntent);
                     if (refinedTarget != null) {
-                        mOnSelectionRefined.accept(refinedTarget);
+                        mRefinementCompletion.setValue(new RefinementCompletion(refinedTarget));
                     } else {
                         Log.e(TAG, "Failed to apply refinement to any matching source intent");
-                        mOnRefinementCancelled.run();
+                        mRefinementCompletion.setValue(new RefinementCompletion(null));
                     }
                 },
                 () -> {
                     destroy();
-                    mOnRefinementCancelled.run();
+                    mRefinementCompletion.setValue(new RefinementCompletion(null));
                 },
-                mContext.getMainThreadHandler());
+                mainHandler);
 
         Intent refinementRequest = makeRefinementRequest(mRefinementResultReceiver, selectedTarget);
         try {
-            mRefinementIntentSender.sendIntent(mContext, 0, refinementRequest, null, null);
+            refinementIntentSender.sendIntent(application, 0, refinementRequest, null, null);
             return true;
-        } catch (SendIntentException e) {
+        } catch (IntentSender.SendIntentException e) {
             Log.e(TAG, "Refinement IntentSender failed to send", e);
         }
-        return false;
+        return true;
+    }
+
+    /** ChooserActivity has stopped */
+    public void onActivityStop(boolean configurationChanging) {
+        mConfigurationChangeInProgress = configurationChanging;
+    }
+
+    /** ChooserActivity has resumed */
+    public void onActivityResume() {
+        if (mConfigurationChangeInProgress) {
+            mConfigurationChangeInProgress = false;
+        } else {
+            if (mRefinementResultReceiver != null) {
+                // This can happen if the refinement activity terminates without ever sending a
+                // response to our `ResultReceiver`. We're probably not prepared to return the user
+                // into a valid Chooser session, so we'll treat it as a cancellation instead.
+                Log.w(TAG, "Chooser resumed while awaiting refinement result; aborting");
+                destroy();
+                mRefinementCompletion.setValue(new RefinementCompletion(null));
+            }
+        }
+    }
+
+    @Override
+    protected void onCleared() {
+        // App lifecycle over, time to clean up.
+        destroy();
     }
 
     /** Clean up any ongoing refinement session. */
-    public void destroy() {
+    private void destroy() {
         if (mRefinementResultReceiver != null) {
             mRefinementResultReceiver.destroyReceiver();
             mRefinementResultReceiver = null;
diff --git a/java/tests/src/com/android/intentresolver/ChooserRefinementManagerTest.kt b/java/tests/src/com/android/intentresolver/ChooserRefinementManagerTest.kt
index 50c37c7f..bd355c86 100644
--- a/java/tests/src/com/android/intentresolver/ChooserRefinementManagerTest.kt
+++ b/java/tests/src/com/android/intentresolver/ChooserRefinementManagerTest.kt
@@ -16,46 +16,227 @@
 
 package com.android.intentresolver
 
-import android.content.Context
+import android.app.Activity
+import android.app.Application
 import android.content.Intent
 import android.content.IntentSender
+import android.os.Bundle
+import android.os.Handler
+import android.os.Looper
+import android.os.Message
+import android.os.ResultReceiver
+import androidx.lifecycle.Observer
+import androidx.test.annotation.UiThreadTest
 import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.intentresolver.ChooserRefinementManager.RefinementCompletion
+import com.android.intentresolver.chooser.ImmutableTargetInfo
 import com.android.intentresolver.chooser.TargetInfo
+import com.google.common.truth.Truth.assertThat
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
+import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.ArgumentCaptor
 import org.mockito.Mockito
-import java.util.function.Consumer
-import org.junit.Assert.assertEquals
 
 @RunWith(AndroidJUnit4::class)
+@UiThreadTest
 class ChooserRefinementManagerTest {
-    @Test
-    fun testMaybeHandleSelection() {
-        val intentSender = mock<IntentSender>()
-        val refinementManager = ChooserRefinementManager(
-            mock<Context>(),
-            intentSender,
-            Consumer<TargetInfo>{},
-            Runnable{})
-
-        val intents = listOf(Intent(Intent.ACTION_VIEW), Intent(Intent.ACTION_EDIT))
-        val targetInfo = mock<TargetInfo>{
-            whenever(allSourceIntents).thenReturn(intents)
+    private val refinementManager = ChooserRefinementManager()
+    private val intentSender = mock<IntentSender>()
+    private val application = mock<Application>()
+    private val exampleSourceIntents =
+        listOf(Intent(Intent.ACTION_VIEW), Intent(Intent.ACTION_EDIT))
+    private val exampleTargetInfo =
+        ImmutableTargetInfo.newBuilder().setAllSourceIntents(exampleSourceIntents).build()
+
+    private val completionObserver =
+        object : Observer<RefinementCompletion> {
+            val failureCountDown = CountDownLatch(1)
+            val successCountDown = CountDownLatch(1)
+            var latestTargetInfo: TargetInfo? = null
+
+            override fun onChanged(completion: RefinementCompletion) {
+                if (completion.consume()) {
+                    val targetInfo = completion.targetInfo
+                    if (targetInfo == null) {
+                        failureCountDown.countDown()
+                    } else {
+                        latestTargetInfo = targetInfo
+                        successCountDown.countDown()
+                    }
+                }
+            }
         }
 
-        refinementManager.maybeHandleSelection(targetInfo)
+    /** Synchronously executes post() calls. */
+    private class FakeHandler(looper: Looper) : Handler(looper) {
+        override fun sendMessageAtTime(msg: Message, uptimeMillis: Long): Boolean {
+            dispatchMessage(msg)
+            return true
+        }
+    }
+
+    @Before
+    fun setup() {
+        refinementManager.refinementCompletion.observeForever(completionObserver)
+    }
+
+    @Test
+    fun testTypicalRefinementFlow() {
+        assertThat(
+                refinementManager.maybeHandleSelection(
+                    exampleTargetInfo,
+                    intentSender,
+                    application,
+                    FakeHandler(Looper.myLooper())
+                )
+            )
+            .isTrue()
 
         val intentCaptor = ArgumentCaptor.forClass(Intent::class.java)
-        Mockito.verify(intentSender).sendIntent(
-            any(), eq(0), intentCaptor.capture(), eq(null), eq(null))
+        Mockito.verify(intentSender)
+            .sendIntent(any(), eq(0), intentCaptor.capture(), eq(null), eq(null))
 
         val intent = intentCaptor.value
-        assertEquals(intents[0], intent.getParcelableExtra(Intent.EXTRA_INTENT, Intent::class.java))
+        assertThat(intent?.getParcelableExtra(Intent.EXTRA_INTENT, Intent::class.java))
+            .isEqualTo(exampleSourceIntents[0])
 
         val alternates =
-            intent.getParcelableArrayExtra(Intent.EXTRA_ALTERNATE_INTENTS, Intent::class.java)
-        assertEquals(1, alternates?.size)
-        assertEquals(intents[1], alternates?.get(0))
+            intent?.getParcelableArrayExtra(Intent.EXTRA_ALTERNATE_INTENTS, Intent::class.java)
+        assertThat(alternates?.size).isEqualTo(1)
+        assertThat(alternates?.get(0)).isEqualTo(exampleSourceIntents[1])
+
+        // Complete the refinement
+        val receiver =
+            intent?.getParcelableExtra(Intent.EXTRA_RESULT_RECEIVER, ResultReceiver::class.java)
+        val bundle = Bundle().apply { putParcelable(Intent.EXTRA_INTENT, exampleSourceIntents[0]) }
+        receiver?.send(Activity.RESULT_OK, bundle)
+
+        assertThat(completionObserver.successCountDown.await(1000, TimeUnit.MILLISECONDS)).isTrue()
+        assertThat(completionObserver.latestTargetInfo?.resolvedIntent?.action)
+            .isEqualTo(Intent.ACTION_VIEW)
+    }
+
+    @Test
+    fun testRefinementCancelled() {
+        assertThat(
+                refinementManager.maybeHandleSelection(
+                    exampleTargetInfo,
+                    intentSender,
+                    application,
+                    FakeHandler(Looper.myLooper())
+                )
+            )
+            .isTrue()
+
+        val intentCaptor = ArgumentCaptor.forClass(Intent::class.java)
+        Mockito.verify(intentSender)
+            .sendIntent(any(), eq(0), intentCaptor.capture(), eq(null), eq(null))
+
+        val intent = intentCaptor.value
+
+        // Complete the refinement
+        val receiver =
+            intent?.getParcelableExtra(Intent.EXTRA_RESULT_RECEIVER, ResultReceiver::class.java)
+        val bundle = Bundle().apply { putParcelable(Intent.EXTRA_INTENT, exampleSourceIntents[0]) }
+        receiver?.send(Activity.RESULT_CANCELED, bundle)
+
+        assertThat(completionObserver.failureCountDown.await(1000, TimeUnit.MILLISECONDS)).isTrue()
+    }
+
+    @Test
+    fun testMaybeHandleSelection_noSourceIntents() {
+        assertThat(
+                refinementManager.maybeHandleSelection(
+                    ImmutableTargetInfo.newBuilder().build(),
+                    intentSender,
+                    application,
+                    FakeHandler(Looper.myLooper())
+                )
+            )
+            .isFalse()
+    }
+
+    @Test
+    fun testMaybeHandleSelection_suspended() {
+        val targetInfo =
+            ImmutableTargetInfo.newBuilder()
+                .setAllSourceIntents(exampleSourceIntents)
+                .setIsSuspended(true)
+                .build()
+
+        assertThat(
+                refinementManager.maybeHandleSelection(
+                    targetInfo,
+                    intentSender,
+                    application,
+                    FakeHandler(Looper.myLooper())
+                )
+            )
+            .isFalse()
+    }
+
+    @Test
+    fun testMaybeHandleSelection_noIntentSender() {
+        assertThat(
+                refinementManager.maybeHandleSelection(
+                    exampleTargetInfo,
+                    /* IntentSender */ null,
+                    application,
+                    FakeHandler(Looper.myLooper())
+                )
+            )
+            .isFalse()
+    }
+
+    @Test
+    fun testConfigurationChangeDuringRefinement() {
+        assertThat(
+                refinementManager.maybeHandleSelection(
+                    exampleTargetInfo,
+                    intentSender,
+                    application,
+                    FakeHandler(Looper.myLooper())
+                )
+            )
+            .isTrue()
+
+        refinementManager.onActivityStop(/* config changing = */ true)
+        refinementManager.onActivityResume()
+
+        assertThat(completionObserver.failureCountDown.count).isEqualTo(1)
+    }
+
+    @Test
+    fun testResumeDuringRefinement() {
+        assertThat(
+                refinementManager.maybeHandleSelection(
+                    exampleTargetInfo,
+                    intentSender,
+                    application,
+                    FakeHandler(Looper.myLooper()!!)
+                )
+            )
+            .isTrue()
+
+        refinementManager.onActivityStop(/* config changing = */ false)
+        // Resume during refinement but not during a config change, so finish the activity.
+        refinementManager.onActivityResume()
+
+        // Call should be synchronous, don't need to await for this one.
+        assertThat(completionObserver.failureCountDown.count).isEqualTo(0)
+    }
+
+    @Test
+    fun testRefinementCompletion() {
+        val refinementCompletion = RefinementCompletion(exampleTargetInfo)
+        assertThat(refinementCompletion.targetInfo).isEqualTo(exampleTargetInfo)
+        assertThat(refinementCompletion.consume()).isTrue()
+        assertThat(refinementCompletion.targetInfo).isEqualTo(exampleTargetInfo)
+
+        // can only consume once.
+        assertThat(refinementCompletion.consume()).isFalse()
     }
 }
diff --git a/java/tests/src/com/android/intentresolver/chooser/ImmutableTargetInfoTest.kt b/java/tests/src/com/android/intentresolver/chooser/ImmutableTargetInfoTest.kt
index cebccaae..504cfd97 100644
--- a/java/tests/src/com/android/intentresolver/chooser/ImmutableTargetInfoTest.kt
+++ b/java/tests/src/com/android/intentresolver/chooser/ImmutableTargetInfoTest.kt
@@ -222,8 +222,8 @@ class ImmutableTargetInfoTest {
             .build()
         val info = originalInfo.tryToCloneWithAppliedRefinement(refinementIntent)
 
-        assertThat(info.baseIntentToSend.getBooleanExtra("ORIGINAL", false)).isTrue()
-        assertThat(info.baseIntentToSend.getBooleanExtra("REFINEMENT", false)).isTrue()
+        assertThat(info?.baseIntentToSend?.getBooleanExtra("ORIGINAL", false)).isTrue()
+        assertThat(info?.baseIntentToSend?.getBooleanExtra("REFINEMENT", false)).isTrue()
     }
 
     @Test
@@ -245,9 +245,9 @@ class ImmutableTargetInfoTest {
 
         val info = infoWithReferrerFillIn.tryToCloneWithAppliedRefinement(refinementIntent)
 
-        assertThat(info.baseIntentToSend.getPackage()).isEqualTo("original")  // Set all along.
-        assertThat(info.baseIntentToSend.action).isEqualTo("REFINE_ME")  // Refinement wins.
-        assertThat(info.baseIntentToSend.type).isEqualTo("test/referrer")  // Left for referrer.
+        assertThat(info?.baseIntentToSend?.getPackage()).isEqualTo("original")  // Set all along.
+        assertThat(info?.baseIntentToSend?.action).isEqualTo("REFINE_ME")  // Refinement wins.
+        assertThat(info?.baseIntentToSend?.type).isEqualTo("test/referrer")  // Left for referrer.
     }
 
     @Test
@@ -266,18 +266,18 @@ class ImmutableTargetInfoTest {
             .build()
 
         val refined1 = originalInfo.tryToCloneWithAppliedRefinement(refinementIntent1)
-        val refined2 = refined1.tryToCloneWithAppliedRefinement(refinementIntent2)  // Cloned clone.
+        val refined2 = refined1?.tryToCloneWithAppliedRefinement(refinementIntent2)  // Cloned clone.
 
         // Both clones get the same values filled in from the referrer intent.
-        assertThat(refined1.baseIntentToSend.getStringExtra("TEST")).isEqualTo("REFERRER")
-        assertThat(refined2.baseIntentToSend.getStringExtra("TEST")).isEqualTo("REFERRER")
+        assertThat(refined1?.baseIntentToSend?.getStringExtra("TEST")).isEqualTo("REFERRER")
+        assertThat(refined2?.baseIntentToSend?.getStringExtra("TEST")).isEqualTo("REFERRER")
         // Each clone has the respective value that was set in their own refinement request.
-        assertThat(refined1.baseIntentToSend.getStringExtra("TEST1")).isEqualTo("1")
-        assertThat(refined2.baseIntentToSend.getStringExtra("TEST2")).isEqualTo("2")
+        assertThat(refined1?.baseIntentToSend?.getStringExtra("TEST1")).isEqualTo("1")
+        assertThat(refined2?.baseIntentToSend?.getStringExtra("TEST2")).isEqualTo("2")
         // The clones don't have the data from each other's refinements, even though the intent
         // field is empty (thus able to be populated by filling-in).
-        assertThat(refined1.baseIntentToSend.getStringExtra("TEST2")).isNull()
-        assertThat(refined2.baseIntentToSend.getStringExtra("TEST1")).isNull()
+        assertThat(refined1?.baseIntentToSend?.getStringExtra("TEST2")).isNull()
+        assertThat(refined2?.baseIntentToSend?.getStringExtra("TEST1")).isNull()
     }
 
     @Test
@@ -301,15 +301,15 @@ class ImmutableTargetInfoTest {
         refinement.putExtra("refinement", true)
 
         val refinedResult = originalInfo.tryToCloneWithAppliedRefinement(refinement)
-        assertThat(refinedResult.baseIntentToSend.getBooleanExtra("refinement", false)).isTrue()
-        assertThat(refinedResult.baseIntentToSend.getBooleanExtra("targetAlternate", false))
+        assertThat(refinedResult?.baseIntentToSend?.getBooleanExtra("refinement", false)).isTrue()
+        assertThat(refinedResult?.baseIntentToSend?.getBooleanExtra("targetAlternate", false))
             .isTrue()
         // None of the other source intents got merged in (not even the later one that matched):
-        assertThat(refinedResult.baseIntentToSend.getBooleanExtra("originalIntent", false))
+        assertThat(refinedResult?.baseIntentToSend?.getBooleanExtra("originalIntent", false))
             .isFalse()
-        assertThat(refinedResult.baseIntentToSend.getBooleanExtra("mismatchedAlternate", false))
+        assertThat(refinedResult?.baseIntentToSend?.getBooleanExtra("mismatchedAlternate", false))
             .isFalse()
-        assertThat(refinedResult.baseIntentToSend.getBooleanExtra("extraMatch", false)).isFalse()
+        assertThat(refinedResult?.baseIntentToSend?.getBooleanExtra("extraMatch", false)).isFalse()
     }
 
     @Test
diff --git a/java/tests/src/com/android/intentresolver/chooser/TargetInfoTest.kt b/java/tests/src/com/android/intentresolver/chooser/TargetInfoTest.kt
index 886e32df..f9d3dd96 100644
--- a/java/tests/src/com/android/intentresolver/chooser/TargetInfoTest.kt
+++ b/java/tests/src/com/android/intentresolver/chooser/TargetInfoTest.kt
@@ -234,15 +234,15 @@ class TargetInfoTest {
 
         val refinedResult = originalInfo.tryToCloneWithAppliedRefinement(refinement)
         // Note `DisplayResolveInfo` targets merge refinements directly into their `resolvedIntent`.
-        assertThat(refinedResult.resolvedIntent.getBooleanExtra("refinement", false)).isTrue()
-        assertThat(refinedResult.resolvedIntent.getBooleanExtra("targetAlternate", false))
+        assertThat(refinedResult?.resolvedIntent?.getBooleanExtra("refinement", false)).isTrue()
+        assertThat(refinedResult?.resolvedIntent?.getBooleanExtra("targetAlternate", false))
             .isTrue()
         // None of the other source intents got merged in (not even the later one that matched):
-        assertThat(refinedResult.resolvedIntent.getBooleanExtra("originalIntent", false))
+        assertThat(refinedResult?.resolvedIntent?.getBooleanExtra("originalIntent", false))
             .isFalse()
-        assertThat(refinedResult.resolvedIntent.getBooleanExtra("mismatchedAlternate", false))
+        assertThat(refinedResult?.resolvedIntent?.getBooleanExtra("mismatchedAlternate", false))
             .isFalse()
-        assertThat(refinedResult.resolvedIntent.getBooleanExtra("extraMatch", false)).isFalse()
+        assertThat(refinedResult?.resolvedIntent?.getBooleanExtra("extraMatch", false)).isFalse()
     }
 
     @Test