[Backport 2.x] fix flaky test of PredictionITTests and RestConnectorT…

…oolIT (#2437) (#2547) * fix flaky test of PredictionITTests and RestConnectorToolIT (#2437) Signed-off-by: Hailong Cui <ihailong@amazon.com> Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com> Co-authored-by: Hailong Cui <ihailong@amazon.com>
opensearch-project · Jun 12, 2024 · 2775942 · 2775942
1 parent 1bb4ec4
commit 2775942
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 3 deletions.
diff --git a/plugin/src/test/java/org/opensearch/ml/action/MLCommonsIntegTestCase.java b/plugin/src/test/java/org/opensearch/ml/action/MLCommonsIntegTestCase.java
@@ -93,15 +93,23 @@
 import org.opensearch.ml.utils.TestData;
 import org.opensearch.plugins.Plugin;
 import org.opensearch.search.builder.SearchSourceBuilder;
-import org.opensearch.test.OpenSearchIntegTestCase;
+import org.opensearch.test.ParameterizedStaticSettingsOpenSearchIntegTestCase;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.gson.Gson;
 
-public class MLCommonsIntegTestCase extends OpenSearchIntegTestCase {
+public class MLCommonsIntegTestCase extends ParameterizedStaticSettingsOpenSearchIntegTestCase {
     private Gson gson = new Gson();
 
+    public MLCommonsIntegTestCase() {
+        super(Settings.EMPTY);
+    }
+
+    public MLCommonsIntegTestCase(Settings nodeSettings) {
+        super(nodeSettings);
+    }
+
     @Override
     protected Collection<Class<? extends Plugin>> nodePlugins() {
         return Collections.singletonList(MachineLearningPlugin.class);

diff --git a/plugin/src/test/java/org/opensearch/ml/action/prediction/PredictionITTests.java b/plugin/src/test/java/org/opensearch/ml/action/prediction/PredictionITTests.java
@@ -5,6 +5,7 @@
 
 package org.opensearch.ml.action.prediction;
 
+import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_SYNC_UP_JOB_INTERVAL_IN_SECONDS;
 import static org.opensearch.ml.utils.TestData.IRIS_DATA_SIZE;
 import static org.opensearch.ml.utils.TestData.TIME_FIELD;
 
@@ -17,6 +18,7 @@
 import org.junit.rules.ExpectedException;
 import org.opensearch.action.ActionRequestValidationException;
 import org.opensearch.common.action.ActionFuture;
+import org.opensearch.common.settings.Settings;
 import org.opensearch.ml.action.MLCommonsIntegTestCase;
 import org.opensearch.ml.common.FunctionName;
 import org.opensearch.ml.common.MLModel;
@@ -50,6 +52,14 @@ public class PredictionITTests extends MLCommonsIntegTestCase {
     private String logisticRegressionModelId;
     private int batchRcfDataSize = 100;
 
+    /**
+     * set ML_COMMONS_SYNC_UP_JOB_INTERVAL_IN_SECONDS to 0 to disable ML_COMMONS_SYNC_UP_JOB
+     * the cluster will be pre-created with the settings at startup
+     */
+    public PredictionITTests() {
+        super(Settings.builder().put(ML_COMMONS_SYNC_UP_JOB_INTERVAL_IN_SECONDS.getKey(), 0).build());
+    }
+
     @Rule
     public ExpectedException exceptionRule = ExpectedException.none();
 

diff --git a/plugin/src/test/java/org/opensearch/ml/rest/MLCommonsRestTestCase.java b/plugin/src/test/java/org/opensearch/ml/rest/MLCommonsRestTestCase.java
@@ -955,10 +955,18 @@ public void waitForTask(String taskId, MLTaskState targetState) throws Interrupt
     }
 
     public String registerConnector(String createConnectorInput) throws IOException, InterruptedException {
-        Response response = RestMLRemoteInferenceIT.createConnector(createConnectorInput);
+        Response response;
+        try {
+            response = RestMLRemoteInferenceIT.createConnector(createConnectorInput);
+        } catch (Throwable throwable) {
+            // Add retry for `The ML encryption master key has not been initialized yet. Please retry after waiting for 10 seconds.`
+            TimeUnit.SECONDS.sleep(10);
+            response = RestMLRemoteInferenceIT.createConnector(createConnectorInput);
+        }
         Map responseMap = parseResponseToMap(response);
         String connectorId = (String) responseMap.get("connector_id");
         return connectorId;
+
     }
 
     public String registerRemoteModel(String createConnectorInput, String modelName, boolean deploy) throws IOException,

diff --git a/plugin/src/test/java/org/opensearch/ml/rest/RestConnectorToolIT.java b/plugin/src/test/java/org/opensearch/ml/rest/RestConnectorToolIT.java
@@ -85,6 +85,9 @@ public void tearDown() throws Exception {
     }
 
     public void testConnectorToolInFlowAgent_WrongAction() throws IOException {
+        if (AWS_ACCESS_KEY_ID == null || AWS_SECRET_ACCESS_KEY == null || AWS_SESSION_TOKEN == null) {
+            return;
+        }
         String registerAgentRequestBody = "{\n"
             + "  \"name\": \"Test agent with connector tool\",\n"
             + "  \"type\": \"flow\",\n"
@@ -110,6 +113,9 @@ public void testConnectorToolInFlowAgent_WrongAction() throws IOException {
     }
 
     public void testConnectorToolInFlowAgent() throws IOException {
+        if (AWS_ACCESS_KEY_ID == null || AWS_SECRET_ACCESS_KEY == null || AWS_SESSION_TOKEN == null) {
+            return;
+        }
         String registerAgentRequestBody = "{\n"
             + "  \"name\": \"Test agent with connector tool\",\n"
             + "  \"type\": \"flow\",\n"