[Backport 2.15] fix flaky test of PredictionITTests and RestConnector…

…ToolIT (#2437) (#2548) (#3039) * fix flaky test of PredictionITTests and RestConnectorToolIT (#2437) Signed-off-by: Bhavana Ramaram <[email protected]> Co-authored-by: Bhavana Ramaram <[email protected]> Co-authored-by: Hailong Cui <[email protected]>
opensearch-project · Oct 2, 2024 · fcbed3f · fcbed3f
1 parent a4d5b33
commit fcbed3f
Show file tree

Hide file tree

Showing 3 changed files with 29 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
@@ -964,10 +964,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,