Merge "Use a ViewState to preserve refinement state" into udc-dev

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