ReplicationStorageIT: Wait for all pushes without order

Some tests don't have a predefined order for which events will be
replicated first. Using a timeout based on a single replication event is
flawed when we don't know the expected order. Instead, use a timeout for
the group of events and ignore the order.

For two events replicating to a single remote with a single thread, we
expect the complete replication to take twice as long. Two events
replicating to two remotes will use one thread each and therefore not
take any longer than the single remote case.

Change-Id: Ieb21b7eee32105eab5b5a15a35159bb4a837e363
diff --git a/src/test/java/com/googlesource/gerrit/plugins/replication/ReplicationDaemon.java b/src/test/java/com/googlesource/gerrit/plugins/replication/ReplicationDaemon.java
index 5e38570..8fb0a19 100644
--- a/src/test/java/com/googlesource/gerrit/plugins/replication/ReplicationDaemon.java
+++ b/src/test/java/com/googlesource/gerrit/plugins/replication/ReplicationDaemon.java
@@ -27,7 +27,9 @@
 import java.nio.file.Path;
 import java.time.Duration;
 import java.util.Arrays;
+import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import org.eclipse.jgit.lib.Ref;
 import org.eclipse.jgit.lib.Repository;
@@ -114,6 +116,31 @@
     }
   }
 
+  protected boolean isPushCompleted(Map<Project.NameKey, String> refsByProject, Duration timeOut) {
+    try {
+      WaitUtil.waitUntil(
+          () -> {
+            Iterator<Map.Entry<Project.NameKey, String>> iterator =
+                refsByProject.entrySet().iterator();
+            while (iterator.hasNext()) {
+              Map.Entry<Project.NameKey, String> entry = iterator.next();
+              try (Repository repo = repoManager.openRepository(entry.getKey())) {
+                if (checkedGetRef(repo, entry.getValue()) != null) {
+                  iterator.remove();
+                }
+              } catch (IOException e) {
+                throw new RuntimeException("Cannot open repo for project" + entry.getKey(), e);
+              }
+            }
+            return refsByProject.isEmpty();
+          },
+          timeOut);
+    } catch (InterruptedException e) {
+      return false;
+    }
+    return true;
+  }
+
   protected Ref checkedGetRef(Repository repo, String branchName) {
     try {
       return repo.getRefDatabase().exactRef(branchName);
diff --git a/src/test/java/com/googlesource/gerrit/plugins/replication/ReplicationStorageIT.java b/src/test/java/com/googlesource/gerrit/plugins/replication/ReplicationStorageIT.java
index fdca243..db0f7ee 100644
--- a/src/test/java/com/googlesource/gerrit/plugins/replication/ReplicationStorageIT.java
+++ b/src/test/java/com/googlesource/gerrit/plugins/replication/ReplicationStorageIT.java
@@ -23,7 +23,9 @@
 import com.google.gerrit.reviewdb.client.Project;
 import com.googlesource.gerrit.plugins.replication.ReplicationTasksStorage.ReplicateRefUpdate;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.regex.Pattern;
 import java.util.stream.Stream;
 import org.junit.Test;
@@ -152,15 +154,20 @@
 
     String changeRef1 = createChange().getPatchSet().getRefName();
     String changeRef2 = createChange().getPatchSet().getRefName();
+    Map<Project.NameKey, String> refsByProject = new HashMap<>();
+    refsByProject.put(target1, changeRef1);
+    refsByProject.put(target2, changeRef1);
+    refsByProject.put(target1, changeRef2);
+    refsByProject.put(target2, changeRef2);
 
     setReplicationDestination(remote1, suffix1, ALL_PROJECTS);
     setReplicationDestination(remote2, suffix2, ALL_PROJECTS);
     reloadConfig();
 
-    assertThat(isPushCompleted(target1, changeRef1, TEST_PUSH_TIMEOUT)).isEqualTo(true);
-    assertThat(isPushCompleted(target2, changeRef1, TEST_PUSH_TIMEOUT)).isEqualTo(true);
-    assertThat(isPushCompleted(target1, changeRef2, TEST_PUSH_TIMEOUT)).isEqualTo(true);
-    assertThat(isPushCompleted(target2, changeRef2, TEST_PUSH_TIMEOUT)).isEqualTo(true);
+    // Wait for completion within the time 2 pushes should take because each remote only has 1
+    // thread and needs to push 2 events
+    assertThat(isPushCompleted(refsByProject, TEST_PUSH_TIMEOUT.plus(TEST_PUSH_TIMEOUT)))
+        .isEqualTo(true);
   }
 
   @Test