[TEST] Cover custom sorting and routing in randomized testing

x-pack/plugin/logsdb/src/javaRestTest/java/org/elasticsearch/xpack/logsdb/qa/AbstractChallengeRestTest.java

@@ -14,7 +14,6 @@
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.CheckedConsumer;
-import org.elasticsearch.core.Tuple;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.test.cluster.ElasticsearchCluster;
@@ -28,8 +27,11 @@
 
 import java.io.IOException;
 import java.util.List;
+import java.util.Map;
 import java.util.function.Supplier;
 
+import static org.hamcrest.Matchers.equalTo;
+
 public abstract class AbstractChallengeRestTest extends ESRestTestCase {
     private final String baselineDataStreamName;
     private final String contenderDataStreamName;
@@ -227,37 +229,69 @@ public void contenderSettings(Settings.Builder builder) {}
 
     public void commonSettings(Settings.Builder builder) {}
 
-    private Response indexDocuments(
-        final String dataStreamName,
-        final CheckedSupplier<List<XContentBuilder>, IOException> documentsSupplier
-    ) throws IOException {
+    public boolean autoGenerateId() {
+        return false;
+    }
+
+    public Map<String, Object> indexContenderDocuments(final CheckedSupplier<List<XContentBuilder>, IOException> documentsSupplier)
+        throws IOException {
         final StringBuilder sb = new StringBuilder();
         int id = 0;
         for (var document : documentsSupplier.get()) {
-            sb.append(Strings.format("{ \"create\": { \"_id\" : \"%d\" } }", id)).append("\n");
+            if (autoGenerateId()) {
+                sb.append("{ \"create\": { } }\n");
+            } else {
+                sb.append(Strings.format("{ \"create\": { \"_id\" : \"%d\" } }\n", id));
+            }
             sb.append(Strings.toString(document)).append("\n");
             id++;
         }
-        var request = new Request("POST", "/" + dataStreamName + "/_bulk");
+        var request = new Request("POST", "/" + getContenderDataStreamName() + "/_bulk");
         request.setJsonEntity(sb.toString());
         request.addParameter("refresh", "true");
-        return client.performRequest(request);
+        var response = client.performRequest(request);
+        assertOK(response);
+        var responseBody = entityAsMap(response);
+        assertThat("errors in contender bulk response:\n " + responseBody, responseBody.get("errors"), equalTo(false));
+        return responseBody;
     }
 
-    public Response indexBaselineDocuments(final CheckedSupplier<List<XContentBuilder>, IOException> documentsSupplier) throws IOException {
-        return indexDocuments(getBaselineDataStreamName(), documentsSupplier);
-    }
-
-    public Response indexContenderDocuments(final CheckedSupplier<List<XContentBuilder>, IOException> documentsSupplier)
-        throws IOException {
-        return indexDocuments(getContenderDataStreamName(), documentsSupplier);
+    @SuppressWarnings("unchecked")
+    public Map<String, Object> indexBaselineDocuments(
+        final CheckedSupplier<List<XContentBuilder>, IOException> documentsSupplier,
+        final Map<String, Object> contenderResponseEntity
+    ) throws IOException {
+        final StringBuilder sb = new StringBuilder();
+        int id = 0;
+        final List<Map<String, Object>> items = contenderResponseEntity != null
+            ? (List<Map<String, Object>>) contenderResponseEntity.get("items")
+            : List.of();
+        for (var document : documentsSupplier.get()) {
+            if (autoGenerateId() && items.isEmpty() == false) {
+                var contenderId = ((Map<String, Object>) items.get(id).get("create")).get("_id");
+                sb.append(Strings.format("{ \"create\": { \"_id\" : \"%s\" } }\n", contenderId));
+            } else {
+                sb.append(Strings.format("{ \"create\": { \"_id\" : \"%d\" } }\n", id));
+            }
+            sb.append(Strings.toString(document)).append("\n");
+            id++;
+        }
+        var request = new Request("POST", "/" + getBaselineDataStreamName() + "/_bulk");
+        request.setJsonEntity(sb.toString());
+        request.addParameter("refresh", "true");
+        var response = client.performRequest(request);
+        assertOK(response);
+        var responseBody = entityAsMap(response);
+        assertThat("errors in baseline bulk response:\n " + responseBody, responseBody.get("errors"), equalTo(false));
+        return responseBody;
     }
 
-    public Tuple<Response, Response> indexDocuments(
+    public void indexDocuments(
         final CheckedSupplier<List<XContentBuilder>, IOException> baselineSupplier,
         final CheckedSupplier<List<XContentBuilder>, IOException> contenderSupplier
     ) throws IOException {
-        return new Tuple<>(indexBaselineDocuments(baselineSupplier), indexContenderDocuments(contenderSupplier));
+        var contenderResponseEntity = indexContenderDocuments(contenderSupplier);
+        indexBaselineDocuments(baselineSupplier, contenderResponseEntity);
     }
 
     public Response queryBaseline(final SearchSourceBuilder search) throws IOException {

x-pack/plugin/logsdb/src/javaRestTest/java/org/elasticsearch/xpack/logsdb/qa/ReindexChallengeRestIT.java

@@ -8,19 +8,20 @@
 package org.elasticsearch.xpack.logsdb.qa;
 
 import org.elasticsearch.client.Request;
-import org.elasticsearch.client.Response;
 import org.elasticsearch.common.CheckedSupplier;
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
 import java.util.List;
 import java.util.Locale;
+import java.util.Map;
 
 import static org.hamcrest.Matchers.equalTo;
 
 public abstract class ReindexChallengeRestIT extends StandardVersusLogsIndexModeRandomDataChallengeRestIT {
     @Override
-    public Response indexContenderDocuments(CheckedSupplier<List<XContentBuilder>, IOException> documentsSupplier) throws IOException {
+    public Map<String, Object> indexContenderDocuments(CheckedSupplier<List<XContentBuilder>, IOException> documentsSupplier)
+        throws IOException {
         var reindexRequest = new Request("POST", "/_reindex?refresh=true");
         reindexRequest.setJsonEntity(String.format(Locale.ROOT, """
             {
@@ -39,6 +40,15 @@ public Response indexContenderDocuments(CheckedSupplier<List<XContentBuilder>, I
         var body = entityAsMap(response);
         assertThat("encountered failures when performing reindex:\n " + body, body.get("failures"), equalTo(List.of()));
 
-        return response;
+        return body;
+    }
+
+    @Override
+    public void indexDocuments(
+        final CheckedSupplier<List<XContentBuilder>, IOException> baselineSupplier,
+        final CheckedSupplier<List<XContentBuilder>, IOException> contenderSupplier
+    ) throws IOException {
+        indexBaselineDocuments(baselineSupplier, null);
+        indexContenderDocuments(contenderSupplier);
     }
 }

x-pack/plugin/logsdb/src/javaRestTest/java/org/elasticsearch/xpack/logsdb/qa/StandardVersusLogsIndexModeChallengeRestIT.java

@@ -11,13 +11,11 @@
 import org.elasticsearch.client.Response;
 import org.elasticsearch.client.ResponseException;
 import org.elasticsearch.client.RestClient;
-import org.elasticsearch.common.CheckedSupplier;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.time.DateFormatter;
 import org.elasticsearch.common.time.FormatNames;
 import org.elasticsearch.common.xcontent.XContentHelper;
 import org.elasticsearch.index.query.QueryBuilders;
-import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.search.aggregations.AggregationBuilders;
 import org.elasticsearch.search.aggregations.bucket.histogram.DateHistogramInterval;
 import org.elasticsearch.search.aggregations.bucket.histogram.HistogramAggregationBuilder;
@@ -29,7 +27,6 @@
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.logsdb.qa.matchers.MatchResult;
 import org.elasticsearch.xpack.logsdb.qa.matchers.Matcher;
-import org.hamcrest.Matchers;
 
 import java.io.IOException;
 import java.time.Instant;
@@ -42,7 +39,6 @@
 import java.util.Map;
 import java.util.TreeMap;
 
-import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
 
 /**
@@ -54,6 +50,8 @@ public class StandardVersusLogsIndexModeChallengeRestIT extends AbstractChalleng
     private final int numShards = randomBoolean() ? randomIntBetween(2, 4) : 0;
     private final int numReplicas = randomBoolean() ? randomIntBetween(1, 3) : 0;
     private final boolean fullyDynamicMapping = randomBoolean();
+    private final boolean useCustomSortConfig = fullyDynamicMapping == false && randomBoolean();
+    private final boolean routeOnSortFields = useCustomSortConfig && randomBoolean();
 
     public StandardVersusLogsIndexModeChallengeRestIT() {
         super("standard-apache-baseline", "logs-apache-contender", "baseline-template", "contender-template", 101, 101);
@@ -159,6 +157,13 @@ public void commonSettings(Settings.Builder builder) {
     @Override
     public void contenderSettings(Settings.Builder builder) {
         builder.put("index.mode", "logsdb");
+        if (useCustomSortConfig) {
+            builder.putList("index.sort.field", "host.name", "method", "@timestamp");
+            builder.putList("index.sort.order", "asc", "asc", "desc");
+            if (routeOnSortFields) {
+                builder.put("index.logsdb.route_on_sort_fields", true);
+            }
+        }
     }
 
     @Override
@@ -169,6 +174,11 @@ public void beforeStart() throws Exception {
         waitForLogs(client());
     }
 
+    @Override
+    public boolean autoGenerateId() {
+        return routeOnSortFields;
+    }
+
     protected static void waitForLogs(RestClient client) throws Exception {
         assertBusy(() -> {
             try {
@@ -330,28 +340,6 @@ public void testFieldCaps() throws IOException {
         assertTrue(matchResult.getMessage(), matchResult.isMatch());
     }
 
-    @Override
-    public Response indexBaselineDocuments(CheckedSupplier<List<XContentBuilder>, IOException> documentsSupplier) throws IOException {
-        var response = super.indexBaselineDocuments(documentsSupplier);
-
-        assertThat(response.getStatusLine().getStatusCode(), Matchers.equalTo(RestStatus.OK.getStatus()));
-        var baselineResponseBody = entityAsMap(response);
-        assertThat("errors in baseline bulk response:\n " + baselineResponseBody, baselineResponseBody.get("errors"), equalTo(false));
-
-        return response;
-    }
-
-    @Override
-    public Response indexContenderDocuments(CheckedSupplier<List<XContentBuilder>, IOException> documentsSupplier) throws IOException {
-        var response = super.indexContenderDocuments(documentsSupplier);
-
-        assertThat(response.getStatusLine().getStatusCode(), Matchers.equalTo(RestStatus.OK.getStatus()));
-        var contenderResponseBody = entityAsMap(response);
-        assertThat("errors in contender bulk response:\n " + contenderResponseBody, contenderResponseBody.get("errors"), equalTo(false));
-
-        return response;
-    }
-
     private List<XContentBuilder> generateDocuments(int numberOfDocuments) throws IOException {
         final List<XContentBuilder> documents = new ArrayList<>();
         // This is static in order to be able to identify documents between test runs.
@@ -383,7 +371,7 @@ private static List<Map<String, Object>> getQueryHits(final Response response) t
         assertThat(hitsList.size(), greaterThan(0));
 
         return hitsList.stream()
-            .sorted(Comparator.comparingInt((Map<String, Object> hit) -> Integer.parseInt((String) hit.get("_id"))))
+            .sorted(Comparator.comparing((Map<String, Object> hit) -> ((String) hit.get("_id"))))
             .map(hit -> (Map<String, Object>) hit.get("_source"))
             .toList();
     }
@@ -404,7 +392,7 @@ private static List<Map<String, Object>> getEsqlSourceResults(final Response res
 
         // Results contain a list of [source, id] lists.
         return values.stream()
-            .sorted(Comparator.comparingInt((List<Object> value) -> Integer.parseInt((String) value.get(1))))
+            .sorted(Comparator.comparing((List<Object> value) -> ((String) value.get(1))))
             .map(value -> (Map<String, Object>) value.get(0))
             .toList();
     }