feat: add wait_for_sync() for reliable sync verification

Author: suxb201

tests/cases/auth_acl.py

@@ -28,7 +28,7 @@ def acl():
     shake = h.Shake(opts)
 
     # wait sync done
-    p.ASSERT_TRUE_TIMEOUT(lambda: shake.is_consistent(), interval=0.01)
+    shake.wait_for_sync(timeout=10)
     p.log(shake.get_status())
 
     # check data

tests/cases/function.py

@@ -24,8 +24,8 @@ def filter_db():
         src.do("select", db)
         src.do("set", "key", "value")
 
-    # wait sync done
-    p.ASSERT_TRUE_TIMEOUT(lambda: shake.is_consistent(), timeout=10, interval=0.01)
+    # wait sync done (use db=1 because db=0 is filtered)
+    shake.wait_for_sync(timeout=10, db=1)
 
     dst.do("select", 0)
     p.ASSERT_EQ(dst.do("get", "key"), None)
@@ -53,15 +53,21 @@ def split_mset_to_set():
     p.log(f"opts: {opts}")
     shake = h.Shake(opts)
     src.do("mset", "k1", "v1", "k2", "v2", "k3", "v3")
-    # wait sync done
+
+    # wait sync done - check both consistency AND data presence
+    # is_consistent() can return true before MSET is captured by AOF streaming
+    def data_synced():
+        if not shake.is_consistent():
+            return False
+        dst.do("select", 1)
+        return dst.do("get", "k1") == b"v1"
+
     try:
-        p.ASSERT_TRUE_TIMEOUT(lambda: shake.is_consistent(), timeout=10, interval=0.01)
+        p.ASSERT_TRUE_TIMEOUT(data_synced, timeout=10, interval=0.01)
     except Exception as e:
         with open(f"{shake.dir}/data/shake.log") as f:
             p.log(f.read())
         raise e
-    dst.do("select", 1)
-    p.ASSERT_EQ(dst.do("get", "k1"), b"v1")
     p.ASSERT_EQ(dst.do("get", "k2"), b"v2")
     p.ASSERT_EQ(dst.do("get", "k3"), b"v3")
 

tests/cases/sync.py

@@ -1,5 +1,3 @@
-import time
-
 import pybbt as p
 
 import helpers as h
@@ -17,8 +15,8 @@ def test(src, dst):
     shake = h.Shake(opts)
 
     # wait sync done
-    try: # HTTPConnectionPool
-        p.ASSERT_TRUE_TIMEOUT(lambda: shake.is_consistent(), timeout=10, interval=0.01)
+    try:
+        shake.wait_for_sync(timeout=10)
     except Exception as e:
         with open(f"{shake.dir}/data/shake.log") as f:
             p.log(f.read())
@@ -28,9 +26,8 @@ def test(src, dst):
     inserter.add_data(src, cross_slots_cmd=cross_slots_cmd)
 
     # wait sync done
-    p.ASSERT_TRUE_TIMEOUT(lambda: shake.is_consistent(), interval=0.01)
+    shake.wait_for_sync(timeout=10)
     p.log(shake.get_status())
-    time.sleep(5)
 
     # check data
     inserter.check_data(src, cross_slots_cmd=cross_slots_cmd)

tests/helpers/shake.py

@@ -1,4 +1,6 @@
 import typing
+import time
+import uuid
 
 import pybbt
 import requests
@@ -21,7 +23,9 @@ def create_sync_opts(src: Redis, dst: Redis) -> typing.Dict:
             "redis_writer": {
                 "cluster": dst.is_cluster(),
                 "address": dst.get_address()
-            }
+            },
+            "__src": src,
+            "__dst": dst,
         }
         return d
 
@@ -35,32 +39,36 @@ def create_scan_opts(src: Redis, dst: Redis) -> typing.Dict:
             "redis_writer": {
                 "cluster": dst.is_cluster(),
                 "address": dst.get_address()
-            }
+            },
+            "__src": src,
+            "__dst": dst,
         }
         return d
 
     @staticmethod
-    def create_rdb_opts(rdb_path: str, dts: Redis) -> typing.Dict:
+    def create_rdb_opts(rdb_path: str, dst: Redis) -> typing.Dict:
         d = {
             "rdb_reader": {"filepath": rdb_path},
             "redis_writer": {
-                "cluster": dts.is_cluster(),
-                "address": dts.get_address()
-            }
+                "cluster": dst.is_cluster(),
+                "address": dst.get_address()
+            },
+            "__dst": dst,
         }
         return d
 
     @staticmethod
-    def create_aof_opts(aof_path: str, dts: Redis, timestamp: int = 0) -> typing.Dict:
+    def create_aof_opts(aof_path: str, dst: Redis, timestamp: int = 0) -> typing.Dict:
         d = {
             "aof_reader": {"filepath": aof_path, "timestamp": timestamp},
             "redis_writer": {
-                "cluster": dts.is_cluster(),
-                "address": dts.get_address()
+                "cluster": dst.is_cluster(),
+                "address": dst.get_address()
             },
             "advanced": {
                 "log_level": "debug"
-            }
+            },
+            "__dst": dst,
         }
         return d
 
@@ -70,6 +78,11 @@ def __init__(self, opts: typing.Dict):
         self.case_ctx = pybbt.get_case_context()
         self.status_port = get_free_port()
         self.status_url = f"http://localhost:{self.status_port}"
+
+        # 提取并保存 src/dst 引用
+        self.src = opts.pop("__src", None)
+        self.dst = opts.pop("__dst", None)
+
         opts["advanced"] = {"status_port": self.status_port, "log_level": "debug"}
 
         self.dir = f"{self.case_ctx.dir}/shake{self.status_port}"
@@ -82,6 +95,9 @@ def __init__(self, opts: typing.Dict):
 
     @staticmethod
     def run_once(opts: typing.Dict):
+        # 移除内部字段
+        opts.pop("__src", None)
+        opts.pop("__dst", None)
         status_port = get_free_port()
         run_dir = f"{pybbt.get_case_context().dir}/shake{status_port}"
         create_empty_dir(run_dir)
@@ -107,3 +123,59 @@ def _wait_start(self, timeout=5):
 
     def is_consistent(self):
         return self.get_status()["consistent"]
+
+    def wait_for_sync(self, timeout: float = 5, db: int = 0):
+        """
+        Wait for data to be fully synced from src to dst.
+
+        1. Wait for is_consistent() - RDB sync done
+        2. Check dbsize matches
+        3. Write marker key to src
+        4. Wait for marker to appear in dst
+        5. Delete marker key
+        6. Wait for consistent again
+        """
+        if self.src is None or self.dst is None:
+            raise ValueError("src and dst not available, use ShakeOpts.create_sync_opts() or create_scan_opts()")
+
+        timer = Timer()
+
+        # 1. Wait for is_consistent()
+        while not self.is_consistent():
+            if timer.elapsed() > timeout:
+                raise TimeoutError(f"is_consistent() not reached within {timeout}s")
+            time.sleep(0.01)
+
+        # 2. Check dbsize (cluster doesn't support SELECT, use db 0)
+        if not self.src.is_cluster():
+            self.src.do("select", db)
+        if not self.dst.is_cluster():
+            self.dst.do("select", db)
+        src_dbsize = self.src.dbsize()
+        while self.dst.dbsize() != src_dbsize:
+            if timer.elapsed() > timeout:
+                raise TimeoutError(f"dbsize mismatch within {timeout}s: src={src_dbsize}, dst={self.dst.dbsize()}")
+            time.sleep(0.01)
+
+        # 3. Write marker key
+        marker_key = f"__shake_sync_marker_{uuid.uuid4()}"
+        marker_value = str(time.time())
+        self.src.do("set", marker_key, marker_value)
+
+        # 4. Wait for marker in dst
+        while True:
+            result = self.dst.do("get", marker_key)
+            if result == marker_value.encode():
+                break
+            if timer.elapsed() > timeout:
+                raise TimeoutError(f"marker key not synced within {timeout}s")
+            time.sleep(0.01)
+
+        # 5. Delete marker key
+        self.src.do("del", marker_key)
+
+        # 6. Wait for consistent again
+        while not self.is_consistent():
+            if timer.elapsed() > timeout:
+                raise TimeoutError(f"final is_consistent() not reached within {timeout}s")
+            time.sleep(0.01)