MB-61592: [cluster_tests] Add support for retry_on_assert...

... in poll_for_condition

Change-Id: If01543399d8fc8a2d23edb563825380a1df78366
Reviewed-on: https://review.couchbase.org/c/ns_server/+/232960
Well-Formed: Restriction Checker
Well-Formed: Build Bot <[email protected]>
Tested-by: Peter Searby <[email protected]>
Reviewed-by: Timofey Barmin <[email protected]>

Author: timofey-barmin

cluster_tests/testlib/testlib.py

@@ -509,28 +509,47 @@ def get_otp_nodes(cluster):
 
 def poll_for_condition(fun, sleep_time, attempts=None, timeout=None,
                        verbose=False, msg="poll for condition",
-                       retry_value=False):
+                       retry_value=False,
+                       retry_on_assert=False):
 
     assert (attempts is not None) or (timeout is not None)
     assert sleep_time > 0, "non-positive sleep_time specified"
     start_time = time.time()
     sleep_time_str = f"{sleep_time:.2f}s"
 
     attempt_count = 0
+    exception_obj = None
     while (attempts is None) or (attempt_count < attempts):
-        if timeout is not None:
-            assert (time.time() - start_time) < timeout, \
-                   f"{msg}: timed-out (timeout: {timeout}s)"
-        value = fun()
+        if timeout is not None and ((time.time() - start_time) >= timeout):
+            error_msg = f"{msg}: timed-out (timeout: {timeout}s)"
+            if exception_obj is None:
+                assert False, error_msg
+            else:
+                maybe_print(error_msg, verbose=verbose)
+                raise exception_obj
+        if retry_on_assert:
+            try:
+                value = fun()
+            except AssertionError as e:
+                exception_obj = e
+                value = retry_value
+        else:
+            value = fun()
         if value is not retry_value:
             maybe_print(f"Time taken for condition to complete: "
                         f"{time.time() - start_time: .2f}s", verbose=verbose)
             return value
         maybe_print(f"Sleeping for {sleep_time_str}", verbose=verbose)
         time.sleep(sleep_time)
         attempt_count += 1
-    assert False, f"{msg} didn't complete in: {attempts} attempts, " \
-                  f"sleep_time: {sleep_time_str}"
+
+    error_msg = f"{msg} didn't complete in: {attempts} attempts, " \
+                f"sleep_time: {sleep_time_str}"
+    if exception_obj is None:
+        assert False, error_msg
+    else:
+        maybe_print(error_msg, verbose=verbose)
+        raise exception_obj
 
 def metakv_get_succ(cluster, key, **kwargs):
     return get_succ(cluster, f"/_metakv{key}", **kwargs)