From ac974739a53b6e3fc7a0b25d5d1c01cba7c2f6e2 Mon Sep 17 00:00:00 2001
From: Lukas Slebodnik <lslebodn@redhat.com>
Date: Tue, 5 Feb 2019 10:42:08 +0100
Subject: [PATCH] Utils: Add decorator for retry

It might be useful in checks which do lot of network communication.
---
 colin/utils/cmd_tools.py |  29 +++++++++++
 tests/unit/test_utils.py | 103 +++++++++++++++++++++++++++++++++++++--
 2 files changed, 129 insertions(+), 3 deletions(-)

diff --git a/colin/utils/cmd_tools.py b/colin/utils/cmd_tools.py
index 2da4118..f3f08a8 100644
--- a/colin/utils/cmd_tools.py
+++ b/colin/utils/cmd_tools.py
@@ -13,9 +13,11 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
+import functools
 import logging
 import subprocess
 import threading
+import time
 
 try:
     import thread
@@ -120,3 +122,30 @@ def inner(*args, **kwargs):
         return inner
 
     return outer
+
+
+def retry(retry_count=5, delay=2):
+    """
+    Use as decorator to retry functions few times with delays
+
+    Exception will be raised if last call fails
+
+    :param retry_count: int could of retries in case of failures. It must be
+                        a positive number
+    :param delay: int delay between retries
+    """
+    if retry_count <= 0:
+        raise ValueError("retry_count have to be positive")
+
+    def decorator(f):
+        @functools.wraps(f)
+        def wrapper(*args, **kwargs):
+            for i in range(retry_count, 0, -1):
+                try:
+                    return f(*args, **kwargs)
+                except Exception:
+                    if i <= 1:
+                        raise
+                time.sleep(delay)
+        return wrapper
+    return decorator
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
index 0d23182..4185a56 100644
--- a/tests/unit/test_utils.py
+++ b/tests/unit/test_utils.py
@@ -13,12 +13,13 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
-from time import sleep
+import time
 
 import pytest
 
 from colin.core.exceptions import ColinException
 from colin.utils.cmd_tools import exit_after
+from colin.utils.cmd_tools import retry
 
 
 @exit_after(1)
@@ -28,7 +29,7 @@ def fast_fce():
 
 @exit_after(1)
 def slow_fce():
-    sleep(2)
+    time.sleep(2)
 
 
 @exit_after(1)
@@ -52,4 +53,100 @@ def test_timeout_bad_fce():
 
 def test_timeout_dirrect():
     with pytest.raises(TimeoutError):
-        exit_after(1)(sleep)(2)
+        exit_after(1)(time.sleep)(2)
+
+
+COUNTER = 0
+
+
+def raise_exception():
+    global COUNTER
+    COUNTER += 1
+
+    raise Exception('I am bad function')
+
+
+def test_no_retry_for_success():
+    global COUNTER
+    COUNTER = 0
+
+    @retry(5, 0)
+    def always_success():
+        global COUNTER
+        COUNTER = COUNTER + 1
+
+        return 42
+
+    assert always_success() == 42
+    assert COUNTER == 1
+
+
+def test_retry_with_exception():
+    global COUNTER
+    COUNTER = 0
+
+    @retry(5, 0)
+    def always_raise_exception():
+        raise_exception()
+
+    with pytest.raises(Exception) as ex:
+        always_raise_exception()
+
+    assert str(ex.value) == 'I am bad function'
+    assert COUNTER == 5
+
+
+def test_wrong_parameter():
+    with pytest.raises(ValueError) as ex:
+        retry(-1, 1)
+    assert str(ex.value) == 'retry_count have to be positive'
+
+    with pytest.raises(ValueError) as ex:
+        retry(0, 1)
+    assert str(ex.value) == 'retry_count have to be positive'
+
+    @retry(5, -1)
+    def fail_negative_sleep():
+        raise_exception()
+
+    with pytest.raises(ValueError) as ex:
+        fail_negative_sleep()
+    assert str(ex.value) == 'sleep length must be non-negative'
+
+
+def test_retry_with_sleep():
+    global COUNTER
+    COUNTER = 0
+
+    @retry(4, .5)
+    def fail_and_sleep():
+        raise_exception()
+
+    time_start = time.time()
+    with pytest.raises(Exception) as ex:
+        fail_and_sleep()
+    time_end = time.time()
+
+    assert str(ex.value) == 'I am bad function'
+    assert COUNTER == 4
+
+    # there are 3 sleeps between 4 delays
+    assert time_end - time_start >= 1.5
+    # there were not 4 sleeps
+    assert time_end - time_start < 4
+
+
+def test_recover_after_few_failures():
+    global COUNTER
+    COUNTER = 0
+
+    @retry(5, 0)
+    def sleep_like_a_baby():
+        global COUNTER
+        if COUNTER < 3:
+            COUNTER += 1
+            raise Exception("sleeping")
+        return []
+
+    assert sleep_like_a_baby() == []
+    assert COUNTER == 3