Simplify CsvTestsDataLoader (#142717)

Author: idegtiarenko

x-pack/plugin/esql/qa/server/multi-clusters/src/javaRestTest/java/org/elasticsearch/xpack/esql/ccq/MultiClusterSpecIT.java

@@ -41,12 +41,12 @@
 import java.util.Locale;
 import java.util.Set;
 import java.util.regex.Pattern;
-import java.util.stream.Collectors;
 
+import static java.util.stream.Collectors.toSet;
 import static org.elasticsearch.xpack.esql.CsvSpecReader.specParser;
 import static org.elasticsearch.xpack.esql.CsvTestUtils.isEnabled;
 import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.CSV_DATASET_MAP;
-import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.ENRICH_SOURCE_INDICES;
+import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.ENRICH_POLICIES;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.classpathResources;
 import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.COMPLETION;
 import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.ENABLE_FORK_FOR_REMOTE_INDICES_V2;
@@ -266,13 +266,11 @@ protected RestClient buildClient(Settings settings, HttpHost[] localHosts) throw
         .stream()
         .filter(td -> td.settingFileName() != null && td.settingFileName().equals("lookup-settings.json"))
         .map(CsvTestsDataLoader.TestDataset::indexName)
-        .collect(Collectors.toSet());
+        .collect(toSet());
 
-    public static final Set<String> LOOKUP_ENDPOINTS = LOOKUP_INDICES.stream().map(i -> "/" + i + "/_bulk").collect(Collectors.toSet());
+    public static final Set<String> LOOKUP_ENDPOINTS = LOOKUP_INDICES.stream().map(i -> "/" + i + "/_bulk").collect(toSet());
 
-    public static final Set<String> ENRICH_ENDPOINTS = ENRICH_SOURCE_INDICES.stream()
-        .map(i -> "/" + i + "/_bulk")
-        .collect(Collectors.toSet());
+    public static final Set<String> ENRICH_ENDPOINTS = ENRICH_POLICIES.stream().map(p -> "/" + p.index() + "/_bulk").collect(toSet());
 
     /**
      * Creates a new mock client that dispatches every request to both the local and remote clusters, excluding _bulk, _query,
@@ -357,7 +355,8 @@ static CsvSpecReader.CsvTestCase convertToRemoteIndices(CsvSpecReader.CsvTestCas
         boolean onlyRemotes = canUseRemoteIndicesOnly() && randomBoolean();
         // Check if query contains enrich source indices - these are loaded into both clusters,
         // so we should use onlyRemotes=true to avoid duplicates
-        if (onlyRemotes == false && EsqlTestUtils.queryContainsIndices(query, Set.copyOf(ENRICH_SOURCE_INDICES))) {
+        var enrichSourceIndices = ENRICH_POLICIES.stream().map(CsvTestsDataLoader.EnrichConfig::index).collect(toSet());
+        if (onlyRemotes == false && EsqlTestUtils.queryContainsIndices(query, enrichSourceIndices)) {
             onlyRemotes = true;
         }
         testCase.query = EsqlTestUtils.addRemoteIndices(testCase.query, LOOKUP_INDICES, onlyRemotes);

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/KnnSemanticTextTestCase.java

@@ -137,21 +137,16 @@ private void setupIndex() throws IOException {
     }
 
     private void setupInferenceEndpoints() throws IOException {
-        CsvTestsDataLoader.createTextEmbeddingInferenceEndpoint(client());
-        CsvTestsDataLoader.createSparseEmbeddingInferenceEndpoint(client());
+        CsvTestsDataLoader.createInferenceEndpoint(client(), CsvTestsDataLoader.INFERENCE_CONFIGS.get("test_dense_inference"));
+        CsvTestsDataLoader.createInferenceEndpoint(client(), CsvTestsDataLoader.INFERENCE_CONFIGS.get("test_sparse_inference"));
     }
 
     @After
     public void tearDown() throws Exception {
         super.tearDown();
         client().performRequest(new Request("DELETE", "semantic-test"));
-
-        if (CsvTestsDataLoader.clusterHasTextEmbeddingInferenceEndpoint(client())) {
-            CsvTestsDataLoader.deleteTextEmbeddingInferenceEndpoint(client());
-        }
-        if (CsvTestsDataLoader.clusterHasSparseEmbeddingInferenceEndpoint(client())) {
-            CsvTestsDataLoader.deleteSparseEmbeddingInferenceEndpoint(client());
-        }
+        CsvTestsDataLoader.deleteInferenceEndpoint(client(), "test_dense_inference");
+        CsvTestsDataLoader.deleteInferenceEndpoint(client(), "test_sparse_inference");
     }
 
     private Map<String, Object> runEsqlQuery(String query) throws IOException {

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/RestRerankTestCase.java

@@ -11,6 +11,7 @@
 import org.elasticsearch.client.ResponseException;
 import org.elasticsearch.test.rest.ESRestTestCase;
 import org.elasticsearch.xpack.esql.AssertWarnings;
+import org.elasticsearch.xpack.esql.CsvTestsDataLoader;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
@@ -19,8 +20,6 @@
 import java.util.List;
 import java.util.Map;
 
-import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.createRerankInferenceEndpoint;
-import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.deleteRerankInferenceEndpoint;
 import static org.hamcrest.core.StringContains.containsString;
 
 public class RestRerankTestCase extends ESRestTestCase {
@@ -36,7 +35,7 @@ public void assertRequestBreakerEmpty() throws Exception {
 
     @Before
     public void setUpInferenceEndpoint() throws IOException {
-        createRerankInferenceEndpoint(adminClient());
+        CsvTestsDataLoader.createInferenceEndpoint(adminClient(), CsvTestsDataLoader.INFERENCE_CONFIGS.get("test_reranker"));
     }
 
     @Before
@@ -77,7 +76,7 @@ public void wipeData() throws IOException {
             }
         }
 
-        deleteRerankInferenceEndpoint(adminClient());
+        CsvTestsDataLoader.deleteInferenceEndpoint(adminClient(), "test_reranker");
     }
 
     public void testRerankWithSingleField() throws IOException {

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeRestTest.java

@@ -433,7 +433,7 @@ private static List<String> originalTypes(Map<String, ?> x) {
 
     private List<String> availableIndices() throws IOException {
         return availableDatasetsForEs(true, supportsSourceFieldMapping(), false, requiresTimeSeries(), cap -> false).stream()
-            .filter(x -> x.requiresInferenceEndpoint() == false)
+            .filter(x -> x.inferenceEndpoints().isEmpty())
             .map(x -> x.indexName())
             .toList();
     }