Retrieve remote model id from registration response in IT to avoid fl…

…aky (#3244) (#3249) Signed-off-by: zane-neo <[email protected]> (cherry picked from commit 1d30671) Co-authored-by: zane-neo <[email protected]>
opensearch-project · Dec 3, 2024 · ee93b69 · ee93b69
1 parent de301b7
commit ee93b69
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/plugin/src/test/java/org/opensearch/ml/rest/MLCommonsRestTestCase.java b/plugin/src/test/java/org/opensearch/ml/rest/MLCommonsRestTestCase.java
@@ -987,15 +987,11 @@ public String registerRemoteModel(String createConnectorInput, String modelName,
         String connectorId = (String) responseMap.get("connector_id");
         response = RestMLRemoteInferenceIT.registerRemoteModel(modelName, modelName, connectorId);
         responseMap = parseResponseToMap(response);
-        String taskId = (String) responseMap.get("task_id");
-        waitForTask(taskId, MLTaskState.COMPLETED);
-        response = RestMLRemoteInferenceIT.getTask(taskId);
-        responseMap = parseResponseToMap(response);
         String modelId = (String) responseMap.get("model_id");
         if (deploy) {
             response = RestMLRemoteInferenceIT.deployRemoteModel(modelId);
             responseMap = parseResponseToMap(response);
-            taskId = (String) responseMap.get("task_id");
+            String taskId = (String) responseMap.get("task_id");
             waitForTask(taskId, MLTaskState.COMPLETED);
         }
         return modelId;