From b13eb809770befd918af905a4900b8ce1b8db49b Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Tue, 26 Sep 2023 13:18:32 -0700 Subject: [PATCH] Integ test cases Signed-off-by: Varun Jain --- .../query/NeuralQueryBuilder.java | 4 +- .../common/BaseNeuralSearchIT.java | 26 ++++++ .../processor/NeuralQueryProcessorIT.java | 83 +++++++++++++++++++ .../resources/processor/IndexMappings.json | 7 ++ .../SearchRequestPipelineConfiguration.json | 2 +- 5 files changed, 120 insertions(+), 2 deletions(-) create mode 100644 src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorIT.java diff --git a/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java index 51263a068..728966428 100644 --- 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 index fac32b538..faafa6d59 100644 --- 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 new file mode 100644 index 000000000..9f845720e --- /dev/null +++ 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 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 index 6cc3fabbe..5464a9311 100644 --- 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 index 3f3675ce3..ad5a0ba52 100644 --- 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"