X-Git-Url: https://code.grnet.gr/git/ganeti-local/blobdiff_plain/79d22269bc9a9fb2c46523cf905d4c0b78958333..415feb2ee0a760e48a21a50488f6072f8afa9cff:/lib/utils/retry.py?ds=sidebyside

diff --git a/lib/utils/retry.py b/lib/utils/retry.py
index b5b7ff3..cc7541c 100644
--- a/lib/utils/retry.py
+++ b/lib/utils/retry.py
@@ -160,7 +160,7 @@ def Retry(fn, delay, timeout, args=None, wait_fn=time.sleep,
   while True:
     retry_args = []
     try:
-      # pylint: disable-msg=W0142
+      # pylint: disable=W0142
       return fn(*args)
     except RetryAgain, err:
       retry_args = err.args
@@ -171,7 +171,7 @@ def Retry(fn, delay, timeout, args=None, wait_fn=time.sleep,
     remaining_time = end_time - _time_fn()
 
     if remaining_time < 0.0:
-      # pylint: disable-msg=W0142
+      # pylint: disable=W0142
       raise RetryTimeout(*retry_args)
 
     assert remaining_time >= 0.0
@@ -182,3 +182,42 @@ def Retry(fn, delay, timeout, args=None, wait_fn=time.sleep,
       current_delay = calc_delay()
       if current_delay > 0.0:
         wait_fn(current_delay)
+
+
+def SimpleRetry(expected, fn, delay, timeout, args=None, wait_fn=time.sleep,
+                _time_fn=time.time):
+  """A wrapper over L{Retry} implementing a simpler interface.
+
+  All the parameters are the same as for L{Retry}, except it has one
+  extra argument: expected, which can be either a value (will be
+  compared with the result of the function, or a callable (which will
+  get the result passed and has to return a boolean). If the test is
+  false, we will retry until either the timeout has passed or the
+  tests succeeds. In both cases, the last result from calling the
+  function will be returned.
+
+  Note that this function is not expected to raise any retry-related
+  exceptions, always simply returning values. As such, the function is
+  designed to allow easy wrapping of code that doesn't use retry at
+  all (e.g. "if fn(args)" replaced with "if SimpleRetry(True, fn,
+  ...)".
+
+  @see: L{Retry}
+
+  """
+  rdict = {}
+
+  def helper(*innerargs):
+    # pylint: disable=W0142
+    result = rdict["result"] = fn(*innerargs)
+    if not ((callable(expected) and expected(result)) or result == expected):
+      raise RetryAgain()
+    return result
+
+  try:
+    result = Retry(helper, delay, timeout, args=args,
+                   wait_fn=wait_fn, _time_fn=_time_fn)
+  except RetryTimeout:
+    assert "result" in rdict
+    result = rdict["result"]
+  return result