2019-09-09 12:10:46 -07:00
|
|
|
use futures_util::{future, pin_mut};
|
|
|
|
use std::future::Future;
|
|
|
|
use tokio_test::{assert_pending, assert_ready_err, assert_ready_ok, task};
|
2018-08-28 11:05:15 -07:00
|
|
|
use tower_retry::Policy;
|
|
|
|
use tower_service::Service;
|
2019-04-08 20:11:09 -07:00
|
|
|
use tower_test::{assert_request_eq, mock};
|
2018-08-28 11:05:15 -07:00
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn retry_errors() {
|
2019-09-09 12:10:46 -07:00
|
|
|
task::mock(|cx| {
|
|
|
|
let (mut service, handle) = new_service(RetryErrors);
|
|
|
|
pin_mut!(handle);
|
2018-08-28 11:05:15 -07:00
|
|
|
|
2019-09-09 12:10:46 -07:00
|
|
|
assert_ready_ok!(service.poll_ready(cx));
|
2018-08-28 11:05:15 -07:00
|
|
|
|
2019-09-09 12:10:46 -07:00
|
|
|
let fut = service.call("hello");
|
|
|
|
pin_mut!(fut);
|
2018-08-28 11:05:15 -07:00
|
|
|
|
2019-09-09 12:10:46 -07:00
|
|
|
assert_request_eq!(handle.as_mut(), "hello").send_error("retry me");
|
2018-08-28 11:05:15 -07:00
|
|
|
|
2019-09-09 12:10:46 -07:00
|
|
|
assert_pending!(fut.as_mut().poll(cx));
|
2018-08-28 11:05:15 -07:00
|
|
|
|
2019-09-09 12:10:46 -07:00
|
|
|
assert_request_eq!(handle.as_mut(), "hello").send_response("world");
|
|
|
|
|
|
|
|
assert_ready_ok!(fut.poll(cx), "world");
|
|
|
|
});
|
2018-08-28 11:05:15 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn retry_limit() {
|
2019-09-09 12:10:46 -07:00
|
|
|
task::mock(|cx| {
|
|
|
|
let (mut service, handle) = new_service(Limit(2));
|
|
|
|
pin_mut!(handle);
|
|
|
|
|
|
|
|
assert_ready_ok!(service.poll_ready(cx));
|
2018-08-28 11:05:15 -07:00
|
|
|
|
2019-09-09 12:10:46 -07:00
|
|
|
let fut = service.call("hello");
|
|
|
|
pin_mut!(fut);
|
2018-08-28 11:05:15 -07:00
|
|
|
|
2019-09-09 12:10:46 -07:00
|
|
|
assert_request_eq!(handle.as_mut(), "hello").send_error("retry 1");
|
|
|
|
assert_pending!(fut.as_mut().poll(cx));
|
2018-08-28 11:05:15 -07:00
|
|
|
|
2019-09-09 12:10:46 -07:00
|
|
|
assert_request_eq!(handle.as_mut(), "hello").send_error("retry 2");
|
|
|
|
assert_pending!(fut.as_mut().poll(cx));
|
2018-08-28 11:05:15 -07:00
|
|
|
|
2019-09-09 12:10:46 -07:00
|
|
|
assert_request_eq!(handle.as_mut(), "hello").send_error("retry 3");
|
|
|
|
assert_eq!(assert_ready_err!(fut.poll(cx)).to_string(), "retry 3");
|
|
|
|
});
|
2018-08-28 11:05:15 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn retry_error_inspection() {
|
2019-09-09 12:10:46 -07:00
|
|
|
task::mock(|cx| {
|
|
|
|
let (mut service, handle) = new_service(UnlessErr("reject"));
|
|
|
|
pin_mut!(handle);
|
2018-08-28 11:05:15 -07:00
|
|
|
|
2019-09-09 12:10:46 -07:00
|
|
|
assert_ready_ok!(service.poll_ready(cx));
|
|
|
|
let fut = service.call("hello");
|
|
|
|
pin_mut!(fut);
|
2018-08-28 11:05:15 -07:00
|
|
|
|
2019-09-09 12:10:46 -07:00
|
|
|
assert_request_eq!(handle.as_mut(), "hello").send_error("retry 1");
|
|
|
|
assert_pending!(fut.as_mut().poll(cx));
|
2018-08-28 11:05:15 -07:00
|
|
|
|
2019-09-09 12:10:46 -07:00
|
|
|
assert_request_eq!(handle.as_mut(), "hello").send_error("reject");
|
|
|
|
assert_eq!(assert_ready_err!(fut.poll(cx)).to_string(), "reject");
|
|
|
|
});
|
2018-08-28 11:05:15 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn retry_cannot_clone_request() {
|
2019-09-09 12:10:46 -07:00
|
|
|
task::mock(|cx| {
|
|
|
|
let (mut service, handle) = new_service(CannotClone);
|
|
|
|
pin_mut!(handle);
|
2018-08-28 11:05:15 -07:00
|
|
|
|
2019-09-09 12:10:46 -07:00
|
|
|
assert_ready_ok!(service.poll_ready(cx));
|
|
|
|
let fut = service.call("hello");
|
|
|
|
pin_mut!(fut);
|
2018-08-28 11:05:15 -07:00
|
|
|
|
2019-09-09 12:10:46 -07:00
|
|
|
assert_request_eq!(handle, "hello").send_error("retry 1");
|
|
|
|
assert_eq!(assert_ready_err!(fut.poll(cx)).to_string(), "retry 1");
|
|
|
|
});
|
2018-08-28 11:05:15 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn success_with_cannot_clone() {
|
2019-09-09 12:10:46 -07:00
|
|
|
task::mock(|cx| {
|
|
|
|
// Even though the request couldn't be cloned, if the first request succeeds,
|
|
|
|
// it should succeed overall.
|
|
|
|
let (mut service, handle) = new_service(CannotClone);
|
|
|
|
pin_mut!(handle);
|
|
|
|
|
|
|
|
assert_ready_ok!(service.poll_ready(cx));
|
|
|
|
let fut = service.call("hello");
|
|
|
|
pin_mut!(fut);
|
|
|
|
|
|
|
|
assert_request_eq!(handle, "hello").send_response("world");
|
|
|
|
assert_ready_ok!(fut.poll(cx), "world");
|
|
|
|
});
|
2018-08-28 11:05:15 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
type Req = &'static str;
|
|
|
|
type Res = &'static str;
|
|
|
|
type InnerError = &'static str;
|
2019-04-08 20:11:09 -07:00
|
|
|
type Error = Box<dyn std::error::Error + Send + Sync>;
|
2019-04-07 20:42:18 -07:00
|
|
|
type Mock = mock::Mock<Req, Res>;
|
|
|
|
type Handle = mock::Handle<Req, Res>;
|
2018-08-28 11:05:15 -07:00
|
|
|
|
|
|
|
#[derive(Clone)]
|
|
|
|
struct RetryErrors;
|
|
|
|
|
|
|
|
impl Policy<Req, Res, Error> for RetryErrors {
|
2019-09-09 12:10:46 -07:00
|
|
|
type Future = future::Ready<Self>;
|
2018-08-28 11:05:15 -07:00
|
|
|
fn retry(&self, _: &Req, result: Result<&Res, &Error>) -> Option<Self::Future> {
|
|
|
|
if result.is_err() {
|
2019-09-09 12:10:46 -07:00
|
|
|
Some(future::ready(RetryErrors))
|
2018-08-28 11:05:15 -07:00
|
|
|
} else {
|
|
|
|
None
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
fn clone_request(&self, req: &Req) -> Option<Req> {
|
|
|
|
Some(*req)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
#[derive(Clone)]
|
|
|
|
struct Limit(usize);
|
|
|
|
|
|
|
|
impl Policy<Req, Res, Error> for Limit {
|
2019-09-09 12:10:46 -07:00
|
|
|
type Future = future::Ready<Self>;
|
2018-08-28 11:05:15 -07:00
|
|
|
fn retry(&self, _: &Req, result: Result<&Res, &Error>) -> Option<Self::Future> {
|
|
|
|
if result.is_err() && self.0 > 0 {
|
2019-09-09 12:10:46 -07:00
|
|
|
Some(future::ready(Limit(self.0 - 1)))
|
2018-08-28 11:05:15 -07:00
|
|
|
} else {
|
|
|
|
None
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
fn clone_request(&self, req: &Req) -> Option<Req> {
|
|
|
|
Some(*req)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
#[derive(Clone)]
|
|
|
|
struct UnlessErr(InnerError);
|
|
|
|
|
|
|
|
impl Policy<Req, Res, Error> for UnlessErr {
|
2019-09-09 12:10:46 -07:00
|
|
|
type Future = future::Ready<Self>;
|
2018-08-28 11:05:15 -07:00
|
|
|
fn retry(&self, _: &Req, result: Result<&Res, &Error>) -> Option<Self::Future> {
|
2019-02-11 15:11:31 -08:00
|
|
|
result.err().and_then(|err| {
|
2019-03-01 13:56:57 -08:00
|
|
|
if err.to_string() != self.0 {
|
2019-09-09 12:10:46 -07:00
|
|
|
Some(future::ready(self.clone()))
|
2019-02-11 15:11:31 -08:00
|
|
|
} else {
|
|
|
|
None
|
|
|
|
}
|
|
|
|
})
|
2018-08-28 11:05:15 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
fn clone_request(&self, req: &Req) -> Option<Req> {
|
|
|
|
Some(*req)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
#[derive(Clone)]
|
|
|
|
struct CannotClone;
|
|
|
|
|
|
|
|
impl Policy<Req, Res, Error> for CannotClone {
|
2019-09-09 12:10:46 -07:00
|
|
|
type Future = future::Ready<Self>;
|
2018-08-28 11:05:15 -07:00
|
|
|
fn retry(&self, _: &Req, _: Result<&Res, &Error>) -> Option<Self::Future> {
|
|
|
|
unreachable!("retry cannot be called since request isn't cloned");
|
|
|
|
}
|
|
|
|
|
|
|
|
fn clone_request(&self, _req: &Req) -> Option<Req> {
|
|
|
|
None
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2019-02-11 15:11:31 -08:00
|
|
|
fn new_service<P: Policy<Req, Res, Error> + Clone>(
|
|
|
|
policy: P,
|
|
|
|
) -> (tower_retry::Retry<P, Mock>, Handle) {
|
2019-04-07 20:42:18 -07:00
|
|
|
let (service, handle) = mock::pair();
|
2018-08-28 11:05:15 -07:00
|
|
|
let service = tower_retry::Retry::new(policy, service);
|
|
|
|
(service, handle)
|
|
|
|
}
|