add retry for creating connector

Signed-off-by: Hailong Cui <[email protected]>
Hailong-am · Jun 11, 2024 · 4aac076 · 4aac076
1 parent bcaee80
commit 4aac076
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 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
@@ -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
@@ -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"