Address PR comments; Added CHANGELOG

Signed-off-by: Tanqiu Liu <[email protected]>
opensearch-project · Nov 8, 2023 · 371a9dd · 371a9dd
1 parent 8ab54c4
commit 371a9dd
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 5 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Features
 ### Enhancements
 ### Bug Fixes
+- Fixed flaky integration tests caused by model_state transition latency. 
 ### Infrastructure
 ### Documentation
 ### Maintenance

diff --git a/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java b/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java
@@ -193,6 +193,8 @@ protected void loadModel(String modelId) throws Exception {
             isComplete = checkComplete(taskQueryResult);
             Thread.sleep(DEFAULT_TASK_RESULT_QUERY_INTERVAL_IN_MILLISECOND);
         }
+        // wait for model state update to DEPLOYED
+        pollForModelState(modelId, Set.of(MLModelState.DEPLOYED), 3000, 5);
     }
 
     /**
@@ -612,7 +614,7 @@ protected void deleteModel(String modelId) {
         );
 
         // wait for model undeploy to complete
-        pollForModelState(modelId, MLModelState.UNDEPLOYED, 3000, 5);
+        pollForModelState(modelId, Set.of(MLModelState.UNDEPLOYED, MLModelState.DEPLOY_FAILED), 3000, 5);
 
         makeRequest(
             client(),
@@ -624,14 +626,26 @@ protected void deleteModel(String modelId) {
         );
     }
 
-    @SneakyThrows
-    protected void pollForModelState(String modelId, MLModelState expectedModelState, int intervalMs, int maxAttempts) {
+    protected void pollForModelState(String modelId, Set<MLModelState> exitModelStates, int intervalMs, int maxAttempts)
+        throws InterruptedException {
+        MLModelState currentState = null;
         for (int i = 0; i < maxAttempts; i++) {
             Thread.sleep(intervalMs);
-            if (expectedModelState.equals(getModelState(modelId))) {
+            currentState = getModelState(modelId);
+            if (exitModelStates.contains(currentState)) {
                 return;
             }
         }
+        fail(
+            String.format(
+                LOCALE,
+                "Model state does not reached exit states %s after %d attempts with interval of %d ms, latest model state: %s.",
+                StringUtils.join(exitModelStates, ","),
+                maxAttempts,
+                intervalMs,
+                currentState
+            )
+        );
     }
 
     public boolean isUpdateClusterSettings() {
@@ -747,7 +761,7 @@ protected Set<String> findDeployedModels() {
             .filter(
                 hitsMap -> !Objects.isNull(hitsMap)
                     && hitsMap.containsKey("model_id")
-                    && MLModelState.DEPLOYED.equals(getModelState(hitsMap.get("model_id").toString()))
+                    && getModelState(hitsMap.get("model_id").toString()) == MLModelState.DEPLOYED
             )
             .map(hitsMap -> (String) hitsMap.get("model_id"))
             .collect(Collectors.toSet());