From 1d306713a226fe84ff7cd2326084edc1beb3f112 Mon Sep 17 00:00:00 2001 From: zane-neo Date: Tue, 3 Dec 2024 17:25:18 +0800 Subject: [PATCH] Retrieve remote model id from registration response in IT to avoid flaky (#3244) Signed-off-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 2092b9f4b4..6f710ea1de 100644 --- a/plugin/src/test/java/org/opensearch/ml/rest/MLCommonsRestTestCase.java +++ b/plugin/src/test/java/org/opensearch/ml/rest/MLCommonsRestTestCase.java @@ -1002,15 +1002,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;