fix(test): Stabilize flaky test by awaiting model deployment

The test testNeuralQueryEnricherProcessor_whenNoModelIdPassed_statsEnabled_thenSuccess
was failing intermittently due to a race condition where the search query
was executed before the model finished deploying.

This change adds a waitUntil block to poll the model's status and ensure
it is in a DEPLOYED state before the test proceeds. This same fix was
applied to other tests in the file to prevent them from failing in the
future.

Fixes #1599

Signed-off-by: Umarov Ismoiljon <[email protected]>

Author: umaarov

CHANGELOG.md

@@ -10,6 +10,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Enhancements
 
 ### Bug Fixes
+- Fix flaky test in `NeuralQueryEnricherProcessorIT` by waiting for model deployment ([#1617](https://github.com/opensearch-project/neural-search/pull/1617))
 
 ### Infrastructure
 

src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java

@@ -57,10 +57,10 @@ public void testNeuralQueryEnricherProcessor_whenNoModelIdPassed_thenSuccess() {
         createPipelineProcessor(modelId, ingest_pipeline, ProcessorType.TEXT_EMBEDDING);
         updateIndexSettings(index, Settings.builder().put("index.search.default_pipeline", search_pipeline));
         NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder()
-            .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1)
-            .queryText("Hello World")
-            .k(1)
-            .build();
+                .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1)
+                .queryText("Hello World")
+                .k(1)
+                .build();
         Map<String, Object> response = search(index, neuralQueryBuilder, 2);
         assertFalse(response.isEmpty());
     }
@@ -88,9 +88,11 @@ public void testNeuralQueryEnricherProcessor_whenGetEmptyQueryBody_thenSuccess()
         updateIndexSettings(index, Settings.builder().put("index.search.default_pipeline", search_pipeline));
         Request request = new Request("POST", "/" + index + "/_search");
         Response response = client().performRequest(request);
-        assertEquals(request.getEndpoint() + ": failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode()));
+        assertEquals(request.getEndpoint() + ": failed", RestStatus.OK,
+                RestStatus.fromCode(response.getStatusLine().getStatusCode()));
         String responseBody = EntityUtils.toString(response.getEntity());
-        Map<String, Object> responseInMap = XContentHelper.convertToMap(XContentType.JSON.xContent(), responseBody, false);
+        Map<String, Object> responseInMap = XContentHelper.convertToMap(XContentType.JSON.xContent(), responseBody,
+                false);
         assertFalse(responseInMap.isEmpty());
         assertEquals(3, ((Map) responseInMap.get("hits")).size());
     }
@@ -104,10 +106,10 @@ public void testNeuralQueryEnricherProcessor_whenHybridQueryBuilderAndNoModelIdP
         createPipelineProcessor(modelId, ingest_pipeline, ProcessorType.TEXT_EMBEDDING);
         updateIndexSettings(index, Settings.builder().put("index.search.default_pipeline", search_pipeline));
         NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder()
-            .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1)
-            .queryText("Hello World")
-            .k(1)
-            .build();
+                .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1)
+                .queryText("Hello World")
+                .k(1)
+                .build();
         HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder();
         hybridQueryBuilder.add(neuralQueryBuilder);
         Map<String, Object> response = search(index, hybridQueryBuilder, 2);
@@ -126,10 +128,10 @@ public void testNeuralQueryEnricherProcessor_whenNoModelIdPassed_statsEnabled_th
         createPipelineProcessor(modelId, ingest_pipeline, ProcessorType.TEXT_EMBEDDING);
         updateIndexSettings(index, Settings.builder().put("index.search.default_pipeline", search_pipeline));
         NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder()
-            .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1)
-            .queryText("Hello World")
-            .k(1)
-            .build();
+                .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1)
+                .queryText("Hello World")
+                .k(1)
+                .build();
         Map<String, Object> response = search(index, neuralQueryBuilder, 2);
         assertFalse(response.isEmpty());
 
@@ -142,7 +144,8 @@ public void testNeuralQueryEnricherProcessor_whenNoModelIdPassed_statsEnabled_th
         Map<String, Object> allNodesStats = parseAggregatedNodeStatsResponse(responseBody);
 
         // Parse json to get stats
-        assertEquals(2, getNestedValue(allNodesStats, EventStatName.NEURAL_QUERY_ENRICHER_PROCESSOR_EXECUTIONS.getFullPath()));
+        assertEquals(2,
+                getNestedValue(allNodesStats, EventStatName.NEURAL_QUERY_ENRICHER_PROCESSOR_EXECUTIONS.getFullPath()));
         assertEquals(1, getNestedValue(stats, InfoStatName.NEURAL_QUERY_ENRICHER_PROCESSORS.getFullPath()));
 
         // Reset stats
@@ -153,21 +156,21 @@ public void testNeuralQueryEnricherProcessor_whenNoModelIdPassed_statsEnabled_th
     private void initializeIndexIfNotExist(String indexName) {
         if (indexName.equals(NeuralQueryEnricherProcessorIT.index) && !indexExists(indexName)) {
             prepareKnnIndex(
-                indexName,
-                Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE))
-            );
+                    indexName,
+                    Collections.singletonList(
+                            new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)));
             addKnnDoc(
-                indexName,
-                "1",
-                Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1),
-                Collections.singletonList(Floats.asList(testVector).toArray())
-            );
+                    indexName,
+                    "1",
+                    Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1),
+                    Collections.singletonList(Floats.asList(testVector).toArray()));
             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)));
+            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));
         }
     }

src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java

@@ -320,8 +320,7 @@ private void doLoadAndWaitForModelToBeReady(String modelId) throws Exception {
             logger.info("Model is already deployed. Skip loading.");
             return;
         }
-        loadModel(modelId);
-        waitForModelToBeReady(modelId);
+        loadAndWaitForModelToBeReady(modelId);
     }
 
     protected void loadAndWaitForModelToBeReady(String modelId) throws Exception {
@@ -370,9 +369,12 @@ protected void waitForModelToBeReady(String modelId) throws Exception {
         throw new RuntimeException("Model " + modelId + " failed to be ready for inference after " + MAX_ATTEMPTS + " attempts");
     }
 
+    /**
+     * LOADED is same as DEPLOYED but it is deprecated at OS2.7. As Neural Search is introduced after OS2.7, we could exclude that state.
+     */
     protected boolean isModelReadyForInference(@NonNull final String modelId) throws IOException, ParseException {
         MLModelState state = getModelState(modelId);
-        return MLModelState.LOADED.equals(state) || MLModelState.DEPLOYED.equals(state);
+        return MLModelState.DEPLOYED.equals(state);
     }
 
     /**
@@ -384,7 +386,7 @@ protected boolean isModelReadyForInference(@NonNull final String modelId) throws
     protected String prepareModel() {
         String requestBody = Files.readString(Path.of(classLoader.getResource("processor/UploadModelRequestBody.json").toURI()));
         String modelId = registerModelGroupAndUploadModel(requestBody);
-        loadModel(modelId);
+        loadAndWaitForModelToBeReady(modelId);
         return modelId;
     }
 
@@ -399,14 +401,14 @@ protected String prepareSparseEncodingModel() {
             Path.of(classLoader.getResource("processor/UploadSparseEncodingModelRequestBody.json").toURI())
         );
         String modelId = registerModelGroupAndUploadModel(requestBody);
-        loadModel(modelId);
+        loadAndWaitForModelToBeReady(modelId);
         return modelId;
     }
 
     protected String prepareSemanticHighlightingLocalModel() throws Exception {
         String requestBody = Files.readString(Path.of(classLoader.getResource("highlight/LocalQuestionAnsweringModel.json").toURI()));
         String modelId = registerModelGroupAndUploadModel(requestBody);
-        loadModel(modelId);
+        loadAndWaitForModelToBeReady(modelId);
         return modelId;
     }