add ut it

Signed-off-by: zhichao-aws <[email protected]>
opensearch-project · Feb 28, 2024 · 168d26c · 168d26c
1 parent 3275439
commit 168d26c
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 0 deletions.
diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java
@@ -9,6 +9,7 @@
 import static org.opensearch.neuralsearch.TestUtils.createRandomVector;
 
 import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 
 import org.apache.hc.core5.http.io.entity.EntityUtils;
@@ -26,13 +27,16 @@
 import com.google.common.primitives.Floats;
 
 import lombok.SneakyThrows;
+import org.opensearch.neuralsearch.query.NeuralSparseQueryBuilder;
 
 public class NeuralQueryEnricherProcessorIT extends BaseNeuralSearchIT {
 
     private static final String index = "my-nlp-index";
+    private static final String sparseIndex = "my-nlp-index-sparse";
     private static final String search_pipeline = "search-pipeline";
     private static final String ingest_pipeline = "nlp-pipeline";
     private static final String TEST_KNN_VECTOR_FIELD_NAME_1 = "test-knn-vector-1";
+    private static final String TEST_RANK_FEATURES_FIELD_NAME_1 = "test-rank-features-1";
     private final float[] testVector = createRandomVector(TEST_DIMENSION);
 
     @Before
@@ -61,6 +65,25 @@ public void testNeuralQueryEnricherProcessor_whenNoModelIdPassed_thenSuccess() {
         }
     }
 
+    @SneakyThrows
+    public void testNeuralQueryEnricherProcessor_whenNoModelIdPassedInNeuralSparseQuery_thenSuccess() {
+        String modelId = null;
+        try {
+            initializeIndexIfNotExist(sparseIndex);
+            modelId = prepareSparseEncodingModel();
+            createSearchRequestProcessor(modelId, search_pipeline);
+            createPipelineProcessor(modelId, ingest_pipeline, ProcessorType.TEXT_EMBEDDING);
+            updateIndexSettings(sparseIndex, Settings.builder().put("index.search.default_pipeline", search_pipeline));
+            NeuralSparseQueryBuilder neuralSparseQueryBuilder = new NeuralSparseQueryBuilder();
+            neuralSparseQueryBuilder.fieldName(TEST_RANK_FEATURES_FIELD_NAME_1);
+            neuralSparseQueryBuilder.queryText("hello");
+            Map<String, Object> response = search(sparseIndex, neuralSparseQueryBuilder, 2);
+            assertFalse(response.isEmpty());
+        } finally {
+            wipeOfTestResources(sparseIndex, ingest_pipeline, modelId, search_pipeline);
+        }
+    }
+
     @SneakyThrows
     public void testNeuralQueryEnricherProcessor_whenGetEmptyQueryBody_thenSuccess() {
         String modelId = null;
@@ -120,5 +143,11 @@ private void initializeIndexIfNotExist(String indexName) {
             );
             assertEquals(1, getDocCount(indexName));
         }
+
+        if (sparseIndex.equals(indexName) && !indexExists(indexName)) {
+            prepareSparseEncodingIndex(indexName, List.of(TEST_RANK_FEATURES_FIELD_NAME_1));
+            addSparseEncodingDoc(indexName, "1", List.of(TEST_RANK_FEATURES_FIELD_NAME_1), List.of(Map.of("hi", 1.0f, "hello", 1.1f)));
+            assertEquals(1, getDocCount(indexName));
+        }
     }
 }
diff --git a/src/test/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitorTests.java b/src/test/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitorTests.java
@@ -39,6 +39,19 @@ public void testAccept_whenNeuralQueryBuilderWithoutFieldModelId_thenSetFieldMod
         assertEquals("bdcvjkcdjvkddcjxdjsc", neuralQueryBuilder.modelId());
     }
 
+    public void testAccept_whenNeuralSparseQueryBuilderWithModelId_theDoNothing() {
+        String modelId1 = "bdcvjkcdjvkddcjxdjsc";
+        String modelId2 = "45dfsnfoiqwrjcjxdjsc";
+        NeuralSparseQueryBuilder neuralSparseQueryBuilder = new NeuralSparseQueryBuilder();
+        neuralSparseQueryBuilder.fieldName("passage_text");
+        neuralSparseQueryBuilder.modelId(modelId1);
+
+        NeuralSearchQueryVisitor neuralSearchQueryVisitor = new NeuralSearchQueryVisitor(modelId2, null);
+        neuralSearchQueryVisitor.accept(neuralSparseQueryBuilder);
+
+        assertEquals(modelId1, neuralSparseQueryBuilder.modelId());
+    }
+
     public void testAccept_whenNeuralSparseQueryBuilderWithoutModelId_thenSetModelId() {
         String modelId = "bdcvjkcdjvkddcjxdjsc";
         NeuralSparseQueryBuilder neuralSparseQueryBuilder = new NeuralSparseQueryBuilder();