Unit test request retry python

Question

I'm trying to retry a request a couple times if the endpoint times out before returning results. Here's the code:

def retry_request(self, params, max_retries=3):
    for i in xrange(max_retries):
        try:
            response = requests.get(params)
            break
        except requests.exceptions.Timeout as e:
            raise e

I'd like to unit test the retries to show the retry logic works. Any thoughts?

Answer 1

I usually prefer not to make actual calls out to the internet in my tests: the remote service could be down or you might need to run your tests offline. Most importantly you want to run your tests FAST, and the network call can slow them down significantly.

I also want to make sure that the retry logic makes the retries I expect, and that it can actually succeeds eventually.

I tried to write a test doing that myself, but I struggled. I asked The Internet, but could not find anything doing what I wanted. I've digged into the magic world of urllib3 and finally got to the bottom of it, but it took me a while.

Since this post came up while searching, I'll leave my solution here for posterity, trying to save someone else the hours I've spent trying:

import urllib3
from http.client import HTTPMessage
from unittest.mock import ANY, Mock, patch, call

import requests


def request_with_retry(*args, **kwargs):
    session = requests.Session()
    adapter = requests.adapters.HTTPAdapter(max_retries=urllib3.Retry(
        raise_on_status=False,
        total=kwargs.pop("max_retries", 3),
        status_forcelist=[429, 500, 502, 503, 504],  # The HTTP response codes to retry on
        allowed_methods=["HEAD", "GET", "PUT", "DELETE", "OPTIONS"],  # The HTTP methods to retry on
    ))
    session.mount("https://", adapter)
    session.mount("http://", adapter)
    return session.request(*args, **kwargs)


@patch("urllib3.connectionpool.HTTPConnectionPool._get_conn")
def test_retry_request(getconn_mock):
    getconn_mock.return_value.getresponse.side_effect = [
        Mock(status=500, msg=HTTPMessage()),
        Mock(status=429, msg=HTTPMessage()),
        Mock(status=200, msg=HTTPMessage()),
    ]

    r = request_with_retry("GET", "http://any.url/testme", max_retries=2)
    r.raise_for_status()

    assert getconn_mock.return_value.request.mock_calls == [
        call("GET", "/testme", body=None, headers=ANY),
        call("GET", "/testme", body=None, headers=ANY),
        call("GET", "/testme", body=None, headers=ANY),
    ]

(Note: if you are making several calls to this method, then you might want to initialize the session object only once, not every time you make a request!)