From ee93b696bb45f6082493f2b03f12ef024af5c489 Mon Sep 17 00:00:00 2001 From: "opensearch-trigger-bot[bot]" <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com> Date: Tue, 3 Dec 2024 12:40:39 -0800 Subject: [PATCH] Retrieve remote model id from registration response in IT to avoid flaky (#3244) (#3249) Signed-off-by: zane-neo (cherry picked from commit 1d306713a226fe84ff7cd2326084edc1beb3f112) Co-authored-by: zane-neo --- .../java/org/opensearch/ml/rest/MLCommonsRestTestCase.java | 6 +----- 1 file changed, 1 insertion(+), 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 index 613e7fb553..c7fa379f39 100644 --- 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;