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/1617.json

@@ -0,0 +1,4 @@
+{
+  \"category\": \"bug\",
+  \"description\": \"Fixes a flaky test (`NeuralQueryEnricherProcessorIT`) by ensuring the test framework waits for a model to be fully deployed.\"
+}

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

@@ -12,6 +12,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 
 import org.apache.hc.core5.http.io.entity.EntityUtils;
 import org.junit.Before;
@@ -31,6 +32,7 @@
 import lombok.SneakyThrows;
 import org.opensearch.neuralsearch.stats.events.EventStatName;
 import org.opensearch.neuralsearch.stats.info.InfoStatName;
+import org.opensearch.ml.common.transport.model.MLModelState;
 
 public class NeuralQueryEnricherProcessorIT extends BaseNeuralSearchIT {
 
@@ -57,10 +59,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 +90,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 +108,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 +130,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 +146,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 +158,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

@@ -372,7 +372,7 @@ protected void waitForModelToBeReady(String modelId) throws Exception {
 
     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);
     }
 
     /**
@@ -385,6 +385,7 @@ protected String prepareModel() {
         String requestBody = Files.readString(Path.of(classLoader.getResource("processor/UploadModelRequestBody.json").toURI()));
         String modelId = registerModelGroupAndUploadModel(requestBody);
         loadModel(modelId);
+        waitForModelToBeReady(modelId);
         return modelId;
     }
 
@@ -400,13 +401,15 @@ protected String prepareSparseEncodingModel() {
         );
         String modelId = registerModelGroupAndUploadModel(requestBody);
         loadModel(modelId);
+        waitForModelToBeReady(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);
+        waitForModelToBeReady(modelId);
         return modelId;
     }