Allow AAB install to be started from a background thread

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

@@ -16,10 +16,11 @@
 
 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;
 
 import android.app.Activity;
-import android.content.Context;
 import android.content.Intent;
 import android.net.Uri;
 import androidx.annotation.GuardedBy;
@@ -42,29 +43,28 @@ class AabUpdater {
   private final HttpsUrlConnectionFactory httpsUrlConnectionFactory;
   private final Executor executor;
 
+  private final Object updateAabLock = new Object();
+
   @GuardedBy("updateAabLock")
   private UpdateTaskImpl cachedUpdateTask;
 
   @GuardedBy("updateAabLock")
   private AppDistributionReleaseInternal aabReleaseInProgress;
 
-  private final Object updateAabLock = new Object();
-  private final Context context;
+  @GuardedBy("updateAabLock")
+  private boolean hasBeenSentToPlayForCurrentTask = false;
 
-  AabUpdater(@NonNull Context context) {
+  AabUpdater() {
     this(
-        context,
         FirebaseAppDistributionLifecycleNotifier.getInstance(),
         new HttpsUrlConnectionFactory(),
         Executors.newSingleThreadExecutor());
   }
 
   AabUpdater(
-      @NonNull Context context,
       @NonNull FirebaseAppDistributionLifecycleNotifier lifecycleNotifier,
       @NonNull HttpsUrlConnectionFactory httpsUrlConnectionFactory,
       @NonNull Executor executor) {
-    this.context = context;
     this.lifecycleNotifier = lifecycleNotifier;
     this.httpsUrlConnectionFactory = httpsUrlConnectionFactory;
     lifecycleNotifier.addOnActivityStartedListener(this::onActivityStarted);
@@ -76,9 +76,13 @@ void onActivityStarted(Activity activity) {
     if (activity instanceof SignInResultActivity || activity instanceof InstallActivity) {
       return;
     }
-    // If app resumes and aab update task is in progress, assume that installation didn't happen so
-    // cancel the task
-    this.tryCancelAabUpdateTask();
+    // If app resumes and update is in progress, assume that installation didn't happen and cancel
+    // the task
+    synchronized (updateAabLock) {
+      if (awaitingUpdateFromPlay()) {
+        this.tryCancelAabUpdateTask();
+      }
+    }
   }
 
   UpdateTaskImpl updateAab(@NonNull AppDistributionReleaseInternal newRelease) {
@@ -89,7 +93,15 @@ UpdateTaskImpl updateAab(@NonNull AppDistributionReleaseInternal newRelease) {
 
       cachedUpdateTask = new UpdateTaskImpl();
       aabReleaseInProgress = newRelease;
-      redirectToPlayForAabUpdate(newRelease.getDownloadUrl());
+      hasBeenSentToPlayForCurrentTask = false;
+
+      // 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()))
+          .addOnFailureListener(this::setUpdateTaskCompletionError);
 
       return cachedUpdateTask;
     }
@@ -138,38 +150,30 @@ private static boolean isRedirectResponse(int responseCode) {
     return responseCode >= 300 && responseCode < 400;
   }
 
-  private void redirectToPlayForAabUpdate(String downloadUrl) {
-    // The 302 redirect is obtained here to open the play store directly and avoid opening chrome
-    executor.execute( // Execute the network calls on a background thread
-        () -> {
-          String redirectUrl;
-          try {
-            redirectUrl = fetchDownloadRedirectUrl(downloadUrl);
-          } catch (FirebaseAppDistributionException e) {
-            setUpdateTaskCompletionError(e);
-            return;
-          }
-
-          Intent updateIntent = new Intent(Intent.ACTION_VIEW);
-          updateIntent.setData(Uri.parse(redirectUrl));
-          updateIntent.addFlags(Intent.FLAG_ACTIVITY_NEW_TASK);
-          LogWrapper.getInstance().v(TAG + "Redirecting to play");
-
-          synchronized (updateAabLock) {
-            context.startActivity(updateIntent);
-            cachedUpdateTask.updateProgress(
-                UpdateProgress.builder()
-                    .setApkBytesDownloaded(-1)
-                    .setApkFileTotalBytes(-1)
-                    .setUpdateStatus(UpdateStatus.REDIRECTED_TO_PLAY)
-                    .build());
-          }
-        });
+  private void openRedirectUrlInPlay(String redirectUrl, Activity hostActivity) {
+    Intent updateIntent = new Intent(Intent.ACTION_VIEW);
+    updateIntent.setData(Uri.parse(redirectUrl));
+    updateIntent.addFlags(Intent.FLAG_ACTIVITY_NEW_TASK);
+    LogWrapper.getInstance().v(TAG + "Redirecting to play");
+
+    // Launch the intent before the synchronized block to avoid failing to update in the rare
+    // scenario where the activity moves to the background while we're awaiting the lock.
+    hostActivity.startActivity(updateIntent);
+
+    synchronized (updateAabLock) {
+      cachedUpdateTask.updateProgress(
+          UpdateProgress.builder()
+              .setApkBytesDownloaded(-1)
+              .setApkFileTotalBytes(-1)
+              .setUpdateStatus(UpdateStatus.REDIRECTED_TO_PLAY)
+              .build());
+      hasBeenSentToPlayForCurrentTask = true;
+    }
   }
 
-  private void setUpdateTaskCompletionError(FirebaseAppDistributionException e) {
+  private void setUpdateTaskCompletionError(Exception e) {
     synchronized (updateAabLock) {
-      safeSetTaskException(cachedUpdateTask, e);
+      safeSetTaskException(cachedUpdateTask, FirebaseAppDistributionException.wrap(e));
     }
   }
 
@@ -183,4 +187,12 @@ void tryCancelAabUpdateTask() {
               ReleaseUtils.convertToAppDistributionRelease(aabReleaseInProgress)));
     }
   }
+
+  private boolean awaitingUpdateFromPlay() {
+    synchronized (updateAabLock) {
+      return cachedUpdateTask != null
+          && !cachedUpdateTask.isComplete()
+          && hasBeenSentToPlayForCurrentTask;
+    }
+  }
 }

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

@@ -102,7 +102,7 @@ public class FirebaseAppDistribution {
             new FirebaseAppDistributionTesterApiClient(),
             firebaseInstallationsApiProvider),
         new ApkUpdater(firebaseApp, new ApkInstaller()),
-        new AabUpdater(firebaseApp.getApplicationContext()),
+        new AabUpdater(),
         signInStorage,
         lifecycleNotifier);
   }

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

@@ -122,7 +122,7 @@ static <T1, T2> CombinedTaskResults<T1, T2> create(T1 first, T2 second) {
    *
    * <pre>{@code
    * runFirstAsyncTask()
-   *   .onSuccessTask(combineWithResultOf(() -> startSecondAsyncTask())
+   *   .onSuccessTask(combineWithResultOf(executor, () -> startSecondAsyncTask())
    *   .addOnSuccessListener(
    *       results ->
    *           doSomethingWithBothResults(results.result1(), results.result2()));

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

@@ -15,24 +15,24 @@
 
 import static com.google.common.truth.Truth.assertThat;
 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.junit.Assert.assertThrows;
 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 androidx.test.core.app.ApplicationProvider;
 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;
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
-import java.util.concurrent.TimeUnit;
 import javax.net.ssl.HttpsURLConnection;
 import org.junit.Before;
 import org.junit.Test;
@@ -42,9 +42,11 @@
 import org.mockito.MockitoAnnotations;
 import org.robolectric.Robolectric;
 import org.robolectric.RobolectricTestRunner;
+import org.robolectric.annotation.LooperMode;
 import org.robolectric.shadows.ShadowActivity;
 
 @RunWith(RobolectricTestRunner.class)
+@LooperMode(PAUSED)
 public class AabUpdaterTest {
   private static final String TEST_URL = "https://test-url";
   private static final String REDIRECT_TO_PLAY = "https://redirect-to-play-url";
@@ -85,11 +87,7 @@ public void setup() throws IOException, FirebaseAppDistributionException {
 
     aabUpdater =
         Mockito.spy(
-            new AabUpdater(
-                ApplicationProvider.getApplicationContext(),
-                mockLifecycleNotifier,
-                mockHttpsUrlConnectionFactory,
-                testExecutor));
+            new AabUpdater(mockLifecycleNotifier, mockHttpsUrlConnectionFactory, testExecutor));
     when(mockLifecycleNotifier.getForegroundActivity()).thenReturn(Tasks.forResult(activity));
   }
 
@@ -100,7 +98,7 @@ public void updateAppTask_whenOpenConnectionFails_setsNetworkFailure()
     when(mockHttpsUrlConnectionFactory.openConnection(TEST_URL)).thenThrow(caughtException);
 
     UpdateTask updateTask = aabUpdater.updateAab(TEST_RELEASE_NEWER_AAB_INTERNAL);
-    testExecutor.awaitTermination(1, TimeUnit.SECONDS);
+    awaitAsyncOperations(testExecutor);
 
     assertTaskFailure(
         updateTask, Status.NETWORK_FAILURE, "Failed to open connection", caughtException);
@@ -112,7 +110,7 @@ public void updateAppTask_isNotRedirectResponse_setsDownloadFailure()
     when(mockHttpsUrlConnection.getResponseCode()).thenReturn(200);
 
     UpdateTask updateTask = aabUpdater.updateAab(TEST_RELEASE_NEWER_AAB_INTERNAL);
-    testExecutor.awaitTermination(1, TimeUnit.SECONDS);
+    awaitAsyncOperations(testExecutor);
 
     assertTaskFailure(updateTask, Status.DOWNLOAD_FAILURE, "Expected redirect");
   }
@@ -123,7 +121,7 @@ public void updateAppTask_missingLocationHeader_setsDownloadFailure()
     when(mockHttpsUrlConnection.getHeaderField("Location")).thenReturn(null);
 
     UpdateTask updateTask = aabUpdater.updateAab(TEST_RELEASE_NEWER_AAB_INTERNAL);
-    testExecutor.awaitTermination(1, TimeUnit.SECONDS);
+    awaitAsyncOperations(testExecutor);
 
     assertTaskFailure(updateTask, Status.DOWNLOAD_FAILURE, "No Location header");
   }
@@ -133,18 +131,18 @@ public void updateAppTask_emptyLocationHeader_setsDownloadFailure() throws Inter
     when(mockHttpsUrlConnection.getHeaderField("Location")).thenReturn("");
 
     UpdateTask updateTask = aabUpdater.updateAab(TEST_RELEASE_NEWER_AAB_INTERNAL);
-    testExecutor.awaitTermination(1, TimeUnit.SECONDS);
+    awaitAsyncOperations(testExecutor);
 
     assertTaskFailure(updateTask, Status.DOWNLOAD_FAILURE, "Empty Location header");
   }
 
   @Test
-  public void updateAppTask_whenAabReleaseAvailable_redirectsToPlay() throws InterruptedException {
+  public void updateAppTask_whenAabReleaseAvailable_redirectsToPlay() throws Exception {
     List<UpdateProgress> progressEvents = new ArrayList<>();
 
     UpdateTask updateTask = aabUpdater.updateAab(TEST_RELEASE_NEWER_AAB_INTERNAL);
     updateTask.addOnProgressListener(testExecutor, progressEvents::add);
-    testExecutor.awaitTermination(1, TimeUnit.SECONDS);
+    awaitAsyncOperations(testExecutor);
 
     assertThat(shadowActivity.getNextStartedActivity().getData())
         .isEqualTo(Uri.parse(REDIRECT_TO_PLAY));
@@ -159,14 +157,13 @@ public void updateAppTask_whenAabReleaseAvailable_redirectsToPlay() throws Inter
   }
 
   @Test
-  public void updateAppTask_onAppResume_setsUpdateCancelled() {
-    TestOnCompleteListener<Void> onCompleteListener = new TestOnCompleteListener<>();
+  public void updateAppTask_onAppResume_setsUpdateCancelled() throws InterruptedException {
     UpdateTask updateTask = aabUpdater.updateAab(TEST_RELEASE_NEWER_AAB_INTERNAL);
-    updateTask.addOnCompleteListener(testExecutor, onCompleteListener);
-
+    awaitAsyncOperations(testExecutor);
     aabUpdater.onActivityStarted(activity);
+
     FirebaseAppDistributionException exception =
-        assertThrows(FirebaseAppDistributionException.class, onCompleteListener::await);
+        assertTaskFailure(updateTask, Status.INSTALLATION_CANCELED, ErrorMessages.UPDATE_CANCELED);
     assertEquals(
         ReleaseUtils.convertToAppDistributionRelease(TEST_RELEASE_NEWER_AAB_INTERNAL),
         exception.getRelease());

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

@@ -26,6 +26,9 @@
  * Helper listener that works around a limitation of the Tasks API where await() cannot be called on
  * the main thread. This listener works around it by running itself on a different thread, thus
  * allowing the main thread to be woken up when the Tasks complete.
+ *
+ * <p>Note: Calling {@link #await()} from a Robolectric test does block the main thread, since those
+ * tests are executed on the main thread.
  */
 public class TestOnCompleteListener<TResult> implements OnCompleteListener<TResult> {
   private static final long TIMEOUT_MS = 5000;

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

@@ -14,25 +14,40 @@
 
 package com.google.firebase.appdistribution;
 
+import static android.os.Looper.getMainLooper;
 import static com.google.common.truth.Truth.assertThat;
+import static org.robolectric.Shadows.shadowOf;
 
 import com.google.android.gms.tasks.Task;
 import com.google.firebase.appdistribution.FirebaseAppDistributionException.Status;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
 
 final class TestUtils {
   private TestUtils() {}
 
-  static void assertTaskFailure(Task task, Status status, String messageSubstring) {
+  static FirebaseAppDistributionException assertTaskFailure(
+      Task task, Status status, String messageSubstring) {
     assertThat(task.isSuccessful()).isFalse();
     assertThat(task.getException()).isInstanceOf(FirebaseAppDistributionException.class);
     FirebaseAppDistributionException e = (FirebaseAppDistributionException) task.getException();
     assertThat(e.getErrorCode()).isEqualTo(status);
     assertThat(e).hasMessageThat().contains(messageSubstring);
+    return e;
   }
 
   static void assertTaskFailure(
       Task task, Status status, String messageSubstring, Throwable cause) {
     assertTaskFailure(task, status, messageSubstring);
     assertThat(task.getException()).hasCauseThat().isEqualTo(cause);
   }
+
+  static void awaitAsyncOperations(ExecutorService executorService) throws InterruptedException {
+    // Await anything enqueued to the executor
+    executorService.awaitTermination(100, TimeUnit.MILLISECONDS);
+
+    // Idle the main looper, which is also running these tests, so any Task or lifecycle callbacks
+    // can be handled. See http://robolectric.org/blog/2019/06/04/paused-looper/ for more info.
+    shadowOf(getMainLooper()).idle();
+  }
 }