fix tests

Author: laur89

retry_decorator/retry_decorator.py

@@ -6,7 +6,6 @@
 import logging
 import time
 import random
-import types
 
 
 def _isiter(i):
@@ -28,16 +27,16 @@ def _deco_retry(f, exc=Exception, tries=10, timeout_secs=1.0, logger=None, callb
     :param callback_by_exception:
     :return:
     """
+
     def f_retry(*args, **kwargs):
         mtries, mdelay = tries, timeout_secs
         run_one_last_time = True
+
         while mtries > 1:
             try:
                 return f(*args, **kwargs)
             except exc as e:
                 # check if this exception is something the caller wants special handling for
-                if isinstance(callback_by_exception, (types.FunctionType, list, tuple)):
-                    callback_by_exception = {Exception: callback_by_exception}
                 callback_errors = callback_by_exception or {}
 
                 for error_type in callback_errors:
@@ -97,6 +96,9 @@ def __init__(
         elif tries < 1:
             raise ValueError("[tries] arg needs to be an int >= 1")
 
+        if callable(callback_by_exception) or isinstance(callback_by_exception, (list, tuple)):
+            callback_by_exception = {Exception: callback_by_exception}
+
         self.exc = exc
         self.tries = tries
         self.timeout_secs = timeout_secs

tests/test_callback.py

@@ -14,14 +14,15 @@ class ExampleTestError(Exception):
     pass
 
 
-class_for_testing = None
+class_for_testing = ClassForTesting()
 
 
 class MyTestCase(unittest.TestCase):
 
-    def setUp():
-        global class_for_testing
-        class_for_testing = ClassForTesting()
+    def setUp(self):
+        class_for_testing.hello = None
+        class_for_testing.cb_counter = 0
+        class_for_testing.exe_counter = 0
 
     def test_callback_invoked_on_configured_exception_type(self):
         try:
@@ -77,7 +78,7 @@ def test_verify_breakout_true_works(self):
         except Exception:
             pass
         self.assertEqual(class_for_testing.hello, 'baz')
-        self.assertEqual(class_for_testing.cb_counter, 6)
+        self.assertEqual(class_for_testing.cb_counter, 6)  # we had 2 handlers, but because of breakout=True only first of them was ever ran
         self.assertEqual(class_for_testing.exe_counter, 7)
 
     def test_verify_run_last_time_false_works(self):
@@ -94,18 +95,29 @@ def test_verify_tries_1_is_ok(self):
             my_test_func_9()
         except Exception:
             pass
+        self.assertEqual(class_for_testing.hello, None)
+        self.assertEqual(class_for_testing.cb_counter, 0)
+        self.assertEqual(class_for_testing.exe_counter, 1)
+
+    def test_verify_run_last_time_false_with_2_tries(self):
+        try:
+            my_test_func_10()
+        except Exception:
+            pass
         self.assertEqual(class_for_testing.hello, 'foo')
+        self.assertEqual(class_for_testing.cb_counter, 1)
+        self.assertEqual(class_for_testing.exe_counter, 1)
 
     def test_verify_tries_0_errors_out(self):
         try:
-            my_test_func_10()
+            retry_decorator.retry(tries=0, callback_by_exception=partial(callback_logic, class_for_testing, 'hello', 'foo'))
             raise Exception('Expected ValueError to be thrown')
         except ValueError:
             pass
 
     def test_verify_tries_not_int_is_error(self):
         try:
-            my_test_func_11()
+            retry_decorator.retry(tries='not int', callback_by_exception=partial(callback_logic, class_for_testing, 'hello', 'foo'))
             raise Exception('Expected TypeError to be thrown')
         except TypeError:
             pass
@@ -117,10 +129,6 @@ def callback_logic(instance, attr_to_set, value_to_set):
     instance.cb_counter += 1
 
 
-def get_callback_tuple_breakout(attr_value, breakout_value=False):
-    return (partial(callback_logic, class_for_testing, 'hello', attr_value), breakout_value)
-
-
 @retry_decorator.retry(exc=ExampleTestError, tries=2, callback_by_exception={
     ExampleTestError: partial(callback_logic, class_for_testing, 'hello', 'world')})
 def my_test_func():
@@ -173,7 +181,7 @@ def my_test_func_7():
 
 @retry_decorator.retry(tries=8, callback_by_exception={
     TypeError: partial(callback_logic, class_for_testing, 'hello', 'foo'),
-    Exception: (partial(callback_logic, class_for_testing, 'hello', 'bar'), False, False)
+    Exception: (partial(callback_logic, class_for_testing, 'hello', 'bar'), (False, False))
     })
 def my_test_func_8():
     class_for_testing.exe_counter += 1
@@ -182,15 +190,14 @@ def my_test_func_8():
 
 @retry_decorator.retry(tries=1, callback_by_exception=partial(callback_logic, class_for_testing, 'hello', 'foo'))
 def my_test_func_9():
+    class_for_testing.exe_counter += 1
     raise TypeError('type oh noes.')
 
 
+@retry_decorator.retry(tries=2, callback_by_exception=(partial(callback_logic, class_for_testing, 'hello', 'foo'), (False, False)))
 def my_test_func_10():
-    retry_decorator.retry(tries=0, callback_by_exception=partial(callback_logic, class_for_testing, 'hello', 'foo'))
-
-
-def my_test_func_11():
-    retry_decorator.retry(tries='not int', callback_by_exception=partial(callback_logic, class_for_testing, 'hello', 'foo'))
+    class_for_testing.exe_counter += 1
+    raise TypeError('type oh noes.')
 
 
 if __name__ == '__main__':