reqwest-middleware/reqwest-retry/tests/all/retry.rs

use paste::paste;
use reqwest::Client;
use reqwest::StatusCode;
use reqwest_middleware::ClientBuilder;
use reqwest_retry::{policies::ExponentialBackoff, RetryTransientMiddleware};
use std::sync::{
    atomic::{AtomicU32, Ordering},
    Arc,
};
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, Respond, ResponseTemplate};

use crate::helpers::SimpleServer;

pub struct RetryResponder(Arc<AtomicU32>, u32, u16);

impl RetryResponder {
    fn new(retries: u32, status_code: u16) -> Self {
        Self(Arc::new(AtomicU32::new(0)), retries, status_code)
    }
}

impl Respond for RetryResponder {
    fn respond(&self, _request: &wiremock::Request) -> ResponseTemplate {
        let mut retries = self.0.load(Ordering::SeqCst);
        retries += 1;
        self.0.store(retries, Ordering::SeqCst);

        if retries + 1 >= self.1 {
            ResponseTemplate::new(200)
        } else {
            ResponseTemplate::new(self.2)
        }
    }
}

macro_rules! assert_retry_succeeds_inner {
    ($x:tt, $name:ident, $status:expr, $retry:tt, $exact:tt, $responder:expr) => {
        #[tokio::test]
        async fn $name() {
            let server = MockServer::start().await;
            let retry_amount: u32 = $retry;
            Mock::given(method("GET"))
                .and(path("/foo"))
                .respond_with($responder)
                .expect($exact)
                .mount(&server)
                .await;

            let reqwest_client = Client::builder().build().unwrap();
            let client = ClientBuilder::new(reqwest_client)
                .with(RetryTransientMiddleware::new_with_policy(
                    ExponentialBackoff {
                        max_n_retries: retry_amount,
                        max_retry_interval: std::time::Duration::from_millis(30),
                        min_retry_interval: std::time::Duration::from_millis(100),
                        backoff_exponent: 2,
                    },
                ))
                .build();

            let resp = client
                .get(&format!("{}/foo", server.uri()))
                .send()
                .await
                .expect("call failed");

            assert_eq!(resp.status(), $status);
        }
    };
}

macro_rules! assert_retry_succeeds {
    ($x:tt, $status:expr) => {
        paste! {
            assert_retry_succeeds_inner!($x, [<assert_retry_succeds_on_ $x>], $status, 3, 2, RetryResponder::new(3 as u32, $x));
        }
    };
}

macro_rules! assert_no_retry {
    ($x:tt, $status:expr) => {
        paste! {
            assert_retry_succeeds_inner!($x, [<assert_no_retry_on_ $x>], $status, 1, 1, ResponseTemplate::new($x));
        }
    };
}

// 2xx.
assert_no_retry!(200, StatusCode::OK);
assert_no_retry!(201, StatusCode::CREATED);
assert_no_retry!(202, StatusCode::ACCEPTED);
assert_no_retry!(203, StatusCode::NON_AUTHORITATIVE_INFORMATION);
assert_no_retry!(204, StatusCode::NO_CONTENT);
assert_no_retry!(205, StatusCode::RESET_CONTENT);
assert_no_retry!(206, StatusCode::PARTIAL_CONTENT);
assert_no_retry!(207, StatusCode::MULTI_STATUS);
assert_no_retry!(226, StatusCode::IM_USED);

// 3xx.
assert_no_retry!(300, StatusCode::MULTIPLE_CHOICES);
assert_no_retry!(301, StatusCode::MOVED_PERMANENTLY);
assert_no_retry!(302, StatusCode::FOUND);
assert_no_retry!(303, StatusCode::SEE_OTHER);
assert_no_retry!(304, StatusCode::NOT_MODIFIED);
assert_no_retry!(307, StatusCode::TEMPORARY_REDIRECT);
assert_no_retry!(308, StatusCode::PERMANENT_REDIRECT);

// 5xx.
assert_retry_succeeds!(500, StatusCode::OK);
assert_retry_succeeds!(501, StatusCode::OK);
assert_retry_succeeds!(502, StatusCode::OK);
assert_retry_succeeds!(503, StatusCode::OK);
assert_retry_succeeds!(504, StatusCode::OK);
assert_retry_succeeds!(505, StatusCode::OK);
assert_retry_succeeds!(506, StatusCode::OK);
assert_retry_succeeds!(507, StatusCode::OK);
assert_retry_succeeds!(508, StatusCode::OK);
assert_retry_succeeds!(510, StatusCode::OK);
assert_retry_succeeds!(511, StatusCode::OK);
// 4xx.
assert_no_retry!(400, StatusCode::BAD_REQUEST);
assert_no_retry!(401, StatusCode::UNAUTHORIZED);
assert_no_retry!(402, StatusCode::PAYMENT_REQUIRED);
assert_no_retry!(403, StatusCode::FORBIDDEN);
assert_no_retry!(404, StatusCode::NOT_FOUND);
assert_no_retry!(405, StatusCode::METHOD_NOT_ALLOWED);
assert_no_retry!(406, StatusCode::NOT_ACCEPTABLE);
assert_no_retry!(407, StatusCode::PROXY_AUTHENTICATION_REQUIRED);
assert_retry_succeeds!(408, StatusCode::OK);
assert_no_retry!(409, StatusCode::CONFLICT);
assert_no_retry!(410, StatusCode::GONE);
assert_no_retry!(411, StatusCode::LENGTH_REQUIRED);
assert_no_retry!(412, StatusCode::PRECONDITION_FAILED);
assert_no_retry!(413, StatusCode::PAYLOAD_TOO_LARGE);
assert_no_retry!(414, StatusCode::URI_TOO_LONG);
assert_no_retry!(415, StatusCode::UNSUPPORTED_MEDIA_TYPE);
assert_no_retry!(416, StatusCode::RANGE_NOT_SATISFIABLE);
assert_no_retry!(417, StatusCode::EXPECTATION_FAILED);
assert_no_retry!(418, StatusCode::IM_A_TEAPOT);
assert_no_retry!(421, StatusCode::MISDIRECTED_REQUEST);
assert_no_retry!(422, StatusCode::UNPROCESSABLE_ENTITY);
assert_no_retry!(423, StatusCode::LOCKED);
assert_no_retry!(424, StatusCode::FAILED_DEPENDENCY);
assert_no_retry!(426, StatusCode::UPGRADE_REQUIRED);
assert_no_retry!(428, StatusCode::PRECONDITION_REQUIRED);
assert_retry_succeeds!(429, StatusCode::OK);
assert_no_retry!(431, StatusCode::REQUEST_HEADER_FIELDS_TOO_LARGE);
assert_no_retry!(451, StatusCode::UNAVAILABLE_FOR_LEGAL_REASONS);

// We assert that we cap retries at 10, which means that we will
// get 11 calls to the RetryResponder.
assert_retry_succeeds_inner!(
    500,
    assert_maximum_retries_is_not_exceeded,
    StatusCode::INTERNAL_SERVER_ERROR,
    100,
    11,
    RetryResponder::new(100_u32, 500)
);

pub struct RetryTimeoutResponder(Arc<AtomicU32>, u32, std::time::Duration);

impl RetryTimeoutResponder {
    fn new(retries: u32, initial_timeout: std::time::Duration) -> Self {
        Self(Arc::new(AtomicU32::new(0)), retries, initial_timeout)
    }
}

impl Respond for RetryTimeoutResponder {
    fn respond(&self, _request: &wiremock::Request) -> ResponseTemplate {
        let mut retries = self.0.load(Ordering::SeqCst);
        retries += 1;
        self.0.store(retries, Ordering::SeqCst);

        if retries + 1 >= self.1 {
            ResponseTemplate::new(200)
        } else {
            ResponseTemplate::new(500).set_delay(self.2)
        }
    }
}

#[tokio::test]
async fn assert_retry_on_request_timeout() {
    let server = MockServer::start().await;
    Mock::given(method("GET"))
        .and(path("/foo"))
        .respond_with(RetryTimeoutResponder::new(
            3,
            std::time::Duration::from_millis(1000),
        ))
        .expect(2)
        .mount(&server)
        .await;

    let reqwest_client = Client::builder().build().unwrap();
    let client = ClientBuilder::new(reqwest_client)
        .with(RetryTransientMiddleware::new_with_policy(
            ExponentialBackoff {
                max_n_retries: 3,
                max_retry_interval: std::time::Duration::from_millis(30),
                min_retry_interval: std::time::Duration::from_millis(100),
                backoff_exponent: 2,
            },
        ))
        .build();

    let resp = client
        .get(&format!("{}/foo", server.uri()))
        .timeout(std::time::Duration::from_millis(10))
        .send()
        .await
        .expect("call failed");

    assert_eq!(resp.status(), 200);
}

#[tokio::test]
async fn assert_retry_on_incomplete_message() {
    let raw_response = "HTTP/1.1 200"; // the full working response is: "HTTP/1.1 200 OK\r\n\r\n"

    let simple_server = SimpleServer::new("127.0.0.1", None, raw_response)
        .expect("Error when creating a simple server");
    tokio::spawn(simple_server.start());

    let reqwest_client = Client::builder().build().unwrap();
    let client = ClientBuilder::new(reqwest_client)
        .with(RetryTransientMiddleware::new_with_policy(
            ExponentialBackoff {
                max_n_retries: 3,
                max_retry_interval: std::time::Duration::from_millis(30),
                min_retry_interval: std::time::Duration::from_millis(100),
                backoff_exponent: 2,
            },
        ))
        .build();

    let resp = client
        .get(&format!("{}/foo", simple_server.uri()))
        .timeout(std::time::Duration::from_millis(10))
        .send()
        .await
        .expect("call failed");

    assert_eq!(resp.status(), 200);
}