ReplicationStorageMPIT: Fix flaky workFromOnlyWaitingIsPerformed

This test was missing a wait for tasks to be removed from running. Also
reduce some code by re-using the list* methods from
ReplicationStorageIT.

Change-Id: Iae941e776c2d3bfa4637af5f5db1ddd017848433
diff --git a/src/test/java/com/googlesource/gerrit/plugins/replication/ReplicationStorageDaemon.java b/src/test/java/com/googlesource/gerrit/plugins/replication/ReplicationStorageDaemon.java
index bcbbc8a..f549f47 100644
--- a/src/test/java/com/googlesource/gerrit/plugins/replication/ReplicationStorageDaemon.java
+++ b/src/test/java/com/googlesource/gerrit/plugins/replication/ReplicationStorageDaemon.java
@@ -23,6 +23,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 /**
  * This class can be extended by any ReplicationStorage*IT class and provides common setup and
@@ -55,17 +56,6 @@
     replicationConfig = plugin.getSysInjector().getInstance(ReplicationConfig.class);
   }
 
-  protected boolean noIncompleteTasks() {
-    Path refUpdates = replicationConfig.getEventsDirectory().resolve("ref-updates");
-    Path runningUpdates = refUpdates.resolve("running");
-    Path waitingUpdates = refUpdates.resolve("waiting");
-    try {
-      return Files.list(runningUpdates).count() == 0 && Files.list(waitingUpdates).count() == 0;
-    } catch (IOException e) {
-      throw new RuntimeException(e.getMessage(), e);
-    }
-  }
-
   protected List<ReplicationTasksStorage.ReplicateRefUpdate> listWaitingReplicationTasks(
       String refRegex) {
     Pattern refmaskPattern = Pattern.compile(refRegex);
@@ -86,4 +76,12 @@
       }
     }
   }
+
+  protected List<ReplicationTasksStorage.ReplicateRefUpdate> listWaiting() {
+    return tasksStorage.streamWaiting().collect(Collectors.toList());
+  }
+
+  protected List<ReplicationTasksStorage.ReplicateRefUpdate> listRunning() {
+    return tasksStorage.streamRunning().collect(Collectors.toList());
+  }
 }
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 753e429..12fba80 100644
--- a/src/test/java/com/googlesource/gerrit/plugins/replication/ReplicationStorageIT.java
+++ b/src/test/java/com/googlesource/gerrit/plugins/replication/ReplicationStorageIT.java
@@ -34,7 +34,6 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.regex.Pattern;
-import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.eclipse.jgit.transport.URIish;
 import org.junit.Test;
@@ -361,12 +360,4 @@
         .filter(task -> changeRef.equals(task.ref()))
         .filter(task -> remote.equals(task.remote()));
   }
-
-  private List<ReplicateRefUpdate> listWaiting() {
-    return tasksStorage.streamWaiting().collect(Collectors.toList());
-  }
-
-  private List<ReplicateRefUpdate> listRunning() {
-    return tasksStorage.streamRunning().collect(Collectors.toList());
-  }
 }
diff --git a/src/test/java/com/googlesource/gerrit/plugins/replication/ReplicationStorageMPIT.java b/src/test/java/com/googlesource/gerrit/plugins/replication/ReplicationStorageMPIT.java
index 65fc4df..e14efe9 100644
--- a/src/test/java/com/googlesource/gerrit/plugins/replication/ReplicationStorageMPIT.java
+++ b/src/test/java/com/googlesource/gerrit/plugins/replication/ReplicationStorageMPIT.java
@@ -15,7 +15,6 @@
 package com.googlesource.gerrit.plugins.replication;
 
 import static com.google.common.truth.Truth.assertThat;
-import static org.junit.Assert.assertTrue;
 
 import com.google.gerrit.acceptance.TestPlugin;
 import com.google.gerrit.acceptance.UseLocalDisk;
@@ -69,6 +68,7 @@
       assertThat(targetBranchRefB.getObjectId()).isEqualTo(masterRef.getObjectId());
     }
 
-    assertTrue(noIncompleteTasks());
+    WaitUtil.waitUntil(() -> listRunning().isEmpty(), TEST_TASK_FINISH_TIMEOUT);
+    assertThat(listWaiting()).isEmpty();
   }
 }