Fix for Flaky test for issue 384 (#559)

* Fix to flaky test 384 Signed-off-by: Varun Jain <[email protected]>
opensearch-project · Jan 30, 2024 · d298e2b · d298e2b
1 parent 51f9d6f
commit d298e2b
Show file tree

Hide file tree

Showing 12 changed files with 1,160 additions and 1,102 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -23,6 +23,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Fixing multiple issues reported in #497 ([#524](https://github.com/opensearch-project/neural-search/pull/524))
 - Fix Flaky test reported in #433 ([#533](https://github.com/opensearch-project/neural-search/pull/533))
 - Enable support for default model id on HybridQueryBuilder ([#541](https://github.com/opensearch-project/neural-search/pull/541))
+- Fix Flaky test reported in #384 ([#559](https://github.com/opensearch-project/neural-search/pull/559))
 ### Infrastructure
 - BWC tests for Neural Search ([#515](https://github.com/opensearch-project/neural-search/pull/515))
 - Github action to run integ tests in secure opensearch cluster ([#535](https://github.com/opensearch-project/neural-search/pull/535))

diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java
@@ -11,7 +11,6 @@
 import java.util.Collections;
 import java.util.Map;
 
-import org.junit.After;
 import org.junit.Before;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.neuralsearch.BaseNeuralSearchIT;
@@ -34,68 +33,65 @@ public class NeuralQueryEnricherProcessorIT extends BaseNeuralSearchIT {
     public void setUp() throws Exception {
         super.setUp();
         updateClusterSettings();
-        prepareModel();
-    }
-
-    @After
-    @SneakyThrows
-    public void tearDown() {
-        super.tearDown();
-        deleteSearchPipeline(search_pipeline);
-        findDeployedModels().forEach(this::deleteModel);
-        deleteIndex(index);
     }
 
     @SneakyThrows
     public void testNeuralQueryEnricherProcessor_whenNoModelIdPassed_thenSuccess() {
-        initializeIndexIfNotExist();
-        String modelId = getDeployedModelId();
-        createSearchRequestProcessor(modelId, search_pipeline);
-        createPipelineProcessor(modelId, ingest_pipeline, ProcessorType.TEXT_EMBEDDING);
-        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());
-
+        String modelId = null;
+        try {
+            initializeIndexIfNotExist(index);
+            modelId = prepareModel();
+            createSearchRequestProcessor(modelId, search_pipeline);
+            createPipelineProcessor(modelId, ingest_pipeline, ProcessorType.TEXT_EMBEDDING);
+            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());
+        } finally {
+            wipeOfTestResources(index, ingest_pipeline, modelId, search_pipeline);
+        }
     }
 
     @SneakyThrows
     public void testNeuralQueryEnricherProcessor_whenHybridQueryBuilderAndNoModelIdPassed_thenSuccess() {
-        initializeIndexIfNotExist();
-        String modelId = getDeployedModelId();
-        createSearchRequestProcessor(modelId, search_pipeline);
-        createPipelineProcessor(modelId, ingest_pipeline, ProcessorType.TEXT_EMBEDDING);
-        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);
-        HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder();
-        hybridQueryBuilder.add(neuralQueryBuilder);
-        Map<String, Object> response = search(index, hybridQueryBuilder, 2);
-
-        assertFalse(response.isEmpty());
-
+        String modelId = null;
+        try {
+            initializeIndexIfNotExist(index);
+            modelId = prepareModel();
+            createSearchRequestProcessor(modelId, search_pipeline);
+            createPipelineProcessor(modelId, ingest_pipeline, ProcessorType.TEXT_EMBEDDING);
+            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);
+            HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder();
+            hybridQueryBuilder.add(neuralQueryBuilder);
+            Map<String, Object> response = search(index, hybridQueryBuilder, 2);
+
+            assertFalse(response.isEmpty());
+        } finally {
+            wipeOfTestResources(index, ingest_pipeline, modelId, search_pipeline);
+        }
     }
 
     @SneakyThrows
-    private void initializeIndexIfNotExist() {
-        if (index.equals(NeuralQueryEnricherProcessorIT.index) && !indexExists(index)) {
+    private void initializeIndexIfNotExist(String indexName) {
+        if (indexName.equals(NeuralQueryEnricherProcessorIT.index) && !indexExists(indexName)) {
             prepareKnnIndex(
-                index,
+                indexName,
                 Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE))
             );
             addKnnDoc(
-                index,
+                indexName,
                 "1",
                 Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1),
                 Collections.singletonList(Floats.asList(testVector).toArray())
             );
-            assertEquals(1, getDocCount(index));
+            assertEquals(1, getDocCount(indexName));
         }
     }
 }