python retry with tenacity, disable `wait` for unittest

Question

I am using the tenacity library to use its @retry decorator.

I am using this to make a function which makes a HTTP-request "repeat" multiple times in case of failure.

Here is a simple code snippet:

@retry(stop=stop_after_attempt(7), wait=wait_random_exponential(multiplier=1, max=60))
def func():
   ...
   requests.post(...)

The function uses the tenacity wait-argument to wait some time in between calls.

The function together with the @retry-decorator seems to work fine.

But I also have a unit-test which checks that the function gets called indeed 7 times in case of a failure. This test takes a lot of time because of this wait in beetween tries.

Is it possible to somehow disable the wait-time only in the unit-test?

Answer 1

The solution came from the maintainer of tenacity himself in this Github issue: https://github.com/jd/tenacity/issues/106

You can simply change the wait function temporarily for your unit test:

from tenacity import wait_none

func.retry.wait = wait_none()

Answer 2

Thanks to discussion here, I found an elegant way to do this based on code from @steveb:

from tenacity import retry, stop_after_attempt, wait_exponential


@retry(reraise=True, stop=stop_after_attempt(5), wait=wait_exponential(multiplier=1, min=4, max=10))
def do_something_flaky(succeed):
    print('Doing something flaky')
    if not succeed:
        print('Failed!')
        raise Exception('Failed!')

And tests:

from unittest import TestCase, mock, skip
from main import do_something_flaky


class TestFlakyRetry(TestCase):
    def test_succeeds_instantly(self):
        try:
            do_something_flaky(True)
        except Exception:
            self.fail('Flaky function should not have failed.')

    def test_raises_exception_immediately_with_direct_mocking(self):
        do_something_flaky.retry.sleep = mock.Mock()
        with self.assertRaises(Exception):
            do_something_flaky(False)

    def test_raises_exception_immediately_with_indirect_mocking(self):
        with mock.patch('main.do_something_flaky.retry.sleep'):
            with self.assertRaises(Exception):
                do_something_flaky(False)

    @skip('Takes way too long to run!')
    def test_raises_exception_after_full_retry_period(self):
        with self.assertRaises(Exception):
            do_something_flaky(False)

Answer 3

After reading the thread in tenacity repo (thanks @DanEEStar for starting it!), I came up with the following code:

@retry(
    stop=stop_after_delay(20.0),
    wait=wait_incrementing(
        start=0,
        increment=0.25,
    ),
    retry=retry_if_exception_type(SomeExpectedException),
    reraise=True,
)
def func() -> None:
    raise SomeExpectedException()


def test_func_should_retry(monkeypatch: MonkeyPatch) -> None:
    # Use monkeypatch to patch retry behavior.
    # It will automatically revert patches when test finishes.
    # Also, it doesn't create nested blocks as `unittest.mock.patch` does.

    # Originally, it was `stop_after_delay` but the test could be
    # unreasonably slow this way. After all, I don't care so much
    # about which policy is applied exactly in this test.
    monkeypatch.setattr(
        func.retry, "stop", stop_after_attempt(3)
    )

    # Disable pauses between retries.
    monkeypatch.setattr(func.retry, "wait", wait_none())

    with pytest.raises(SomeExpectedException):
        func()

    # Ensure that there were retries.
    stats: Dict[str, Any] = func.retry.statistics
    assert "attempt_number" in stats
    assert stats["attempt_number"] == 3

I use pytest-specific features in this test. Probably, it could be useful as an example for someone, at least for future me.

Answer 4

mock the base class wait func with:

mock.patch('tenacity.BaseRetrying.wait', side_effect=lambda *args, **kwargs: 0)

it always not wait