Add retry_until() to take a callable and enforce the timeout

Author: categulario

libtmux/test.py

@@ -5,10 +5,13 @@
 import tempfile
 import time
 
+from .exc import WaitTimeout
+
 logger = logging.getLogger(__name__)
 
 TEST_SESSION_PREFIX = "libtmux_"
 RETRY_TIMEOUT_SECONDS = int(os.getenv("RETRY_TIMEOUT_SECONDS", 8))
+RETRY_INTERVAL_SECONDS = float(os.getenv("RETRY_INTERVAL_SECONDS", 0.05))
 
 namer = tempfile._RandomNameSequence()
 current_dir = os.path.abspath(os.path.dirname(__file__))
@@ -43,6 +46,52 @@ def retry(seconds=RETRY_TIMEOUT_SECONDS):
     return (lambda: time.time() < time.time() + seconds)()
 
 
+def retry_until(
+    fun,
+    seconds=RETRY_TIMEOUT_SECONDS,
+    *,
+    interval=RETRY_INTERVAL_SECONDS,
+    raises=True,
+):
+    """
+    Retry a function until a condition meets or the specified time passes.
+
+    Parameters
+    ----------
+    fun : callable
+        A function that will be called repeatedly until it returns ``True``  or
+        the specified time passes.
+    seconds : int
+        Seconds to retry. Defaults to ``8``, which is configurable via
+        ``RETRY_TIMEOUT_SECONDS`` environment variables.
+    interval : float
+        Time in seconds to wait between calls. Defaults to ``0.05`` and is
+        configurable via ``RETRY_INTERVAL_SECONDS`` environment variable.
+    raises : bool
+        Wether or not to raise an exception on timeout. Defaults to ``True``.
+
+    Examples
+    --------
+
+    >>> def f():
+    ...     p = w.attached_pane
+    ...     p.server._update_panes()
+    ...     return p.current_path == pane_path
+    ...
+    ... retry(f)
+    """
+    ini = time.time()
+
+    while not fun():
+        end = time.time()
+        if end - ini >= seconds:
+            if raises:
+                raise WaitTimeout()
+            else:
+                break
+        time.sleep(interval)
+
+
 def get_test_session_name(server, prefix=TEST_SESSION_PREFIX):
     """
     Faker to create a session name that doesn't exist.

tests/test_test.py

@@ -0,0 +1,53 @@
+from time import time
+
+import pytest
+
+from libtmux.test import WaitTimeout, retry_until
+
+
+def test_retry_three_times():
+    ini = time()
+    value = 0
+
+    def call_me_three_times():
+        nonlocal value
+
+        if value == 2:
+            return True
+
+        value += 1
+
+        return False
+
+    retry_until(call_me_three_times, 1)
+
+    end = time()
+
+    assert abs((end - ini) - 0.1) < 0.01
+
+
+def test_function_times_out():
+    ini = time()
+
+    def never_true():
+        return False
+
+    with pytest.raises(WaitTimeout):
+        retry_until(never_true, 1)
+
+    end = time()
+
+    assert abs((end - ini) - 1.0) < 0.01
+
+
+def test_function_times_out_no_rise():
+    ini = time()
+
+    def never_true():
+        return False
+
+    retry_until(never_true, 1, raises=False)
+
+    end = time()
+
+    assert abs((end - ini) - 1.0) < 0.01