Modify to raise exception if retries unsuccessful, add decorator

Author: Doru Irimescu

src/exception_with_retry/source.py

@@ -1,4 +1,6 @@
+import functools
 import time
+from typing import Any, Callable
 
 __author__ = "Doru Irimescu"
 __copyright__ = "Doru Irimescu"
@@ -27,7 +29,7 @@ def __init__(self, method, n_retry=5, sleep_time_s=0.5):
         self._max_retries = n_retry
         self._sleep_time_s = sleep_time_s
 
-    def run(self, args):
+    def run(self, *args, **kwargs):
         """Execute the wrapped method, and retry for _max_retries.
 
         Args:
@@ -36,11 +38,38 @@ def run(self, args):
         Returns:
             [method return type]: value which the wrapped method returns.
         """
-        if self._retries < self._max_retries:
+
+        try:
             self._retries = self._retries + 1
-            try:
-                return self._method(*args)
-            except Exception as e:
+            return self._method(*args, **kwargs)
+        except Exception as e:
+            if self._retries < self._max_retries:
                 print(e)
                 time.sleep(self._sleep_time_s)
-                return self.run(args)
+                return self.run(*args, **kwargs)
+            else:
+                raise e
+
+
+def exception_with_retry(n_retry: int, sleep_time_s: float):
+    """Use this decorator to retry calling your function n_retry times,
+    with sleep_time_s in between unsuccessful calls.
+    Will finally throw the original error if the last retry fails.
+
+    Args:
+        n_retry (int): number of retries
+        sleep_time_s (float): time to sleep between unsuccessful retries, in seconds
+
+    Returns:
+        Callable[..., Any]: decorated function
+    """
+
+    def dec(func: Callable[..., Any]) -> Callable[..., Any]:
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            ewr = ExceptionWithRetry(func, n_retry, sleep_time_s)
+            return ewr.run(*args, **kwargs)
+
+        return wrapper
+
+    return dec

tests/test_1.py

@@ -1,6 +1,7 @@
 import unittest
+from typing import List
 
-from exception_with_retry.source import ExceptionWithRetry
+from exception_with_retry.source import ExceptionWithRetry, exception_with_retry
 
 __author__ = "Doru Irimescu"
 __copyright__ = "Doru Irimescu"
@@ -31,39 +32,88 @@ def method_caller(arg):
     method(arg)
 
 
+def wrapper_test_method(n: int):
+    if n < 0:
+        raise Exception("Not gonna happen")
+    else:
+        return 15
+
+
 class TestExceptionWithRetry(unittest.TestCase):
     def test_1(self):
-        ewr = ExceptionWithRetry(method, 5, 0.0)
-        arg = [0]
-        ewr.run([arg])
-        self.assertEqual(arg, [0, 1, 2, 3, 4, 5])
+        with self.assertRaises(Exception):
+            ewr = ExceptionWithRetry(method, 5, 0.0)
+            arg = [0]
+            ewr.run(arg)
+            self.assertEqual(arg, [0, 1, 2, 3, 4, 5])
 
     def test_2(self):
-        ewr = ExceptionWithRetry(method, 0, 0.0)
-        arg = [0]
-        ewr.run([arg])
-        self.assertEqual(arg, [0])
+        with self.assertRaises(Exception):
+            ewr = ExceptionWithRetry(method, 0, 0.0)
+            arg = [0]
+            ewr.run(arg)
+            self.assertEqual(arg, [0])
 
     def test_3(self):
-        ewr = ExceptionWithRetry(method2, 2, 0.0)
-        arg = [0]
-        ewr.run([arg, 1])
-        self.assertEqual(arg, [0, 1, 2])
+        with self.assertRaises(Exception):
+            ewr = ExceptionWithRetry(method2, 2, 0.0)
+            arg = [0]
+            ewr.run(arg, 1)
+            self.assertEqual(arg, [0, 1, 2])
 
     def test_4(self):
-        ewr = ExceptionWithRetry(method_caller, 2, 0.0)
-        arg = [1, 2]
-        ewr.run([arg])
-        self.assertEqual(arg, [1, 2, 3, 4])
+        with self.assertRaises(Exception):
+            ewr = ExceptionWithRetry(method_caller, 2, 0.0)
+            arg = [1, 2]
+            ewr.run(arg)
+            self.assertEqual(arg, [1, 2, 3, 4])
 
     def test_5(self):
-        t = TestMethods()
-        ewr = ExceptionWithRetry(t.method, 5, 0.0)
-        arg = [0]
-        ewr.run([arg])
-        self.assertEqual(arg, [0, 20, 40, 60, 80, 100])
+        with self.assertRaises(Exception):
+            t = TestMethods()
+            ewr = ExceptionWithRetry(t.method, 5, 0.0)
+            arg = [0]
+            ewr.run(arg)
+            self.assertEqual(arg, [0, 20, 40, 60, 80, 100])
 
     def test_6(self):
         ewr = ExceptionWithRetry(method3, 5, 0.0)
-        result = ewr.run([])
+        result = ewr.run()
         self.assertEqual(result, 25)
+
+
+@exception_with_retry(3, 0.1)
+def wrapped_method(number_1: int, number_2: int = 0, calls: List = []):
+    calls.append(1)
+    if number_1 < 0:
+        raise Exception("Not gonna happen")
+    else:
+        return number_1 + number_2
+
+
+class TestExceptionWithRetryDecorator(unittest.TestCase):
+    def test_no_exception(self):
+        result = wrapped_method(number_1=1)
+        self.assertEqual(1, result)
+
+        calls = []
+        result = wrapped_method(number_1=1, number_2=10, calls=calls)
+        self.assertEqual(11, result)
+        self.assertEqual(1, len(calls))
+
+    def test_throws_exception(self):
+        with self.assertRaises(Exception):
+            calls = []
+            result = wrapped_method(-1, 0, calls)
+            self.assertEqual(15, result)
+            self.assertEqual(3, len(calls))
+
+        with self.assertRaises(Exception):
+            calls = []
+            wrapped_method(number_1=-1, number_2=2)
+            self.assertEqual(3, len(calls))
+
+        with self.assertRaises(Exception):
+            calls = []
+            wrapped_method(-1, 2)
+            self.assertEqual(3, len(calls))