From 4aac076323f3a5b13d4e95a5c3f983fd54ce7156 Mon Sep 17 00:00:00 2001 From: Hailong Cui Date: Tue, 11 Jun 2024 11:00:14 +0800 Subject: [PATCH] add retry for creating connector Signed-off-by: Hailong Cui --- .../org/opensearch/ml/rest/MLCommonsRestTestCase.java | 10 +++++++++- .../org/opensearch/ml/rest/RestConnectorToolIT.java | 6 ++++++ 2 files changed, 15 insertions(+), 1 deletion(-) 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 cf1f87e09e..678b3d65bf 100644 --- a/plugin/src/test/java/org/opensearch/ml/rest/MLCommonsRestTestCase.java +++ b/plugin/src/test/java/org/opensearch/ml/rest/MLCommonsRestTestCase.java @@ -962,10 +962,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 4ae9653d60..76feea2e18 100644 --- a/plugin/src/test/java/org/opensearch/ml/rest/RestConnectorToolIT.java +++ b/plugin/src/test/java/org/opensearch/ml/rest/RestConnectorToolIT.java @@ -84,6 +84,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" @@ -109,6 +112,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"