Copy whole internal state when retry_with (#233)

tenacity/__init__.py

@@ -181,6 +181,10 @@ def __init__(self, sleep):
 _unset = object()
 
 
+def _first_set(first, second):
+    return second if first is _unset else first
+
+
 class RetryError(Exception):
     """Encapsulates the last attempt instance right before giving up."""
 
@@ -257,19 +261,23 @@ def copy(
         after=_unset,
         before_sleep=_unset,
         reraise=_unset,
+        retry_error_cls=_unset,
+        retry_error_callback=_unset,
     ):
         """Copy this object with some parameters changed if needed."""
-        if before_sleep is _unset:
-            before_sleep = self.before_sleep
         return self.__class__(
-            sleep=self.sleep if sleep is _unset else sleep,
-            stop=self.stop if stop is _unset else stop,
-            wait=self.wait if wait is _unset else wait,
-            retry=self.retry if retry is _unset else retry,
-            before=self.before if before is _unset else before,
-            after=self.after if after is _unset else after,
-            before_sleep=before_sleep,
-            reraise=self.reraise if after is _unset else reraise,
+            sleep=_first_set(sleep, self.sleep),
+            stop=_first_set(stop, self.stop),
+            wait=_first_set(wait, self.wait),
+            retry=_first_set(retry, self.retry),
+            before=_first_set(before, self.before),
+            after=_first_set(after, self.after),
+            before_sleep=_first_set(before_sleep, self.before_sleep),
+            reraise=_first_set(reraise, self.reraise),
+            retry_error_cls=_first_set(retry_error_cls, self.retry_error_cls),
+            retry_error_callback=_first_set(
+                retry_error_callback, self.retry_error_callback
+            ),
         )
 
     def __repr__(self):

tenacity/tests/test_tenacity.py

@@ -882,15 +882,6 @@ def test_with_wait(self):
         self.assertGreaterEqual(t, 250)
         self.assertTrue(result)
 
-    def test_retry_with(self):
-        start = current_time_ms()
-        result = _retryable_test_with_wait.retry_with(wait=tenacity.wait_fixed(0.1))(
-            NoneReturnUntilAfterCount(5)
-        )
-        t = current_time_ms() - start
-        self.assertGreaterEqual(t, 500)
-        self.assertTrue(result)
-
     def test_with_stop_on_return_value(self):
         try:
             _retryable_test_with_stop(NoneReturnUntilAfterCount(5))
@@ -1041,6 +1032,47 @@ def __call__(self):
         self.assertEqual(h(), "Hello")
 
 
+class TestRetryWith:
+    def test_redefine_wait(self):
+        start = current_time_ms()
+        result = _retryable_test_with_wait.retry_with(wait=tenacity.wait_fixed(0.1))(
+            NoneReturnUntilAfterCount(5)
+        )
+        t = current_time_ms() - start
+        assert t >= 500
+        assert result is True
+
+    def test_redefine_stop(self):
+        result = _retryable_test_with_stop.retry_with(
+            stop=tenacity.stop_after_attempt(5)
+        )(NoneReturnUntilAfterCount(4))
+        assert result is True
+
+    def test_retry_error_cls_should_be_preserved(self):
+        @retry(stop=tenacity.stop_after_attempt(10), retry_error_cls=ValueError)
+        def _retryable():
+            raise Exception("raised for test purposes")
+
+        with pytest.raises(Exception) as exc_ctx:
+            _retryable.retry_with(stop=tenacity.stop_after_attempt(2))()
+
+        assert exc_ctx.type is ValueError, "Should remap to specific exception type"
+
+    def test_retry_error_callback_should_be_preserved(self):
+        def return_text(retry_state):
+            return "Calling %s keeps raising errors after %s attempts" % (
+                retry_state.fn.__name__,
+                retry_state.attempt_number,
+            )
+
+        @retry(stop=tenacity.stop_after_attempt(10), retry_error_callback=return_text)
+        def _retryable():
+            raise Exception("raised for test purposes")
+
+        result = _retryable.retry_with(stop=tenacity.stop_after_attempt(5))()
+        assert result == "Calling _retryable keeps raising errors after 5 attempts"
+
+
 class TestBeforeAfterAttempts(unittest.TestCase):
     _attempt_number = 0