Add to inference service and crud IT rerank tests

Author: Evgenii-Kazannik

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceBaseRestTest.java

@@ -171,6 +171,20 @@ static String mockDenseServiceModelConfig() {
             """;
     }
 
+    static String mockRerankServiceModelConfig() {
+        return """
+            {
+              "service": "test_reranking_service",
+              "service_settings": {
+                 "model_id": "my_model",
+                 "api_key": "abc64"
+              },
+              "task_settings": {
+              }
+            }
+            """;
+    }
+
     static void deleteModel(String modelId) throws IOException {
         var request = new Request("DELETE", "_inference/" + modelId);
         var response = client().performRequest(request);

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceCrudIT.java

@@ -53,9 +53,12 @@ public void testCRUD() throws IOException {
         for (int i = 0; i < 4; i++) {
             putModel("te_model_" + i, mockDenseServiceModelConfig(), TaskType.TEXT_EMBEDDING);
         }
+        for (int i = 0; i < 3; i++) {
+            putModel("re-model-" + i, mockRerankServiceModelConfig(), TaskType.RERANK);
+        }
 
         var getAllModels = getAllModels();
-        int numModels = 12;
+        int numModels = 15;
         assertThat(getAllModels, hasSize(numModels));
 
         var getSparseModels = getModels("_all", TaskType.SPARSE_EMBEDDING);
@@ -71,6 +74,13 @@ public void testCRUD() throws IOException {
         for (var denseModel : getDenseModels) {
             assertEquals("text_embedding", denseModel.get("task_type"));
         }
+
+        var getRerankModels = getModels("_all", TaskType.RERANK);
+        int numRerankModels = 4;
+        assertThat(getRerankModels, hasSize(numRerankModels));
+        for (var denseModel : getRerankModels) {
+            assertEquals("rerank", denseModel.get("task_type"));
+        }
         String oldApiKey;
         {
             var singleModel = getModels("se_model_1", TaskType.SPARSE_EMBEDDING);
@@ -100,6 +110,9 @@ public void testCRUD() throws IOException {
         for (int i = 0; i < 4; i++) {
             deleteModel("te_model_" + i, TaskType.TEXT_EMBEDDING);
         }
+        for (int i = 0; i < 3; i++) {
+            deleteModel("re-model-" + i, TaskType.RERANK);
+        }
     }
 
     public void testGetModelWithWrongTaskType() throws IOException {

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/MockRerankInferenceServiceIT.java

@@ -0,0 +1,67 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference;
+
+import org.elasticsearch.inference.TaskType;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+
+public class MockRerankInferenceServiceIT extends InferenceBaseRestTest {
+
+    @SuppressWarnings("unchecked")
+    public void testMockService() throws IOException {
+        String inferenceEntityId = "test-mock";
+        var putModel = putModel(inferenceEntityId, mockRerankServiceModelConfig(), TaskType.RERANK);
+        var model = getModels(inferenceEntityId, TaskType.RERANK).get(0);
+
+        for (var modelMap : List.of(putModel, model)) {
+            assertEquals(inferenceEntityId, modelMap.get("inference_id"));
+            assertEquals(TaskType.RERANK, TaskType.fromString((String) modelMap.get("task_type")));
+            assertEquals("test_reranking_service", modelMap.get("service"));
+        }
+
+        List<String> input = List.of(randomAlphaOfLength(10));
+        var inference = infer(inferenceEntityId, input);
+        assertNonEmptyInferenceResults(inference, 1, TaskType.RERANK);
+        assertEquals(inference, infer(inferenceEntityId, input));
+        assertNotEquals(inference, infer(inferenceEntityId, randomValueOtherThan(input, () -> List.of(randomAlphaOfLength(10)))));
+    }
+
+    public void testMockServiceWithMultipleInputs() throws IOException {
+        String inferenceEntityId = "test-mock-with-multi-inputs";
+        putModel(inferenceEntityId, mockRerankServiceModelConfig(), TaskType.RERANK);
+        var queryParams = Map.of("timeout", "120s");
+
+        var inference = infer(
+            inferenceEntityId,
+            TaskType.RERANK,
+            List.of(randomAlphaOfLength(5), randomAlphaOfLength(10)),
+            "What if?",
+            queryParams
+        );
+
+        assertNonEmptyInferenceResults(inference, 2, TaskType.RERANK);
+    }
+
+    @SuppressWarnings("unchecked")
+    public void testMockService_DoesNotReturnSecretsInGetResponse() throws IOException {
+        String inferenceEntityId = "test-mock";
+        var putModel = putModel(inferenceEntityId, mockRerankServiceModelConfig(), TaskType.RERANK);
+        var model = getModels(inferenceEntityId, TaskType.RERANK).get(0);
+
+        var serviceSettings = (Map<String, Object>) model.get("service_settings");
+        assertNull(serviceSettings.get("api_key"));
+        assertNotNull(serviceSettings.get("model_id"));
+
+        var putServiceSettings = (Map<String, Object>) putModel.get("service_settings");
+        assertNull(putServiceSettings.get("api_key"));
+        assertNotNull(putServiceSettings.get("model_id"));
+    }
+}