Rework to use listener to avoid EventAccumulator / timeout complication.

Author: Michael Lehenbauer

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/testutil/EventAccumulator.java

@@ -14,7 +14,6 @@
 
 package com.google.firebase.firestore.testutil;
 
-import static com.google.firebase.firestore.testutil.IntegrationTestUtil.OPERATION_WAIT_TIMEOUT_MS;
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.waitFor;
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 
@@ -48,30 +47,22 @@ public EventListener<T> listener() {
     };
   }
 
-  public List<T> await(int numEvents, long timeoutMs) {
+  public List<T> await(int numEvents) {
     synchronized (this) {
       hardAssert(completion == null, "calling await while another await is running");
       completion = new TaskCompletionSource<>();
       maxEvents = maxEvents + numEvents;
       checkFulfilled();
     }
 
-    waitFor(completion.getTask(), timeoutMs);
+    waitFor(completion.getTask());
     completion = null;
     return events.subList(maxEvents - numEvents, maxEvents);
   }
 
-  public List<T> await(int numEvents) {
-    return await(numEvents, OPERATION_WAIT_TIMEOUT_MS);
-  }
-
   // Await 1 event.
-  public T await(long timeoutMs) {
-    return await(1, timeoutMs).get(0);
-  }
-
   public T await() {
-    return await(1, OPERATION_WAIT_TIMEOUT_MS).get(0);
+    return await(1).get(0);
   }
 
   /** Waits for a snapshot with pending writes. */

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/testutil/IntegrationTestUtil.java

@@ -16,7 +16,6 @@
 
 import static com.google.firebase.firestore.testutil.TestUtil.map;
 import static com.google.firebase.firestore.util.Util.autoId;
-import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertNull;
 import static junit.framework.Assert.fail;
 
@@ -64,7 +63,7 @@ public class IntegrationTestUtil {
   private static final Map<FirebaseFirestore, Boolean> firestoreStatus = new HashMap<>();
 
   /** Default amount of time to wait for a given operation to complete, used by waitFor() helper. */
-  static final long OPERATION_WAIT_TIMEOUT_MS = 30000;
+  private static final long OPERATION_WAIT_TIMEOUT_MS = 30000;
 
   /**
    * Firestore databases can be subject to a ~30s "cold start" delay if they have not been used
@@ -112,35 +111,43 @@ public static FirebaseFirestore testFirestore() {
    */
   public static FirebaseFirestore testFirestore(FirebaseFirestoreSettings settings) {
     FirebaseFirestore firestore = testFirestore(provider.projectId(), Level.DEBUG, settings);
-    if (!backendPrimed) {
-      backendPrimed = true;
-      primeBackend();
-    }
+    primeBackend();
     return firestore;
   }
 
   private static void primeBackend() {
-    EventAccumulator<DocumentSnapshot> accumulator = new EventAccumulator<>();
-    DocumentReference docRef = testDocument();
-    ListenerRegistration listenerRegistration = docRef.addSnapshotListener(accumulator.listener());
-
-    // Wait for watch to initialize and deliver first event.
-    accumulator.awaitRemoteEvent();
-
-    // Use a transaction to perform a write without triggering any local events.
-    docRef
-        .getFirestore()
-        .runTransaction(
-            transaction -> {
-              transaction.set(docRef, map("value", "done"));
-              return null;
-            });
-
-    // Wait to see the write on the watch stream.
-    DocumentSnapshot docSnap = accumulator.await(PRIMING_TIMEOUT_MS);
-    assertEquals("done", docSnap.get("value"));
-
-    listenerRegistration.remove();
+    if (!backendPrimed) {
+      backendPrimed = true;
+      TaskCompletionSource<Void> watchInitialized = new TaskCompletionSource<>();
+      TaskCompletionSource<Void> watchUpdateReceived = new TaskCompletionSource<>();
+      DocumentReference docRef = testDocument();
+      ListenerRegistration listenerRegistration =
+          docRef.addSnapshotListener(
+              (snapshot, error) -> {
+                if ("done".equals(snapshot.get("value"))) {
+                  watchUpdateReceived.setResult(null);
+                } else {
+                  watchInitialized.setResult(null);
+                }
+              });
+
+      // Wait for watch to initialize and deliver first event.
+      waitFor(watchInitialized.getTask());
+
+      // Use a transaction to perform a write without triggering any local events.
+      docRef
+          .getFirestore()
+          .runTransaction(
+              transaction -> {
+                transaction.set(docRef, map("value", "done"));
+                return null;
+              });
+
+      // Wait to see the write on the watch stream.
+      waitFor(watchUpdateReceived.getTask(), PRIMING_TIMEOUT_MS);
+
+      listenerRegistration.remove();
+    }
   }
 
   /** Initializes a new Firestore instance that uses a non-existing default project. */