Use a latch for the other usage of awaitProgressEvents()

Author: lfkellogg

firebase-appdistribution/src/test/java/com/google/firebase/appdistribution/impl/FirebaseAppDistributionServiceImplTest.java

@@ -86,6 +86,7 @@
 import com.google.firebase.installations.InstallationTokenResult;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -364,11 +365,29 @@ public void updateIfNewReleaseAvailable_whenNoReleaseAvailable_updateDialogNotSh
   @Test
   public void updateIfNewReleaseAvailable_whenActivityBackgrounded_updateDialogNotShown()
       throws InterruptedException {
-    when(mockNewReleaseFetcher.checkForNewRelease()).thenReturn(Tasks.forResult(null));
-
+    // Block thread making the request on a latch, which gives us time to add listeners to the
+    // returned UpdateTask in time to get all the progress updates
+    CountDownLatch countDownLatch = new CountDownLatch(1);
+    when(mockNewReleaseFetcher.checkForNewRelease())
+        .thenAnswer(
+            invocation -> {
+              try {
+                countDownLatch.await();
+              } catch (InterruptedException e) {
+                throw new AssertionError("Interrupted while waiting in mock");
+              }
+              return Tasks.forResult(null);
+            });
+
+    // Start the update
     UpdateTask task = firebaseAppDistribution.updateIfNewReleaseAvailable();
 
-    List<UpdateProgress> progressEvents = awaitProgressEvents(task, 1);
+    // Listen for progress events
+    List<UpdateProgress> progressEvents = new ArrayList<>();
+    task.addOnProgressListener(FirebaseExecutors.directExecutor(), progressEvents::add);
+    countDownLatch.countDown();
+    awaitProgressEvents(progressEvents, 1);
+
     assertEquals(UpdateStatus.NEW_RELEASE_NOT_AVAILABLE, progressEvents.get(0).getUpdateStatus());
     assertNull(ShadowAlertDialog.getLatestAlertDialog());
   }

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

@@ -29,16 +29,12 @@
 import com.google.firebase.appdistribution.FirebaseAppDistributionException;
 import com.google.firebase.appdistribution.FirebaseAppDistributionException.Status;
 import com.google.firebase.appdistribution.UpdateProgress;
-import com.google.firebase.appdistribution.UpdateTask;
 import com.google.firebase.appdistribution.impl.FirebaseAppDistributionLifecycleNotifier.ActivityConsumer;
 import com.google.firebase.appdistribution.impl.FirebaseAppDistributionLifecycleNotifier.ActivityFunction;
-import com.google.firebase.concurrent.FirebaseExecutors;
 import com.google.firebase.concurrent.TestOnlyExecutors;
 import java.io.IOException;
 import java.io.InputStream;
-import java.util.ArrayList;
 import java.util.Collection;
-import java.util.List;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -102,15 +98,6 @@ static void awaitAsyncOperations(ExecutorService executorService) throws Interru
     shadowOf(getMainLooper()).idle();
   }
 
-  /** Await a specified number of progress events on an {@link UpdateTask}. */
-  static List<UpdateProgress> awaitProgressEvents(UpdateTask updateTask, int count)
-      throws InterruptedException {
-    List<UpdateProgress> progressEvents = new ArrayList<>();
-    updateTask.addOnProgressListener(FirebaseExecutors.directExecutor(), progressEvents::add);
-    awaitProgressEvents(progressEvents, count);
-    return progressEvents;
-  }
-
   /** Await a specified number of progress events being added to the given collection. */
   static void awaitProgressEvents(Collection<UpdateProgress> progressEvents, int count)
       throws InterruptedException {