Replace getForegroundActivity with applyToForegroundActivity[Task]

firebase-appdistribution/src/main/java/com/google/firebase/appdistribution/AabUpdater.java

@@ -16,7 +16,6 @@
 
 import static com.google.firebase.appdistribution.FirebaseAppDistributionException.Status.DOWNLOAD_FAILURE;
 import static com.google.firebase.appdistribution.FirebaseAppDistributionException.Status.NETWORK_FAILURE;
-import static com.google.firebase.appdistribution.TaskUtils.combineWithResultOf;
 import static com.google.firebase.appdistribution.TaskUtils.runAsyncInTask;
 import static com.google.firebase.appdistribution.TaskUtils.safeSetTaskException;
 
@@ -97,10 +96,10 @@ UpdateTaskImpl updateAab(@NonNull AppDistributionReleaseInternal newRelease) {
 
       // On a background thread, fetch the redirect URL and open it in the Play app
       runAsyncInTask(executor, () -> fetchDownloadRedirectUrl(newRelease.getDownloadUrl()))
-          .onSuccessTask(combineWithResultOf(() -> lifecycleNotifier.getForegroundActivity()))
-          .addOnSuccessListener(
-              urlAndActivity ->
-                  openRedirectUrlInPlay(urlAndActivity.first(), urlAndActivity.second()))
+          .onSuccessTask(
+              redirectUrl ->
+                  lifecycleNotifier.applyToForegroundActivity(
+                      activity -> openRedirectUrlInPlay(redirectUrl, activity)))
           .addOnFailureListener(this::setUpdateTaskCompletionError);
 
       return cachedUpdateTask;

firebase-appdistribution/src/main/java/com/google/firebase/appdistribution/ApkUpdater.java

@@ -113,8 +113,8 @@ UpdateTaskImpl updateApk(
 
   private void installApk(File file, boolean showDownloadNotificationManager) {
     lifeCycleNotifier
-        .getForegroundActivity()
-        .onSuccessTask(taskExecutor, activity -> apkInstaller.installApk(file.getPath(), activity))
+        .applyToForegroundActivityTask(
+            activity -> apkInstaller.installApk(file.getPath(), activity))
         .addOnSuccessListener(
             taskExecutor,
             unused -> {

firebase-appdistribution/src/main/java/com/google/firebase/appdistribution/FirebaseAppDistribution.java

@@ -142,8 +142,7 @@ public UpdateTask updateIfNewReleaseAvailable() {
     }
 
     lifecycleNotifier
-        .getForegroundActivity()
-        .onSuccessTask(this::showSignInConfirmationDialog)
+        .applyToForegroundActivityTask(this::showSignInConfirmationDialog)
         // TODO(rachelprince): Revisit this comment once changes to checkForNewRelease are reviewed
         // Even though checkForNewRelease() calls signInTester(), we explicitly call signInTester
         // here both for code clarifty, and because we plan to remove the signInTester() call
@@ -173,9 +172,8 @@ public UpdateTask updateIfNewReleaseAvailable() {
                 setCachedUpdateIfNewReleaseResult();
                 return Tasks.forResult(null);
               }
-              return lifecycleNotifier
-                  .getForegroundActivity()
-                  .onSuccessTask((activity) -> showUpdateAlertDialog(activity, release));
+              return lifecycleNotifier.applyToForegroundActivityTask(
+                  activity -> showUpdateAlertDialog(activity, release));
             })
         .onSuccessTask(
             unused ->

firebase-appdistribution/src/main/java/com/google/firebase/appdistribution/FirebaseAppDistributionLifecycleNotifier.java

@@ -20,14 +20,25 @@
 import androidx.annotation.GuardedBy;
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
+import com.google.android.gms.tasks.SuccessContinuation;
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.TaskCompletionSource;
 import com.google.android.gms.tasks.Tasks;
 import java.util.ArrayDeque;
 import java.util.Queue;
+import java.util.concurrent.Executor;
 
 class FirebaseAppDistributionLifecycleNotifier implements Application.ActivityLifecycleCallbacks {
 
+  /** An {@link Executor} that runs tasks on the current thread. */
+  private static final Executor DIRECT_EXECUTOR = Runnable::run;
+
+  /** A functional interface for a function that takes an activity and does something with it. */
+  interface ActivityConsumer<T> {
+    void consume(Activity activity);
+  }
+
   private static FirebaseAppDistributionLifecycleNotifier instance;
   private final Object lock = new Object();
 
@@ -54,7 +65,8 @@ class FirebaseAppDistributionLifecycleNotifier implements Application.ActivityLi
   @GuardedBy("lock")
   private final Queue<OnActivityDestroyedListener> onDestroyedListeners = new ArrayDeque<>();
 
-  private FirebaseAppDistributionLifecycleNotifier() {}
+  @VisibleForTesting
+  FirebaseAppDistributionLifecycleNotifier() {}
 
   static synchronized FirebaseAppDistributionLifecycleNotifier getInstance() {
     if (instance == null) {
@@ -84,13 +96,43 @@ interface OnActivityDestroyedListener {
   }
 
   /**
-   * Get a {@link Task} that will succeed with a result of the app's foregrounded {@link Activity},
-   * when one is available.
-   *
-   * <p>The returned task will never fail. It will instead remain pending indefinitely until some
-   * activity comes to the foreground.
+   * Apply a function to a foreground activity, when one is available, returning a {@link Task} that
+   * will complete immediately after the function is applied.
+   */
+  Task<Void> applyToForegroundActivity(ActivityConsumer consumer) {
+    return getForegroundActivity()
+        .onSuccessTask(
+            // Use direct executor to ensure the consumer is called while Activity is in foreground
+            DIRECT_EXECUTOR,
+            activity -> {
+              try {
+                consumer.consume(activity);
+                return Tasks.forResult(null);
+              } catch (Throwable t) {
+                return Tasks.forException(FirebaseAppDistributionException.wrap(t));
+              }
+            });
+  }
+
+  /**
+   * Apply a function to a foreground activity, when one is available, returning a {@link Task} that
+   * will complete with the result of the Task returned by that function.
    */
-  Task<Activity> getForegroundActivity() {
+  <T> Task<T> applyToForegroundActivityTask(SuccessContinuation<Activity, T> continuation) {
+    return getForegroundActivity()
+        .onSuccessTask(
+            // Use direct executor to ensure the consumer is called while Activity is in foreground
+            DIRECT_EXECUTOR,
+            activity -> {
+              try {
+                return continuation.then(activity);
+              } catch (Throwable t) {
+                return Tasks.forException(FirebaseAppDistributionException.wrap(t));
+              }
+            });
+  }
+
+  private Task<Activity> getForegroundActivity() {
     synchronized (lock) {
       if (currentActivity != null) {
         return Tasks.forResult(currentActivity);

firebase-appdistribution/src/main/java/com/google/firebase/appdistribution/TaskUtils.java

@@ -14,11 +14,9 @@
 
 package com.google.firebase.appdistribution;
 
-import com.google.android.gms.tasks.SuccessContinuation;
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.TaskCompletionSource;
 import com.google.android.gms.tasks.Tasks;
-import com.google.auto.value.AutoValue;
 import com.google.firebase.appdistribution.FirebaseAppDistributionException.Status;
 import com.google.firebase.appdistribution.internal.LogWrapper;
 import java.util.concurrent.Executor;
@@ -34,11 +32,6 @@ interface Operation<TResult> {
     TResult run() throws FirebaseAppDistributionException;
   }
 
-  /** A functional interface to wrap a function that produces a {@link Task}. */
-  interface TaskSource<TResult> {
-    Task<TResult> get();
-  }
-
   /**
    * Runs a long running operation inside a {@link Task}, wrapping any errors in {@link
    * FirebaseAppDistributionException}.
@@ -92,58 +85,6 @@ static <TResult> Task<TResult> handleTaskFailure(Task<TResult> task) {
     return task;
   }
 
-  /**
-   * An @{link AutoValue} class to hold the result of two Tasks, combined using {@link
-   * #combineWithResultOf}.
-   *
-   * @param <T1> The result type of the first task
-   * @param <T2> The result type of the second task
-   */
-  @AutoValue
-  abstract static class CombinedTaskResults<T1, T2> {
-    abstract T1 first();
-
-    abstract T2 second();
-
-    static <T1, T2> CombinedTaskResults<T1, T2> create(T1 first, T2 second) {
-      return new AutoValue_TaskUtils_CombinedTaskResults(first, second);
-    }
-  }
-
-  /**
-   * Returns a {@link SuccessContinuation} to be chained off of a {@link Task}, that will run
-   * another task in sequence and combine both results together.
-   *
-   * <p>This is useful when you want to run two tasks and use the results of each, but those tasks
-   * need to be run sequentially. If they can be run in parallel, use {@link Tasks#whenAll} or one
-   * of its variations.
-   *
-   * <p>Usage:
-   *
-   * <pre>{@code
-   * runFirstAsyncTask()
-   *   .onSuccessTask(combineWithResultOf(executor, () -> startSecondAsyncTask())
-   *   .addOnSuccessListener(
-   *       results ->
-   *           doSomethingWithBothResults(results.result1(), results.result2()));
-   * }</pre>
-   *
-   * @param secondTaskSource A {@link TaskSource} providing the next task to run
-   * @param <T1> The result type of the first task
-   * @param <T2> The result type of the second task
-   * @return A {@link SuccessContinuation} that will return a new task with result type {@link
-   *     CombinedTaskResults}, combining the results of both tasks
-   */
-  static <T1, T2> SuccessContinuation<T1, CombinedTaskResults<T1, T2>> combineWithResultOf(
-      TaskSource<T2> secondTaskSource) {
-    return firstResult ->
-        secondTaskSource
-            .get()
-            .onSuccessTask(
-                secondResult ->
-                    Tasks.forResult(CombinedTaskResults.create(firstResult, secondResult)));
-  }
-
   static void safeSetTaskException(TaskCompletionSource taskCompletionSource, Exception e) {
     if (taskCompletionSource != null && !taskCompletionSource.getTask().isComplete()) {
       taskCompletionSource.setException(e);

firebase-appdistribution/src/main/java/com/google/firebase/appdistribution/TesterSignInManager.java

@@ -15,7 +15,6 @@
 package com.google.firebase.appdistribution;
 
 import static com.google.firebase.appdistribution.FirebaseAppDistributionException.Status.AUTHENTICATION_CANCELED;
-import static com.google.firebase.appdistribution.TaskUtils.combineWithResultOf;
 import static com.google.firebase.appdistribution.TaskUtils.safeSetTaskException;
 import static com.google.firebase.appdistribution.TaskUtils.safeSetTaskResult;
 
@@ -141,28 +140,30 @@ public Task<Void> signInTester() {
           .getId()
           .addOnFailureListener(
               handleTaskFailure(ErrorMessages.AUTHENTICATION_ERROR, Status.AUTHENTICATION_FAILURE))
-          .onSuccessTask(combineWithResultOf(() -> lifecycleNotifier.getForegroundActivity()))
-          .addOnSuccessListener(
-              fidAndActivity -> {
-                // Launch the intent outside of the synchronized block because we don't need to wait
-                // for the lock, and we don't want to risk the activity leaving the foreground in
-                // the meantime.
-                openSignInFlowInBrowser(fidAndActivity.first(), fidAndActivity.second());
-                // This synchronized block is required by the @GuardedBy annotation, but is not
-                // practically required in this case because the only reads of this variable are on
-                // the main thread, which this callback is also running on.
-                synchronized (signInTaskLock) {
-                  hasBeenSentToBrowserForCurrentTask = true;
-                }
-              })
-          // No failures expected here, since getForegroundActivity() will wait indefinitely for a
-          // foreground activity, but catch any unexpected failures to be safe.
+          .onSuccessTask(this::getForegroundActivityAndOpenSignInFlow)
+          // Catch any unexpected failures to be safe.
           .addOnFailureListener(handleTaskFailure(ErrorMessages.UNKNOWN_ERROR, Status.UNKNOWN));
 
       return signInTaskCompletionSource.getTask();
     }
   }
 
+  private Task<Void> getForegroundActivityAndOpenSignInFlow(String fid) {
+    return lifecycleNotifier.applyToForegroundActivity(
+        activity -> {
+          // Launch the intent outside of the synchronized block because we don't need to wait
+          // for the lock, and we don't want to risk the activity leaving the foreground in
+          // the meantime.
+          openSignInFlowInBrowser(fid, activity);
+          // This synchronized block is required by the @GuardedBy annotation, but is not
+          // practically required in this case because the only reads of this variable are on
+          // the main thread, which this callback is also running on.
+          synchronized (signInTaskLock) {
+            hasBeenSentToBrowserForCurrentTask = true;
+          }
+        });
+  }
+
   private OnFailureListener handleTaskFailure(String message, Status status) {
     return e -> {
       LogWrapper.getInstance().e(TAG + message, e);

firebase-appdistribution/src/test/java/com/google/firebase/appdistribution/AabUpdaterTest.java

@@ -14,16 +14,17 @@
 package com.google.firebase.appdistribution;
 
 import static com.google.common.truth.Truth.assertThat;
+import static com.google.firebase.appdistribution.TestUtils.applyToForegroundActivityAnswer;
 import static com.google.firebase.appdistribution.TestUtils.assertTaskFailure;
 import static com.google.firebase.appdistribution.TestUtils.awaitAsyncOperations;
 import static org.junit.Assert.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.when;
 import static org.robolectric.Shadows.shadowOf;
 import static org.robolectric.annotation.LooperMode.Mode.PAUSED;
 
 import android.app.Activity;
 import android.net.Uri;
-import com.google.android.gms.tasks.Tasks;
 import com.google.firebase.FirebaseApp;
 import com.google.firebase.appdistribution.Constants.ErrorMessages;
 import com.google.firebase.appdistribution.FirebaseAppDistributionException.Status;
@@ -88,7 +89,9 @@ public void setup() throws IOException, FirebaseAppDistributionException {
     aabUpdater =
         Mockito.spy(
             new AabUpdater(mockLifecycleNotifier, mockHttpsUrlConnectionFactory, testExecutor));
-    when(mockLifecycleNotifier.getForegroundActivity()).thenReturn(Tasks.forResult(activity));
+
+    when(mockLifecycleNotifier.applyToForegroundActivity(any()))
+        .thenAnswer(applyToForegroundActivityAnswer(activity));
   }
 
   @Test

firebase-appdistribution/src/test/java/com/google/firebase/appdistribution/ApkUpdaterTest.java

@@ -15,6 +15,7 @@
 package com.google.firebase.appdistribution;
 
 import static com.google.common.truth.Truth.assertThat;
+import static com.google.firebase.appdistribution.TestUtils.applyToForegroundActivityTaskAnswer;
 import static org.junit.Assert.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doReturn;
@@ -102,7 +103,8 @@ public void setup() throws IOException, FirebaseAppDistributionException {
     when(mockFile.length()).thenReturn(TEST_FILE_LENGTH);
     when(mockHttpsUrlConnectionFactory.openConnection(TEST_URL)).thenReturn(mockHttpsUrlConnection);
     when(mockHttpsUrlConnection.getResponseCode()).thenReturn(200);
-    when(mockLifecycleNotifier.getForegroundActivity()).thenReturn(Tasks.forResult(activity));
+    when(mockLifecycleNotifier.applyToForegroundActivityTask(any()))
+        .thenAnswer(applyToForegroundActivityTaskAnswer(activity));
     onCompleteListener = new TestOnCompleteListener<>();
 
     apkUpdater =
@@ -169,36 +171,6 @@ public void updateApk_whenInstallSuccessful_setsResult() throws Exception {
     assertThat(updateTask.isSuccessful()).isTrue();
   }
 
-  @Test
-  public void updateApk_getForegroundActivityFails_setsError() {
-    boolean showNotification = true;
-    doReturn(Tasks.forResult(mockFile))
-        .when(apkUpdater)
-        .downloadApk(TEST_RELEASE, showNotification);
-    TaskCompletionSource<Activity> getForegroundActivityCompletionSource =
-        new TaskCompletionSource<>();
-    when(mockLifecycleNotifier.getForegroundActivity())
-        .thenReturn(getForegroundActivityCompletionSource.getTask());
-    UpdateTask updateTask = apkUpdater.updateApk(TEST_RELEASE, showNotification);
-    updateTask.addOnCompleteListener(testExecutor, onCompleteListener);
-    List<UpdateProgress> progressEvents = new ArrayList<>();
-    updateTask.addOnProgressListener(testExecutor, progressEvents::add);
-
-    getForegroundActivityCompletionSource.setException(
-        new FirebaseAppDistributionException(
-            Constants.ErrorMessages.APK_INSTALLATION_FAILED,
-            FirebaseAppDistributionException.Status.INSTALLATION_FAILURE));
-
-    assertThat(updateTask.isComplete()).isFalse();
-    FirebaseAppDistributionException e =
-        assertThrows(FirebaseAppDistributionException.class, () -> onCompleteListener.await());
-    assertThat(e.getErrorCode()).isEqualTo(Status.INSTALLATION_FAILURE);
-    assertThat(progressEvents).hasSize(1);
-    assertThat(progressEvents.get(0).getUpdateStatus()).isEqualTo(UpdateStatus.INSTALL_FAILED);
-    assertThat(updateTask.isSuccessful()).isFalse();
-    verify(mockNotificationsManager).updateNotification(1000, 1000, UpdateStatus.INSTALL_FAILED);
-  }
-
   @Test
   public void updateApk_whenInstallFailed_setsError() {
     boolean showNotification = true;