diff --git a/marklogic-client-api-functionaltests/src/test/java/com/marklogic/client/fastfunctest/datamovement/SetMaxBatchesTest.java b/marklogic-client-api-functionaltests/src/test/java/com/marklogic/client/fastfunctest/datamovement/SetMaxBatchesTest.java index 576f8750f..b834ae268 100644 --- a/marklogic-client-api-functionaltests/src/test/java/com/marklogic/client/fastfunctest/datamovement/SetMaxBatchesTest.java +++ b/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 results = new ArrayList<>(); + + List 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()); + } }