Search code examples
pythonpython-unittestretry-logic

How to test Retry attempts in python using the request library


I have the following:

from requests.adapters import HTTPAdapter
from requests.packages.urllib3.util.retry import Retry

retry_strategy = Retry(
    total=3,
    status_forcelist=[429, 500, 502, 503, 504],
    method_whitelist=["HEAD", "GET", "OPTIONS"]
)
adapter = HTTPAdapter(max_retries=retry_strategy)
http = requests.Session()
http.mount("https://", adapter)
http.mount("http://", adapter)

response = http.get("some_endpoint")

How can unit testing that effectively I'm retrying N times (in this case, 3)?


Solution

  • Wondered the same and I could not find any suggestion, not even on Stack Overflow. But I really wanted to make sure it was working as I expected...

    It took me a while to figure out where and how to patch urllib3 to test it properly, but this is the best solution I came up with so far, saving you a few hours of research:

    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!)