KAFKA-9509: Fixing flakiness of MirrorConnectorsIntegrationTest.testReplication

connect/mirror/src/test/java/org/apache/kafka/connect/mirror/MirrorConnectorsIntegrationTest.java

@@ -16,9 +16,14 @@
  */
 package org.apache.kafka.connect.mirror;
 
+import java.util.Optional;
+import org.apache.kafka.connect.runtime.AbstractStatus;
+import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo;
+import org.apache.kafka.connect.runtime.rest.errors.ConnectRestException;
 import org.apache.kafka.connect.util.clusters.EmbeddedConnectCluster;
 import org.apache.kafka.connect.util.clusters.EmbeddedKafkaCluster;
 import org.apache.kafka.test.IntegrationTest;
+import org.apache.kafka.test.TestUtils;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.admin.Admin;
@@ -34,9 +39,9 @@
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.Properties;
-import java.util.concurrent.TimeoutException;
 import java.time.Duration;
 
 import static org.junit.Assert.assertEquals;
@@ -60,13 +65,14 @@ public class MirrorConnectorsIntegrationTest {
     private static final int NUM_PARTITIONS = 10;
     private static final int RECORD_TRANSFER_DURATION_MS = 10_000;
     private static final int CHECKPOINT_DURATION_MS = 20_000;
+    private static final int MIN_TASKS = 1;
 
     private MirrorMakerConfig mm2Config; 
     private EmbeddedConnectCluster primary;
     private EmbeddedConnectCluster backup;
 
     @Before
-    public void setup() throws IOException {
+    public void setup() throws IOException, InterruptedException {
         Properties brokerProps = new Properties();
         brokerProps.put("auto.create.topics.enable", "false");
 
@@ -150,23 +156,67 @@ public void setup() throws IOException {
         mm2Props.put("backup.bootstrap.servers", backup.kafka().bootstrapServers());
         mm2Config = new MirrorMakerConfig(mm2Props);
 
+        // we wait for the connector and tasks to come up for each connector, so that when we do the
+        // actual testing, we are certain that the tasks are up and running; this will prevent
+        // flaky tests where the connector and tasks didn't start up in time for the tests to be
+        // run
         backup.configureConnector("MirrorSourceConnector", mm2Config.connectorBaseConfig(new SourceAndTarget("primary", "backup"),
             MirrorSourceConnector.class));
+        TestUtils.waitForCondition(() -> assertConnectorAndTasksRunning(backup,
+            "MirrorSourceConnector", MIN_TASKS).isPresent(),
+            "Timed out trying to verify connector MirrorSourceConnector was up on"
+                + " backup cluster!");
 
         backup.configureConnector("MirrorCheckpointConnector", mm2Config.connectorBaseConfig(new SourceAndTarget("primary", "backup"),
             MirrorCheckpointConnector.class));
+        TestUtils.waitForCondition(() -> assertConnectorAndTasksRunning(backup,
+            "MirrorCheckpointConnector", MIN_TASKS).isPresent(),
+            "Timed out trying to verify connector MirrorCheckpointConnector was up on"
+                + " backup cluster!");
 
         backup.configureConnector("MirrorHeartbeatConnector", mm2Config.connectorBaseConfig(new SourceAndTarget("primary", "backup"),
             MirrorHeartbeatConnector.class));
+        TestUtils.waitForCondition(() -> assertConnectorAndTasksRunning(backup,
+            "MirrorHeartbeatConnector", MIN_TASKS).isPresent(),
+            "Timed out trying to verify connector MirrorHeartbeatConnector was up on"
+                + " backup cluster!");
 
         primary.configureConnector("MirrorSourceConnector", mm2Config.connectorBaseConfig(new SourceAndTarget("backup", "primary"),
             MirrorSourceConnector.class));
+        TestUtils.waitForCondition(() -> assertConnectorAndTasksRunning(primary,
+            "MirrorSourceConnector", MIN_TASKS).isPresent(),
+            "Timed out trying to verify connector MirrorSourceConnector was up on"
+                + " primary cluster!");
 
         primary.configureConnector("MirrorCheckpointConnector", mm2Config.connectorBaseConfig(new SourceAndTarget("backup", "primary"),
             MirrorCheckpointConnector.class));
+        TestUtils.waitForCondition(() -> assertConnectorAndTasksRunning(primary,
+            "MirrorCheckpointConnector", MIN_TASKS).isPresent(),
+            "Timed out trying to verify connector MirrorCheckpointConnector was up on"
+                + " primary cluster!");
 
         primary.configureConnector("MirrorHeartbeatConnector", mm2Config.connectorBaseConfig(new SourceAndTarget("backup", "primary"),
             MirrorHeartbeatConnector.class));
+        TestUtils.waitForCondition(() -> assertConnectorAndTasksRunning(primary,
+            "MirrorHeartbeatConnector", MIN_TASKS).isPresent(),
+            "Timed out trying to verify connector MirrorHeartbeatConnector was up on"
+                + " primary cluster!");
+    }
+
+    private Optional<Boolean> assertConnectorAndTasksRunning(EmbeddedConnectCluster connectCluster,
+        String connectorName, int numTasks) {
+        try {
+            ConnectorStateInfo info = connectCluster.connectorStatus(connectorName);
+            boolean result = info != null
+                && info.tasks().size() == numTasks
+                && info.connector().state().equals(AbstractStatus.State.RUNNING.toString())
+                && info.tasks().stream().allMatch(s -> s.state().equals(AbstractStatus.State.RUNNING.toString()));
+            log.info("Found connector and tasks running: {}", result);
+            return Optional.of(result);
+        } catch (Exception e) {
+            log.error("Could not check connector state info.", e);
+            return Optional.empty();
+        }
     }
 
     @After
@@ -184,7 +234,7 @@ public void close() throws IOException {
     }
 
     @Test
-    public void testReplication() throws InterruptedException, TimeoutException {
+    public void testReplication() throws InterruptedException {
         MirrorClient primaryClient = new MirrorClient(mm2Config.clientConfig("primary"));
         MirrorClient backupClient = new MirrorClient(mm2Config.clientConfig("backup"));