From aa09014abdd0a890c9ea3a55bf472ecc7eb7e480 Mon Sep 17 00:00:00 2001 From: Hailong Cui Date: Fri, 10 May 2024 07:53:13 +0800 Subject: [PATCH] Fix flaky test of visualization tool (#2416) * add wait for model to be undeployed Signed-off-by: Hailong Cui * spotless Signed-off-by: Hailong Cui * update model undeploy status Signed-off-by: Hailong Cui * reduce total wait time Signed-off-by: Hailong Cui --------- Signed-off-by: Hailong Cui --- .../ml/tools/ToolIntegrationWithLLMTest.java | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/plugin/src/test/java/org/opensearch/ml/tools/ToolIntegrationWithLLMTest.java b/plugin/src/test/java/org/opensearch/ml/tools/ToolIntegrationWithLLMTest.java index 0caf7e964c..fe033645d6 100644 --- a/plugin/src/test/java/org/opensearch/ml/tools/ToolIntegrationWithLLMTest.java +++ b/plugin/src/test/java/org/opensearch/ml/tools/ToolIntegrationWithLLMTest.java @@ -8,8 +8,11 @@ import java.io.IOException; import java.util.List; import java.util.Locale; +import java.util.Map; +import java.util.Set; import java.util.UUID; import java.util.concurrent.TimeUnit; +import java.util.function.Predicate; import org.apache.hc.core5.http.ParseException; import org.junit.After; @@ -17,15 +20,21 @@ import org.opensearch.client.Response; import org.opensearch.core.rest.RestStatus; import org.opensearch.ml.common.MLModel; +import org.opensearch.ml.common.model.MLModelState; import org.opensearch.ml.rest.RestBaseAgentToolsIT; import org.opensearch.ml.utils.TestHelper; import com.sun.net.httpserver.HttpServer; +import lombok.SneakyThrows; import lombok.extern.log4j.Log4j2; @Log4j2 public abstract class ToolIntegrationWithLLMTest extends RestBaseAgentToolsIT { + + private static final int MAX_TASK_RESULT_QUERY_TIME_IN_SECOND = 30; + private static final int DEFAULT_TASK_RESULT_QUERY_INTERVAL_IN_MILLISECOND = 1000; + protected HttpServer server; protected String modelId; protected String agentId; @@ -63,9 +72,39 @@ public void stopMockLLM() { @After public void deleteModel() throws IOException { undeployModel(modelId); + waitModelUndeployed(modelId); deleteModel(client(), modelId, null); } + @SneakyThrows + private void waitModelUndeployed(String modelId) { + Predicate condition = response -> { + try { + Map responseInMap = parseResponseToMap(response); + MLModelState state = MLModelState.from(responseInMap.get(MLModel.MODEL_STATE_FIELD).toString()); + return Set.of(MLModelState.UNDEPLOYED, MLModelState.DEPLOY_FAILED).contains(state); + } catch (Exception e) { + return false; + } + }; + waitResponseMeetingCondition("GET", "/_plugins/_ml/models/" + modelId, null, condition); + } + + @SneakyThrows + protected Response waitResponseMeetingCondition(String method, String endpoint, String jsonEntity, Predicate condition) { + for (int i = 0; i < MAX_TASK_RESULT_QUERY_TIME_IN_SECOND; i++) { + Response response = TestHelper.makeRequest(client(), method, endpoint, null, jsonEntity, null); + assertEquals(RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + if (condition.test(response)) { + return response; + } + logger.info("The {}-th response: {}", i, response.toString()); + Thread.sleep(DEFAULT_TASK_RESULT_QUERY_INTERVAL_IN_MILLISECOND); + } + fail("The response failed to meet condition after " + MAX_TASK_RESULT_QUERY_TIME_IN_SECOND + " seconds."); + return null; + } + private String setUpConnectorWithRetry(int maxRetryTimes) throws InterruptedException { int retryTimes = 0; String connectorId = null;