diff --git a/plugin/src/test/java/org/opensearch/ml/action/MLCommonsIntegTestCase.java b/plugin/src/test/java/org/opensearch/ml/action/MLCommonsIntegTestCase.java index c4fb3f906c..9e913f98d4 100644 --- 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> 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 index f9a911b2a6..7df3fff3d5 100644 --- 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 index 2a0f5db613..0d2dfa1fa6 100644 --- 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 index 7c8f6aab84..0f75be90ac 100644 --- 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"