From a85bcdf93bc6baf2bbb5c2cebcfb5048b3007f76 Mon Sep 17 00:00:00 2001 From: Hailong Cui Date: Wed, 8 May 2024 17:24:10 +0800 Subject: [PATCH] update model undeploy status Signed-off-by: Hailong Cui --- .../opensearch/ml/tools/ToolIntegrationWithLLMTest.java | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) 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 2c6bae22dd..aeac24ef3a 100644 --- a/plugin/src/test/java/org/opensearch/ml/tools/ToolIntegrationWithLLMTest.java +++ b/plugin/src/test/java/org/opensearch/ml/tools/ToolIntegrationWithLLMTest.java @@ -9,6 +9,7 @@ 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; @@ -80,11 +81,9 @@ private void waitModelUndeployed(String modelId) { Predicate condition = response -> { try { Map responseInMap = parseResponseToMap(response); - String state = responseInMap.get(MLModel.MODEL_STATE_FIELD).toString(); - return !state.equals(MLModelState.DEPLOYED.toString()) - && !state.equals(MLModelState.DEPLOYING.toString()) - && !state.equals(MLModelState.PARTIALLY_DEPLOYED.toString()); - } catch (IOException e) { + 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; } };