diff --git a/x-pack/plugin/logsdb/src/javaRestTest/java/org/elasticsearch/xpack/logsdb/qa/AbstractChallengeRestTest.java b/x-pack/plugin/logsdb/src/javaRestTest/java/org/elasticsearch/xpack/logsdb/qa/AbstractChallengeRestTest.java index f4ccb4f64ae8e..e10fdbf3f85c3 100644 --- a/x-pack/plugin/logsdb/src/javaRestTest/java/org/elasticsearch/xpack/logsdb/qa/AbstractChallengeRestTest.java +++ b/x-pack/plugin/logsdb/src/javaRestTest/java/org/elasticsearch/xpack/logsdb/qa/AbstractChallengeRestTest.java @@ -60,6 +60,7 @@ public abstract class AbstractChallengeRestTest extends ESRestTestCase { .setting("xpack.security.autoconfiguration.enabled", "false") .setting("xpack.license.self_generated.type", "trial") .setting("cluster.logsdb.enabled", "true") + .setting("xpack.ml.enabled", "false") .build(); @Override diff --git a/x-pack/plugin/logsdb/src/javaRestTest/java/org/elasticsearch/xpack/logsdb/qa/BulkChallengeRestIT.java b/x-pack/plugin/logsdb/src/javaRestTest/java/org/elasticsearch/xpack/logsdb/qa/BulkChallengeRestIT.java index 3f5b69c170f51..fe532f7001420 100644 --- a/x-pack/plugin/logsdb/src/javaRestTest/java/org/elasticsearch/xpack/logsdb/qa/BulkChallengeRestIT.java +++ b/x-pack/plugin/logsdb/src/javaRestTest/java/org/elasticsearch/xpack/logsdb/qa/BulkChallengeRestIT.java @@ -12,8 +12,10 @@ import org.elasticsearch.xcontent.XContentBuilder; import java.io.IOException; +import java.util.ArrayList; import java.util.List; import java.util.Map; +import java.util.function.IntFunction; /** * Challenge test that uses bulk indexing for both baseline and contender sides. @@ -23,6 +25,8 @@ */ public class BulkChallengeRestIT extends StandardVersusLogsIndexModeChallengeRestIT { + private static final int BULK_BATCH_SIZE = 20; + public BulkChallengeRestIT() {} protected BulkChallengeRestIT(DataGenerationHelper dataGenerationHelper) { @@ -38,40 +42,62 @@ public void indexDocuments( indexBaselineDocuments(baselineSupplier, contenderResponseEntity); } - private Map indexContenderDocuments(final CheckedSupplier, IOException> documentsSupplier) + private List> indexContenderDocuments(final CheckedSupplier, IOException> documentsSupplier) throws IOException { - final StringBuilder sb = new StringBuilder(); - int id = 0; - for (var document : documentsSupplier.get()) { - if (autoGenerateId()) { - sb.append("{ \"create\": { } }\n"); - } else { - sb.append(Strings.format("{ \"create\": { \"_id\" : \"%d\" } }\n", id)); - } - sb.append(Strings.toString(document)).append("\n"); - id++; - } - return performBulkRequest(sb.toString(), false); + final IntFunction bulkActionGenerator = id -> autoGenerateId() + ? "{ \"create\": { } }\n" + : Strings.format("{ \"create\": { \"_id\" : \"%d\" } }\n", id); + + return indexDocumentsInBatches(documentsSupplier.get(), bulkActionGenerator, false); } @SuppressWarnings("unchecked") private void indexBaselineDocuments( final CheckedSupplier, IOException> documentsSupplier, - final Map contenderResponseEntity + final List> contenderItems ) throws IOException { - final StringBuilder sb = new StringBuilder(); - int id = 0; - final List> items = (List>) contenderResponseEntity.get("items"); - for (var document : documentsSupplier.get()) { + final IntFunction bulkActionGenerator = id -> { if (autoGenerateId()) { - var contenderId = ((Map) items.get(id).get("create")).get("_id"); - sb.append(Strings.format("{ \"create\": { \"_id\" : \"%s\" } }\n", contenderId)); + final var contenderId = ((Map) contenderItems.get(id).get("create")).get("_id"); + return Strings.format("{ \"create\": { \"_id\" : \"%s\" } }\n", contenderId); } else { - sb.append(Strings.format("{ \"create\": { \"_id\" : \"%d\" } }\n", id)); + return Strings.format("{ \"create\": { \"_id\" : \"%d\" } }\n", id); } + }; + + indexDocumentsInBatches(documentsSupplier.get(), bulkActionGenerator, true); + } + + @SuppressWarnings("unchecked") + private List> indexDocumentsInBatches( + final List documents, + final IntFunction bulkActionGenerator, + final boolean isBaseline + ) throws IOException { + final List> allItems = new ArrayList<>(); + final StringBuilder sb = new StringBuilder(); + int id = 0; + int batchCount = 0; + + for (final var document : documents) { + sb.append(bulkActionGenerator.apply(id)); sb.append(Strings.toString(document)).append("\n"); id++; + batchCount++; + + if (batchCount >= BULK_BATCH_SIZE) { + final Map response = performBulkRequest(sb.toString(), isBaseline); + allItems.addAll((List>) response.get("items")); + sb.setLength(0); + batchCount = 0; + } } - performBulkRequest(sb.toString(), true); + + if (batchCount > 0) { + final Map response = performBulkRequest(sb.toString(), isBaseline); + allItems.addAll((List>) response.get("items")); + } + + return allItems; } }