test new retry() function and add parameters

Author: categulario

libtmux/test.py

@@ -9,22 +9,41 @@
 
 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__))
 example_dir = os.path.abspath(os.path.join(current_dir, "..", "examples"))
 fixtures_dir = os.path.realpath(os.path.join(current_dir, "fixtures"))
 
 
-def retry(fun, seconds=RETRY_TIMEOUT_SECONDS):
+class WaitTimeout(Exception):
+    '''Function timed out without meeting condition'''
+
+
+def retry(
+        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`` environmental variables.
+        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``.
 
     Returns
     -------
@@ -46,7 +65,11 @@ def retry(fun, seconds=RETRY_TIMEOUT_SECONDS):
     while not fun():
         end = time.time()
         if end - ini >= seconds:
-            raise Exception('Function timed out without meeting condition')
+            if raises:
+                raise WaitTimeout()
+            else:
+                break
+        time.sleep(interval)
 
 
 def get_test_session_name(server, prefix=TEST_SESSION_PREFIX):

tests/test_test.py

@@ -0,0 +1,52 @@
+from time import time
+import pytest
+
+from libtmux.test import retry, WaitTimeout
+
+
+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(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(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(never_true, 1, raises=False)
+
+    end = time()
+
+    assert abs((end - ini) - 1.0) < 0.01