add a helper for while loops

Signed-off-by: Alfredo Deza <alfredo.deza@inktank.com>

diff --git a/teuthology/contextutil.py b/teuthology/contextutil.py
index bfcb830ae04b7127a4953388c4775d240bb2121f..90c10196270ad127987c160aa3db87503e2aabde 100644 (file)
--- a/teuthology/contextutil.py
+++ b/teuthology/contextutil.py
@@ -1,6 +1,7 @@
 import contextlib
 import sys
 import logging
+import time
 
 log = logging.getLogger(__name__)
 
@@ -41,3 +42,74 @@ def nested(*managers):
             # the right information. Another exception may
             # have been raised and caught by an exit method
             raise exc[0], exc[1], exc[2]
+
+
+class MaxWhileTries(Exception):
+    pass
+
+
+class safe_while(object):
+    """
+    A contect manager to remove boiler plate code that deals with `while` loops
+    that need a given number of tries and some seconds to sleep between each
+    one of those tries.
+
+    The most simple example possible will try 5 times sleeping for 5 seconds
+    and increasing the sleep time by 5 each time::
+
+        >>> from teuthology.contexutil import safe_while
+        >>> with safe_while() as bomb:
+        ...    while 1:
+        ...        bomb()
+        ...        # repetitive code here
+        ...
+        Traceback (most recent call last):
+        ...
+        MaxWhileTries: reached maximum tries (5) after waiting for 75 seconds
+
+    Yes, this adds yet another level of indentation but it allows you to
+    implement while loops exactly the same as before with just 1 more
+    indentation level and one extra call. Everything else stays the same,
+    code-wise. So adding this helper to existing code is simpler.
+
+    The defaults are to start the sleeping time in seconds at 5s and to add
+    5 more seconds at every point in the loop. Setting the increment value to
+    0 makes the sleep time in seconds stay the same throughout the calls.
+
+    """
+
+    def __init__(self, sleep=5, increment=5, tries=5):
+        self.sleep = sleep
+        self.increment = increment
+        self.tries = tries
+        self.counter = 0
+        self.sleep_current = sleep
+
+    def _make_error_msg(self):
+        """
+        Sum the total number of seconds we waited while providing the number
+        of tries we attempted
+        """
+        total_seconds_waiting = sum(
+            map(
+                lambda x: x * self.increment,
+                range(1, self.tries+1)
+            )
+        )
+        return 'reached maximum tries (%s) after waiting for %s seconds' % (
+            self.tries, total_seconds_waiting
+        )
+
+    def __call__(self):
+        self.counter += 1
+        if self.counter > self.tries:
+            error_msg = self._make_error_msg()
+            raise MaxWhileTries(error_msg)
+        time.sleep(self.sleep_current)
+        self.sleep_current += self.increment
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return False