From 6eeddb4b531f5b2d1ab32566fd14f7168b5a7254 Mon Sep 17 00:00:00 2001
From: Alex Duan <417921451@qq.com>
Date: Fri, 30 Dec 2022 15:19:44 +0800
Subject: [PATCH] add case feature function with loop check data

---
 tests/pytest/util/sql.py                    | 64 +++++++++++++++++++++
 tests/system-test/6-cluster/5dnode1mnode.py | 15 ++---
 2 files changed, 70 insertions(+), 9 deletions(-)

diff --git a/tests/pytest/util/sql.py b/tests/pytest/util/sql.py
index 9cfd1d368e..bdf3f20e15 100644
--- a/tests/pytest/util/sql.py
+++ b/tests/pytest/util/sql.py
@@ -261,6 +261,70 @@ class TDSql:
 
         tdLog.info(f"sql:{self.sql}, row:{row} col:{col} data:{self.queryResult[row][col]} == expect:{data}")
 
+    # return true or false replace exit, no print out
+    def checkRowColNoExit(self, row, col):
+        caller = inspect.getframeinfo(inspect.stack()[2][0])
+        if row < 0:
+            args = (caller.filename, caller.lineno, self.sql, row)
+            return False
+        if col < 0:
+            args = (caller.filename, caller.lineno, self.sql, row)
+            return False
+        if row > self.queryRows:
+            args = (caller.filename, caller.lineno, self.sql, row, self.queryRows)
+            return False
+        if col > self.queryCols:
+            args = (caller.filename, caller.lineno, self.sql, col, self.queryCols)
+            return False
+            
+        return True
+
+
+    # return true or false replace exit, no print out
+    def checkDataNoExit(self, row, col, data):
+        if self.checkRowColNoExit(row, col) == False:
+            return False
+        if self.queryResult[row][col] != data:
+            if self.cursor.istype(col, "TIMESTAMP"):
+                # suppose user want to check nanosecond timestamp if a longer data passed
+                if (len(data) >= 28):
+                    if pd.to_datetime(self.queryResult[row][col]) == pd.to_datetime(data):
+                        return True
+                else:
+                    if self.queryResult[row][col] == _parse_datetime(data):
+                        return True
+                return False
+
+            if str(self.queryResult[row][col]) == str(data):
+                return True
+            elif isinstance(data, float):
+                if abs(data) >= 1 and abs((self.queryResult[row][col] - data) / data) <= 0.000001:
+                    return True
+                elif abs(data) < 1 and abs(self.queryResult[row][col] - data) <= 0.000001:
+                    return True
+                else:
+                    return False
+            else:
+                return False
+                
+        return True
+
+
+    # loop execute sql then sleep(waitTime) , if checkData ok break loop
+    def checkDataLoop(self, row, col, data, sql, loopCount, waitTime):
+        # loop check util checkData return true
+        for i in range(loopCount):
+            self.query(sql)
+            if self.checkDataNoExit(row, col, data) :
+                self.checkData(row, col, data)
+                return
+            time.sleep(waitTime)
+
+        # last check
+        self.query(sql)
+        self.checkData(row, col, data)
+
+
     def getData(self, row, col):
         self.checkRowCol(row, col)
         return self.queryResult[row][col]
diff --git a/tests/system-test/6-cluster/5dnode1mnode.py b/tests/system-test/6-cluster/5dnode1mnode.py
index 548dd6c247..6f3ca8f20b 100644
--- a/tests/system-test/6-cluster/5dnode1mnode.py
+++ b/tests/system-test/6-cluster/5dnode1mnode.py
@@ -140,24 +140,21 @@ class TDTestCase:
     def check_alive(self):
         # check cluster alive
         tdLog.printNoPrefix("======== test cluster alive: ")
-        tdSql.query("show cluster alive;")
-        tdSql.checkData(0, 0, 1)
+        tdSql.checkDataLoop(0, 0, 1, "show cluster alive;", 10, 0.5)
 
         tdSql.query("show db.alive;")
         tdSql.checkData(0, 0, 1)
 
         # stop 5 dnode 
-        self.TDDnodes[4].stoptaosd()
-        tdSql.query("show cluster alive;")
-        tdSql.checkData(0, 0, 2)
-
+        self.TDDnodes.stoptaosd(5)
+        tdSql.checkDataLoop(0, 0, 2, "show cluster alive;", 10, 0.5)
+        
         tdSql.query("show db.alive;")
         tdSql.checkData(0, 0, 2)
 
         # stop 2 dnode
-        self.TDDnodes[1].stoptaosd()
-        tdSql.query("show cluster alive;")
-        tdSql.checkData(0, 0, 0)
+        self.TDDnodes.stoptaosd(2)
+        tdSql.checkDataLoop(0, 0, 0, "show cluster alive;", 10, 0.5)
 
         tdSql.query("show db.alive;")
         tdSql.checkData(0, 0, 0)