From f28bb740e87cd520190b003011303d5a85d477d6 Mon Sep 17 00:00:00 2001 From: Hailong Cui Date: Thu, 13 Jun 2024 06:30:57 +0800 Subject: [PATCH] fix flaky test of PredictionITTests and RestConnectorToolIT (#2437) * fix flaky test of prediction Signed-off-by: Hailong Cui * add retry for creating connector Signed-off-by: Hailong Cui * disable SyncUpJob only for PredictionITTests Signed-off-by: Hailong Cui --------- Signed-off-by: Hailong Cui --- .../opensearch/ml/action/MLCommonsIntegTestCase.java | 12 ++++++++++-- .../ml/action/prediction/PredictionITTests.java | 10 ++++++++++ .../opensearch/ml/rest/MLCommonsRestTestCase.java | 10 +++++++++- .../org/opensearch/ml/rest/RestConnectorToolIT.java | 6 ++++++ 4 files changed, 35 insertions(+), 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 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 937167f00c..f697aa6fd4 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; @@ -16,6 +17,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; @@ -49,6 +51,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 7321c27769..8e013d4305 100644 --- a/plugin/src/test/java/org/opensearch/ml/rest/MLCommonsRestTestCase.java +++ b/plugin/src/test/java/org/opensearch/ml/rest/MLCommonsRestTestCase.java @@ -970,10 +970,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 76a1c20e61..58abc5fb55 100644 --- a/plugin/src/test/java/org/opensearch/ml/rest/RestConnectorToolIT.java +++ b/plugin/src/test/java/org/opensearch/ml/rest/RestConnectorToolIT.java @@ -86,6 +86,9 @@ public void tearDown() throws Exception { } public void testConnectorToolInFlowAgent_WrongAction() throws IOException, ParseException { + 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" @@ -111,6 +114,9 @@ public void testConnectorToolInFlowAgent_WrongAction() throws IOException, Parse } public void testConnectorToolInFlowAgent() throws IOException, ParseException { + 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"