Test: Improved test speed by refactoring getProduced

getProduced can now check if it contains the expected results and
return quicker, rather than waiting for max retry count.

Author: rozza

src/integrationTest/java/com/mongodb/kafka/connect/MongoSourceConnectorTest.java

@@ -488,7 +488,7 @@ void testSourceSurvivesARestartWhenCopyingExisting() {
     List<ChangeStreamOperation> inserts = createInserts(1, 5000);
     inserts.addAll(createInserts(10001, 10050));
 
-    assertEventuallyProduces(inserts, coll1);
+    assertProduced(inserts, coll1.getNamespace().getFullName(), 60);
   }
 
   @Test

src/integrationTest/java/com/mongodb/kafka/connect/embedded/EmbeddedKafka.java

@@ -280,14 +280,14 @@ public void beforeAll(final ExtensionContext context) throws Exception {
   }
 
   @Override
-  public void afterEach(final ExtensionContext context) throws InterruptedException {
+  public void afterEach(final ExtensionContext context) {
     deleteSinkConnector();
     deleteSourceConnector();
-    deleteTopicsAndWait(Duration.ofMinutes(4));
   }
 
   @Override
-  public void afterAll(final ExtensionContext context) {
+  public void afterAll(final ExtensionContext context) throws InterruptedException {
+    deleteTopicsAndWait(Duration.ofMinutes(4));
     stop();
   }
 

src/integrationTest/java/com/mongodb/kafka/connect/mongodb/ChangeStreamOperations.java

@@ -22,6 +22,8 @@
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
+import org.apache.kafka.common.utils.Bytes;
+
 import org.bson.Document;
 
 public class ChangeStreamOperations {
@@ -55,7 +57,9 @@ public static ChangeStreamOperation createInsert(final int id) {
     return new Insert(id);
   }
 
-  public static ChangeStreamOperation createChangeStreamOperation(final String changeStreamJson) {
+  public static ChangeStreamOperation createChangeStreamOperation(
+      final Bytes changeStreamJsonBytes) {
+    String changeStreamJson = changeStreamJsonBytes.toString();
     Document document = Document.parse(changeStreamJson);
     ChangeStreamOperation changeStreamOperation;
     switch (document.get("operationType", "unknown").toLowerCase()) {

src/integrationTest/java/com/mongodb/kafka/connect/mongodb/MongoKafkaTestCase.java

@@ -16,19 +16,19 @@
 package com.mongodb.kafka.connect.mongodb;
 
 import static com.mongodb.kafka.connect.mongodb.ChangeStreamOperations.ChangeStreamOperation;
-import static com.mongodb.kafka.connect.mongodb.ChangeStreamOperations.createChangeStreamOperation;
 import static java.lang.String.format;
 import static java.util.Collections.singletonList;
 import static org.apache.kafka.common.utils.Utils.sleep;
-import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertIterableEquals;
 
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Properties;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
 import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
 import io.confluent.connect.avro.AvroConverter;
 
@@ -57,6 +57,8 @@ public class MongoKafkaTestCase {
   protected static final Logger LOGGER = LoggerFactory.getLogger(MongoKafkaTestCase.class);
   protected static final AtomicInteger POSTFIX = new AtomicInteger();
 
+  private static final int DEFAULT_MAX_RETRIES = 15;
+
   @RegisterExtension public static final EmbeddedKafka KAFKA = new EmbeddedKafka();
   @RegisterExtension public static final MongoDBHelper MONGODB = new MongoDBHelper();
 
@@ -108,88 +110,90 @@ public boolean isGreaterThanThreeDotSix() {
   }
 
   public void assertProduced(final String topicName, final int expectedCount) {
-    assertEquals(expectedCount, getProduced(topicName, expectedCount).size());
+    List<Integer> expected = IntStream.range(1, expectedCount).boxed().collect(Collectors.toList());
+    AtomicInteger counter = new AtomicInteger();
+    List<Integer> produced =
+        getProduced(topicName, b -> counter.addAndGet(1), expected, DEFAULT_MAX_RETRIES);
+    assertIterableEquals(expected, produced);
   }
 
   public void assertProduced(
       final List<ChangeStreamOperation> operationTypes, final MongoCollection<?> coll) {
-    assertProduced(operationTypes, coll.getNamespace().getFullName());
+    assertProduced(operationTypes, coll, DEFAULT_MAX_RETRIES);
   }
 
   public void assertProduced(
-      final List<ChangeStreamOperation> operationTypes, final String topicName) {
-    List<ChangeStreamOperation> produced =
-        getProduced(topicName, operationTypes.size()).stream()
-            .map((b) -> createChangeStreamOperation(b.toString()))
-            .collect(Collectors.toList());
-    assertIterableEquals(operationTypes, produced);
+      final List<ChangeStreamOperation> operationTypes,
+      final MongoCollection<?> coll,
+      final int maxRetryCount) {
+    assertProduced(operationTypes, coll.getNamespace().getFullName(), maxRetryCount);
   }
 
-  public void assertEventuallyProduces(
-      final List<ChangeStreamOperation> operationTypes, final MongoCollection<?> coll) {
-    assertEventuallyProduces(operationTypes, coll.getNamespace().getFullName());
+  public void assertProduced(
+      final List<ChangeStreamOperation> operationTypes, final String topicName) {
+    assertProduced(operationTypes, topicName, DEFAULT_MAX_RETRIES);
   }
 
-  public void assertEventuallyProduces(
-      final List<ChangeStreamOperation> operationTypes, final String topicName) {
+  public void assertProduced(
+      final List<ChangeStreamOperation> operationTypes,
+      final String topicName,
+      final int maxRetryCount) {
     List<ChangeStreamOperation> produced =
-        getProduced(topicName, Integer.MAX_VALUE).stream()
-            .map((b) -> createChangeStreamOperation(b.toString()))
-            .collect(Collectors.toList());
-
-    if (produced.size() > operationTypes.size()) {
-      boolean startsWith =
-          produced
-              .get(operationTypes.size() - 1)
-              .equals(operationTypes.get(operationTypes.size() - 1));
-      if (startsWith) {
-        assertIterableEquals(operationTypes, produced.subList(0, operationTypes.size()));
-      } else {
-        assertIterableEquals(
+        getProduced(
+            topicName,
+            ChangeStreamOperations::createChangeStreamOperation,
             operationTypes,
-            produced.subList(produced.lastIndexOf(operationTypes.get(0)), produced.size()));
-      }
-    } else {
-      assertIterableEquals(operationTypes, produced);
-    }
+            maxRetryCount);
+    assertIterableEquals(operationTypes, produced);
   }
 
   public void assertProducedDocs(final List<Document> docs, final MongoCollection<?> coll) {
-    assertEquals(
-        docs,
-        getProduced(coll.getNamespace().getFullName(), docs.size()).stream()
-            .map((b) -> Document.parse(b.toString()))
-            .collect(Collectors.toList()));
+    List<Document> produced =
+        getProduced(
+            coll.getNamespace().getFullName(),
+            b -> Document.parse(b.toString()),
+            docs,
+            DEFAULT_MAX_RETRIES);
+    assertIterableEquals(docs, produced);
   }
 
-  public List<Bytes> getProduced(final String topicName, final int expectedCount) {
-    if (expectedCount != Integer.MAX_VALUE) {
-      LOGGER.info("Subscribing to {} expecting to see #{}", topicName, expectedCount);
-    } else {
-      LOGGER.info("Subscribing to {} getting all messages", topicName);
-    }
+  public <T> List<T> getProduced(
+      final String topicName,
+      final Function<Bytes, T> mapper,
+      final List<T> expected,
+      final int maxRetryCount) {
+    LOGGER.info("Subscribing to {}", topicName);
 
     try (KafkaConsumer<?, ?> consumer = createConsumer()) {
       consumer.subscribe(singletonList(topicName));
-      List<Bytes> data = new ArrayList<>();
+      List<T> data = new ArrayList<>();
+      T firstExpected = expected.isEmpty() ? null : expected.get(0);
+      T lastExpected = expected.isEmpty() ? null : expected.get(expected.size() - 1);
       int counter = 0;
       int retryCount = 0;
       int previousDataSize;
-      while (data.size() < expectedCount && retryCount < 30) {
+
+      while (retryCount < maxRetryCount) {
         counter++;
-        LOGGER.info("Polling {} ({}) seen: #{}", topicName, counter, data.size());
         previousDataSize = data.size();
-
         consumer
             .poll(Duration.ofSeconds(2))
             .records(topicName)
-            .forEach((r) -> data.add((Bytes) r.value()));
+            .forEach((r) -> data.add(mapper.apply((Bytes) r.value())));
+
+        int firstExpectedIndex = data.lastIndexOf(firstExpected);
+        int lastExpectedIndex = data.lastIndexOf(lastExpected);
+        int dataSize = lastExpectedIndex - firstExpectedIndex + 1;
+        if (firstExpectedIndex > -1 && lastExpectedIndex > -1 && dataSize == expected.size()) {
+          return data.subList(firstExpectedIndex, lastExpectedIndex + 1);
+        }
 
         // Wait at least 3 minutes for the first set of data to arrive
-        if (data.size() > 0 || counter > 90) {
-          retryCount = data.size() == previousDataSize ? retryCount + 1 : 0;
+        if (expected.size() == 0 || data.size() > 0 || counter > 90) {
+          retryCount += previousDataSize == data.size() ? 1 : 0;
         }
       }
+
       return data;
     }
   }