(iotdb) 01/01: AINode: add retry logic in model status check in AINode IT (#13774)

[ycycse]

This is an automated email from the ASF dual-hosted git repository.

ycycse pushed a commit to branch cp133
in repository https://gitbox.apache.org/repos/asf/iotdb.git

commit a87dcc205de40b013eeb38b30275b88cb6a6b7d4
Author: YangCaiyin <ycy...@gmail.com>
AuthorDate: Wed Oct 16 10:24:23 2024 +0800

    AINode: add retry logic in model status check in AINode IT (#13774)
    
    (cherry picked from commit 4558493e6b7e2039a05e466a2c50dc47c215a384)
---
 .../org/apache/iotdb/ainode/it/AINodeBasicIT.java  | 39 +++++++++++++---------
 1 file changed, 24 insertions(+), 15 deletions(-)

diff --git 
a/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeBasicIT.java 
b/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeBasicIT.java
index 387a7cc3fd8..17404a50e9d 100644
--- 
a/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeBasicIT.java
+++ 
b/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeBasicIT.java
@@ -135,27 +135,36 @@ public class AINodeBasicIT {
     try (Connection connection = EnvFactory.getEnv().getConnection();
         Statement statement = connection.createStatement()) {
       statement.execute(registerSql);
-      try (ResultSet resultSet = statement.executeQuery(showSql)) {
-        ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
-        checkHeader(resultSetMetaData, 
"ModelId,ModelType,State,Configs,Notes");
-        int count = 0;
-        while (resultSet.next()) {
-          String modelName = resultSet.getString(1);
-          String modelType = resultSet.getString(2);
-          String status = resultSet.getString(3);
-
-          assertEquals("operationTest", modelName);
-          assertEquals("USER_DEFINED", modelType);
-          assertEquals("ACTIVE", status);
-          count++;
+      boolean loading = true;
+      int count = 0;
+      while (loading) {
+        try (ResultSet resultSet = statement.executeQuery(showSql)) {
+          ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
+          checkHeader(resultSetMetaData, 
"ModelId,ModelType,State,Configs,Notes");
+          while (resultSet.next()) {
+            String modelName = resultSet.getString(1);
+            String modelType = resultSet.getString(2);
+            String status = resultSet.getString(3);
+
+            assertEquals("operationTest", modelName);
+            assertEquals("USER_DEFINED", modelType);
+            if (status.equals("ACTIVE")) {
+              loading = false;
+              count++;
+            } else if (status.equals("LOADING")) {
+              break;
+            } else {
+              fail("Unexpected status of model: " + status);
+            }
+          }
         }
-        assertEquals(1, count);
       }
+      assertEquals(1, count);
       statement.execute(dropSql);
       try (ResultSet resultSet = statement.executeQuery(showSql)) {
         ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
         checkHeader(resultSetMetaData, 
"ModelId,ModelType,State,Configs,Notes");
-        int count = 0;
+        count = 0;
         while (resultSet.next()) {
           count++;
         }