fixed a bug where classes not extending from the Python exception hierarchy could slip through, preparing for 1.0.1 release

4 files changed, 99 insertions, 4 deletions

diff --git a/HISTORY.rst b/HISTORY.rst
index 85be6b0..c971b0b 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -3,6 +3,11 @@
 History
 -------
 
+1.0.1 (2013-03-20)
+++++++++++++++++++
+- Fixed a bug where classes not extending from the Python exception hierarchy could slip through
+- Update test suite for custom Python exceptions
+
 1.0.0 (2013-01-21)
 ++++++++++++++++++
 - First stable, tested version now exists
diff --git a/retrying.py b/retrying.py
index 55d09ac..c624d09 100644
--- a/retrying.py
+++ b/retrying.py
@@ -13,6 +13,7 @@
 ## limitations under the License.
 
 import random
+import sys
 import time
 
 # sys.maxint / 2, since Python 3.2 doesn't have a sys.maxint...
@@ -150,7 +151,8 @@ class Retrying:
         while True:
             try:
                 attempt = Attempt(fn(*args, **kwargs), attempt_number, False)
-            except BaseException as e:
+            except:
+                e = sys.exc_info()[1]
                 attempt = Attempt(e, attempt_number, True)
 
             if not self.should_reject(attempt):
diff --git a/setup.py b/setup.py
index c95a801..cd2c7d5 100644
--- a/setup.py
+++ b/setup.py
@@ -18,7 +18,7 @@ if sys.argv[-1] == 'publish':
 
 settings.update(
     name='retrying',
-    version='1.0.0',
+    version='1.0.1',
     description='Retrying',
     long_description=open('README.rst').read() + '\n\n' +
                      open('HISTORY.rst').read(),
diff --git a/test_retrying.py b/test_retrying.py
index 4de06d7..bfb311b 100644
--- a/test_retrying.py
+++ b/test_retrying.py
@@ -162,6 +162,34 @@ class NoNameErrorAfterCount:
             raise NameError()
         return True
 
+class CustomError:
+    """
+    This is a custom exception class that doesn't inherit from any of the Python base Exception hierarchy.
+    """
+    
+    def __init__(self, value):
+        self.value = value
+
+    def __str__(self):
+        return repr(self.value)
+
+class NoCustomErrorAfterCount:
+    """
+    This class holds counter state for invoking a method several times in a row.
+    """
+
+    def __init__(self, count):
+        self.counter = 0
+        self.count = count
+
+    def go(self):
+        """
+        Raise a CustomError until after count threshold has been crossed, then return True.
+        """
+        if self.counter < self.count:
+            self.counter += 1
+            raise CustomError("This is a Custom exception class")
+        return True
 
 def retry_if_result_none(result):
     return result is None
@@ -213,6 +241,29 @@ def _retryable_default(thing):
 def _retryable_default_f(thing):
     return thing.go()
 
+@retry(retry_on_exception=retry_if_exception_of_type(CustomError))
+def _retryable_test_with_exception_type_custom(thing):
+    return thing.go()
+
+@retry(retry_on_exception=retry_if_exception_of_type(CustomError), wrap_exception=True)
+def _retryable_test_with_exception_type_custom_wrap(thing):
+    return thing.go()
+
+@retry(
+    stop='stop_after_attempt',
+    stop_max_attempt_number=3,
+    retry_on_exception=retry_if_exception_of_type(CustomError))
+def _retryable_test_with_exception_type_custom_attempt_limit(thing):
+    return thing.go()
+
+@retry(
+    stop='stop_after_attempt',
+    stop_max_attempt_number=3,
+    retry_on_exception=retry_if_exception_of_type(CustomError),
+    wrap_exception=True)
+def _retryable_test_with_exception_type_custom_attempt_limit_wrap(thing):
+    return thing.go()
+
 class TestDecoratorWrapper(unittest.TestCase):
 
     def test_with_wait(self):
@@ -240,13 +291,31 @@ class TestDecoratorWrapper(unittest.TestCase):
 
         try:
             _retryable_test_with_exception_type_io_attempt_limit(NoIOErrorAfterCount(5))
-            self.fail("RetryError expected")
+            self.fail("Expected RetryError")
         except RetryError as re:
             self.assertEqual(3, re.last_attempt.attempt_number)
             self.assertTrue(re.last_attempt.has_exception)
             self.assertTrue(isinstance(re.last_attempt.value, IOError))
 
+        self.assertTrue(_retryable_test_with_exception_type_custom(NoCustomErrorAfterCount(5)))
+
+        try:
+            _retryable_test_with_exception_type_custom(NoNameErrorAfterCount(5))
+            self.fail("Expected NameError")
+        except NameError as n:
+            self.assertTrue(isinstance(n, NameError))
+
+        try:
+            _retryable_test_with_exception_type_custom_attempt_limit(NoCustomErrorAfterCount(5))
+            self.fail("Expected RetryError")
+        except RetryError as re:
+            self.assertEqual(3, re.last_attempt.attempt_number)
+            self.assertTrue(re.last_attempt.has_exception)
+            self.assertTrue(isinstance(re.last_attempt.value, CustomError))
+
     def test_wrapped_exception(self):
+
+        # base exception cases
         self.assertTrue(_retryable_test_with_exception_type_io_wrap(NoIOErrorAfterCount(5)))
 
         try:
@@ -257,15 +326,34 @@ class TestDecoratorWrapper(unittest.TestCase):
 
         try:
             _retryable_test_with_exception_type_io_attempt_limit_wrap(NoIOErrorAfterCount(5))
-            self.fail("RetryError expected")
+            self.fail("Expected RetryError")
         except RetryError as re:
             self.assertEqual(3, re.last_attempt.attempt_number)
             self.assertTrue(re.last_attempt.has_exception)
             self.assertTrue(isinstance(re.last_attempt.value, IOError))
 
+        # custom error cases
+        self.assertTrue(_retryable_test_with_exception_type_custom_wrap(NoCustomErrorAfterCount(5)))
+
+        try:
+            _retryable_test_with_exception_type_custom_wrap(NoNameErrorAfterCount(5))
+            self.fail("Expected RetryError")
+        except RetryError as r:
+            self.assertTrue(isinstance(r.last_attempt.value, NameError))
+
+        try:
+            _retryable_test_with_exception_type_custom_attempt_limit_wrap(NoCustomErrorAfterCount(5))
+            self.fail("Expected RetryError")
+        except RetryError as re:
+            self.assertEqual(3, re.last_attempt.attempt_number)
+            self.assertTrue(re.last_attempt.has_exception)
+            self.assertTrue(isinstance(re.last_attempt.value, CustomError))
+
     def test_defaults(self):
         self.assertTrue(_retryable_default(NoNameErrorAfterCount(5)))
         self.assertTrue(_retryable_default_f(NoNameErrorAfterCount(5)))
+        self.assertTrue(_retryable_default(NoCustomErrorAfterCount(5)))
+        self.assertTrue(_retryable_default_f(NoCustomErrorAfterCount(5)))
 
 if __name__ == '__main__':
     unittest.main()
\ No newline at end of file