Improved SetMaxBatchesTest

marklogic-client-api-functionaltests/src/test/java/com/marklogic/client/fastfunctest/datamovement/SetMaxBatchesTest.java

@@ -6,15 +6,18 @@
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
-@Disabled("Getting this test in place to verify that setMaxBatches does not appear to work yet.")
 public class SetMaxBatchesTest extends AbstractFunctionalTest {
 
+	@Disabled("Getting this test in place to verify that setMaxBatches does not work when used with a query.")
 	@Test
-	void test() {
+	void testWithQuery() {
 		writeJsonDocs(50, "max-batches-test");
 
 		DataMovementManager dmm = client.newDataMovementManager();
@@ -33,4 +36,32 @@ void test() {
 			"expect 20 URIs back. But through 6.2.2 (and probably going back much further), all URIs are returned. " +
 			"Modifying the thread count and batch size do not appear to affect this at all.");
 	}
+
+	/**
+	 * This verifies that setMaxBatches works with an iterator. The feature appears to have been introduced in 5.1.0,
+	 * so since then, it's only ever worked for an iterator. It does not work when an actual query is involved.
+	 */
+	@Test
+	void iteratorTest() {
+		List<String> results = new ArrayList<>();
+
+		List<String> input = new ArrayList<>();
+		for (int i = 1; i <= 100; i++) {
+			input.add(i + "");
+		}
+
+		DataMovementManager dmm = client.newDataMovementManager();
+		QueryBatcher queryBatcher = dmm
+			.newQueryBatcher(input.iterator())
+			.withThreadCount(4)
+			.withBatchSize(10)
+			.onUrisReady(batch -> results.addAll(Arrays.asList(batch.getItems())));
+
+		queryBatcher.setMaxBatches(2);
+		dmm.startJob(queryBatcher);
+		queryBatcher.awaitCompletion();
+		dmm.stopJob(queryBatcher);
+
+		assertEquals(20, results.size());
+	}
 }