Integ test cases

Signed-off-by: Varun Jain <[email protected]>
opensearch-project · Sep 26, 2023 · b13eb80 · b13eb80
1 parent dee67ca
commit b13eb80
Show file tree

Hide file tree

Showing 5 changed files with 120 additions and 2 deletions.
diff --git a/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java
@@ -123,7 +123,9 @@ protected void doXContent(XContentBuilder xContentBuilder, Params params) throws
         xContentBuilder.startObject(NAME);
         xContentBuilder.startObject(fieldName);
         xContentBuilder.field(QUERY_TEXT_FIELD.getPreferredName(), queryText);
-        xContentBuilder.field(MODEL_ID_FIELD.getPreferredName(), modelId);
+        if (!isClusterOnOrAfterMinRequiredVersion() || (isClusterOnOrAfterMinRequiredVersion() && modelId != null)) {
+            xContentBuilder.field(MODEL_ID_FIELD.getPreferredName(), modelId);
+        }
         xContentBuilder.field(K_FIELD.getPreferredName(), k);
         if (filter != null) {
             xContentBuilder.field(FILTER_FIELD.getPreferredName(), filter);

diff --git a/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java b/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java
@@ -39,6 +39,8 @@
 import org.opensearch.client.Response;
 import org.opensearch.client.RestClient;
 import org.opensearch.client.WarningsHandler;
+import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.settings.Settings;
 import org.opensearch.common.xcontent.XContentFactory;
 import org.opensearch.common.xcontent.XContentHelper;
 import org.opensearch.common.xcontent.XContentType;
@@ -48,10 +50,16 @@
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.knn.index.SpaceType;
 import org.opensearch.neuralsearch.OpenSearchSecureRestTestCase;
+import org.opensearch.neuralsearch.util.NeuralSearchClusterUtil;
+import org.opensearch.test.ClusterServiceUtils;
+import org.opensearch.threadpool.TestThreadPool;
+import org.opensearch.threadpool.ThreadPool;
 
 import com.carrotsearch.randomizedtesting.RandomizedTest;
+import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope;
 import com.google.common.collect.ImmutableList;
 
+@ThreadLeakScope(ThreadLeakScope.Scope.NONE)
 public abstract class BaseNeuralSearchIT extends OpenSearchSecureRestTestCase {
 
     private static final Locale LOCALE = Locale.ROOT;
@@ -66,11 +74,29 @@ public abstract class BaseNeuralSearchIT extends OpenSearchSecureRestTestCase {
 
     protected final ClassLoader classLoader = this.getClass().getClassLoader();
 
+    protected ThreadPool threadPool;
+    protected ClusterService clusterService;
+
     @Before
     public void setupSettings() {
+        threadPool = setUpThreadPool();
+        clusterService = createClusterService(threadPool);
         if (isUpdateClusterSettings()) {
             updateClusterSettings();
         }
+        NeuralSearchClusterUtil.instance().initialize(clusterService);
+    }
+
+    protected ThreadPool setUpThreadPool() {
+        return new TestThreadPool(getClass().getName(), threadPoolSettings());
+    }
+
+    public Settings threadPoolSettings() {
+        return Settings.EMPTY;
+    }
+
+    public static ClusterService createClusterService(ThreadPool threadPool) {
+        return ClusterServiceUtils.createClusterService(threadPool);
     }
 
     protected void updateClusterSettings() {

diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorIT.java
@@ -0,0 +1,83 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.neuralsearch.processor;
+
+import static org.opensearch.neuralsearch.TestUtils.createRandomVector;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.Map;
+
+import lombok.SneakyThrows;
+
+import org.junit.After;
+import org.junit.Before;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.knn.index.SpaceType;
+import org.opensearch.neuralsearch.common.BaseNeuralSearchIT;
+import org.opensearch.neuralsearch.query.NeuralQueryBuilder;
+
+import com.google.common.primitives.Floats;
+
+public class NeuralQueryProcessorIT extends BaseNeuralSearchIT {
+
+    public static final String index = "my-nlp-index";
+    public static final String search_pipeline = "search-pipeline";
+    public static final String ingest_pipeline = "nlp-pipeline";
+    private static final String TEST_KNN_VECTOR_FIELD_NAME_1 = "test-knn-vector-1";
+    private static final int TEST_DIMENSION = 768;
+    private static final SpaceType TEST_SPACE_TYPE = SpaceType.L2;
+    private final float[] testVector = createRandomVector(TEST_DIMENSION);
+
+    @Before
+    public void setUp() throws Exception {
+        super.setUp();
+        updateClusterSettings();
+        prepareModel();
+    }
+
+    @After
+    @SneakyThrows
+    public void tearDown() {
+        super.tearDown();
+        deleteSearchPipeline(search_pipeline);
+        findDeployedModels().forEach(this::deleteModel);
+    }
+
+    public void testNeuralQueryProcessor() throws Exception {
+        initializeIndexIfNotExist(index);
+        String modelId = getDeployedModelId();
+        createSearchRequestProcessor(modelId, search_pipeline);
+        createPipelineProcessor(modelId, ingest_pipeline);
+        updateIndexSettings(index, Settings.builder().put("index.search.default_pipeline", search_pipeline));
+        NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder();
+        neuralQueryBuilder.fieldName(TEST_KNN_VECTOR_FIELD_NAME_1);
+        neuralQueryBuilder.queryText("Hello World");
+        neuralQueryBuilder.k(1);
+        Map<String, Object> response = search(index, neuralQueryBuilder, 2);
+
+        assertFalse(response.isEmpty());
+
+        assertEquals(modelId, neuralQueryBuilder.modelId());
+
+    }
+
+    private void initializeIndexIfNotExist(String indexName) throws IOException {
+        if (index.equals(indexName) && !indexExists(index)) {
+            prepareKnnIndex(
+                index,
+                Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE))
+            );
+            addKnnDoc(
+                index,
+                "1",
+                Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1),
+                Collections.singletonList(Floats.asList(testVector).toArray())
+            );
+            assertEquals(1, getDocCount(index));
+        }
+    }
+}
diff --git a/src/test/resources/processor/IndexMappings.json b/src/test/resources/processor/IndexMappings.json
@@ -67,6 +67,13 @@
             }
           }
         }
+      },
+      "passage_embedding": {
+        "type": "knn_vector",
+        "dimension": 768
+      },
+      "passage_text": {
+        "type": "text"
       }
     }
   }

diff --git a/src/test/resources/processor/SearchRequestPipelineConfiguration.json b/src/test/resources/processor/SearchRequestPipelineConfiguration.json
@@ -1,7 +1,7 @@
 {
   "request_processors": [
     {
-      "default_query": {
+      "neural_query": {
         "tag": "tag1",
         "description": "This processor is going to restrict to publicly visible documents",
         "default_model_id": "%s"