From 30660828de40062768ccc1b8dbe4cfef02a7a9ee Mon Sep 17 00:00:00 2001
From: will-hwang <sang7239@gmail.com>
Date: Tue, 7 Jan 2025 22:04:52 -0800
Subject: [PATCH] add retry

---
 .../neuralsearch/bwc/MultiModalSearchIT.java  | 29 ++++++++++++-------
 .../neuralsearch/BaseNeuralSearchIT.java      |  4 +--
 2 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/MultiModalSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/MultiModalSearchIT.java
index a84bf6990..d109f25c4 100644
--- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/MultiModalSearchIT.java
+++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/MultiModalSearchIT.java
@@ -10,6 +10,8 @@
 import static org.opensearch.neuralsearch.util.TestUtils.NODES_BWC_CLUSTER;
 import static org.opensearch.neuralsearch.util.TestUtils.TEXT_IMAGE_EMBEDDING_PROCESSOR;
 import static org.opensearch.neuralsearch.util.TestUtils.getModelId;
+
+import org.opensearch.client.Response;
 import org.opensearch.neuralsearch.query.NeuralQueryBuilder;
 
 public class MultiModalSearchIT extends AbstractRollingUpgradeTestCase {
@@ -24,15 +26,23 @@ public class MultiModalSearchIT extends AbstractRollingUpgradeTestCase {
     private static final String TEST_IMAGE_TEXT_UPGRADED = "/9j/4AAQSkZJR8eydhgfwceocvlk";
 
     private static final int NUM_DOCS_PER_ROUND = 1;
+    private static String modelId = "";
 
     // Test rolling-upgrade test image embedding processor
     // Create Text Image Embedding Processor, Ingestion Pipeline and add document
     // Validate process , pipeline and document count in rolling-upgrade scenario
     public void testTextImageEmbeddingProcessor_E2EFlow() throws Exception {
-        waitForClusterHealthGreen(NODES_BWC_CLUSTER);
+        int count = 0;
+        int maxCount = 3;
+        Response response;
+        do {
+            response = waitForClusterHealthGreen(NODES_BWC_CLUSTER);
+            count++;
+        } while (count < maxCount && response.getStatusLine().getStatusCode() == 200 );
+
         switch (getClusterType()) {
             case OLD:
-                String modelId = uploadTextImageEmbeddingModel();
+                modelId = uploadTextImageEmbeddingModel();
                 loadModel(modelId);
                 createPipelineForTextImageProcessor(modelId, PIPELINE_NAME);
                 createIndexWithConfiguration(
@@ -43,27 +53,26 @@ public void testTextImageEmbeddingProcessor_E2EFlow() throws Exception {
                 addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT);
                 break;
             case MIXED:
-                String modelId2 = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR);
+                modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR);
                 int totalDocsCountMixed;
                 if (isFirstMixedRound()) {
                     totalDocsCountMixed = NUM_DOCS_PER_ROUND;
-                    validateTestIndexOnUpgrade(totalDocsCountMixed, modelId2, TEXT, TEST_IMAGE_TEXT);
+                    validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, TEXT, TEST_IMAGE_TEXT);
                     addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_MIXED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_MIXED);
                 } else {
                     totalDocsCountMixed = 2 * NUM_DOCS_PER_ROUND;
-                    validateTestIndexOnUpgrade(totalDocsCountMixed, modelId2, TEXT_MIXED, TEST_IMAGE_TEXT_MIXED);
+                    validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, TEXT_MIXED, TEST_IMAGE_TEXT_MIXED);
                 }
                 break;
             case UPGRADED:
-                String modelId3 = null;
                 try {
-                    modelId3 = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR);
+                    modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR);
                     int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND;
-                    loadModel(modelId3);
+                    loadModel(modelId);
                     addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_UPGRADED);
-                    validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId3, TEXT_UPGRADED, TEST_IMAGE_TEXT_UPGRADED);
+                    validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, TEXT_UPGRADED, TEST_IMAGE_TEXT_UPGRADED);
                 } finally {
-                    wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId3, null);
+                    wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null);
                 }
                 break;
             default:
diff --git a/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java b/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java
index 454c10fb2..e96a1e985 100644
--- a/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java
+++ b/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java
@@ -1280,13 +1280,13 @@ protected String registerModelGroup(final String modelGroupRegisterRequestBody)
     }
 
     // Method that waits till the health of nodes in the cluster goes green
-    protected void waitForClusterHealthGreen(final String numOfNodes) throws IOException {
+    protected Response waitForClusterHealthGreen(final String numOfNodes) throws IOException {
         Request waitForGreen = new Request("GET", "/_cluster/health");
         waitForGreen.addParameter("wait_for_nodes", numOfNodes);
         waitForGreen.addParameter("wait_for_status", "green");
         waitForGreen.addParameter("cluster_manager_timeout", "60s");
         waitForGreen.addParameter("timeout", "60s");
-        client().performRequest(waitForGreen);
+        return client().performRequest(waitForGreen);
     }
 
     /**