From a4897c91558304c61fbd0f00e090e1790a7ceae5 Mon Sep 17 00:00:00 2001
From: Nik Everett <nik9000@gmail.com>
Date: Wed, 28 May 2025 16:25:25 -0400
Subject: [PATCH] ESQL: Speed up semantic_text tests (#128591)

Speeds up the semantic_text tests by using a test inference plugin. That
skips downloading the normal inference setup.

Closes #128511 Closes #128513 Closes #128571 Closes #128572 Closes
#128573 Closes #128574
---
 .../esql/qa/single_node/PushQueriesIT.java    | 55 ++++++++++++++++---
 1 file changed, 47 insertions(+), 8 deletions(-)

diff --git a/x-pack/plugin/esql/qa/server/single-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/single_node/PushQueriesIT.java b/x-pack/plugin/esql/qa/server/single-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/single_node/PushQueriesIT.java
index 671950b682506..5fa86d2f12189 100644
--- a/x-pack/plugin/esql/qa/server/single-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/single_node/PushQueriesIT.java
+++ b/x-pack/plugin/esql/qa/server/single-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/single_node/PushQueriesIT.java
@@ -9,9 +9,7 @@
 
 import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
 import com.carrotsearch.randomizedtesting.annotations.ThreadLeakFilters;
-import com.carrotsearch.randomizedtesting.annotations.TimeoutSuite;
 
-import org.apache.lucene.tests.util.TimeUnits;
 import org.elasticsearch.client.Request;
 import org.elasticsearch.client.Response;
 import org.elasticsearch.client.ResponseException;
@@ -25,6 +23,7 @@
 import org.elasticsearch.xpack.esql.AssertWarnings;
 import org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase;
 import org.hamcrest.Matcher;
+import org.junit.Before;
 import org.junit.ClassRule;
 
 import java.io.IOException;
@@ -51,10 +50,9 @@
  * Tests for pushing queries to lucene.
  */
 @ThreadLeakFilters(filters = TestClustersThreadFilter.class)
-@TimeoutSuite(millis = 10 * TimeUnits.MINUTE) // semantic_text can take a long, long time to start in CI
 public class PushQueriesIT extends ESRestTestCase {
     @ClassRule
-    public static ElasticsearchCluster cluster = Clusters.testCluster();
+    public static ElasticsearchCluster cluster = Clusters.testCluster(spec -> spec.plugin("inference-service-test"));
 
     @ParametersFactory(argumentFormatting = "%1s")
     public static List<Object[]> args() {
@@ -288,13 +286,30 @@ private void indexValue(String value) throws IOException {
                   "number_of_shards": 1
                 }
               }""";
-        if (false == "auto".equals(type)) {
-            json += """
+        json += switch (type) {
+            case "auto" -> "";
+            case "semantic_text" -> """
                 ,
                 "mappings": {
                   "properties": {
                     "test": {
-                      "type": "%type",
+                      "type": "semantic_text",
+                      "inference_id": "test",
+                      "fields": {
+                        "keyword": {
+                          "type": "keyword",
+                          "ignore_above": 256
+                        }
+                      }
+                    }
+                  }
+                }""";
+            default -> """
+                  ,
+                  "mappings": {
+                    "properties": {
+                      "test": {
+                        "type": "%type",
                         "fields": {
                           "keyword": {
                             "type": "keyword",
@@ -305,7 +320,7 @@ private void indexValue(String value) throws IOException {
                     }
                   }
                 }""".replace("%type", type);
-        }
+        };
         json += "}";
         createIndex.setJsonEntity(json);
         Response createResponse = client().performRequest(createIndex);
@@ -347,4 +362,28 @@ protected boolean preserveClusterUponCompletion() {
         // Preserve the cluser to speed up the semantic_text tests
         return true;
     }
+
+    private static boolean setupEmbeddings = false;
+
+    @Before
+    public void setUpTextEmbeddingInferenceEndpoint() throws IOException {
+        if (type.equals("semantic_text") == false || setupEmbeddings) {
+            return;
+        }
+        setupEmbeddings = true;
+        Request request = new Request("PUT", "_inference/text_embedding/test");
+        request.setJsonEntity("""
+                  {
+                   "service": "text_embedding_test_service",
+                   "service_settings": {
+                     "model": "my_model",
+                     "api_key": "abc64",
+                     "dimensions": 128
+                   },
+                   "task_settings": {
+                   }
+                 }
+            """);
+        adminClient().performRequest(request);
+    }
 }