zebra_test/
service_extensions.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
//! Extension traits for [`Service`] types to help with testing.

use std::task::Poll;

use futures::future::{BoxFuture, FutureExt};
use tower::{Service, ServiceExt};

/// An extension trait to check if a [`Service`] is immediately ready to be called.
pub trait IsReady<Request>: Service<Request> {
    /// Poll the [`Service`] once, and return true if it is immediately ready to be called.
    #[allow(clippy::wrong_self_convention)]
    fn is_ready(&mut self) -> BoxFuture<bool>;

    /// Poll the [`Service`] once, and return true if it is pending.
    #[allow(clippy::wrong_self_convention)]
    fn is_pending(&mut self) -> BoxFuture<bool>;

    /// Poll the [`Service`] once, and return true if it has failed.
    #[allow(clippy::wrong_self_convention)]
    fn is_failed(&mut self) -> BoxFuture<bool>;
}

impl<S, Request> IsReady<Request> for S
where
    S: Service<Request> + Send,
    Request: 'static,
{
    fn is_ready(&mut self) -> BoxFuture<bool> {
        async move {
            let ready_result = futures::poll!(self.ready());
            matches!(ready_result, Poll::Ready(Ok(_)))
        }
        .boxed()
    }

    fn is_pending(&mut self) -> BoxFuture<bool> {
        async move {
            let ready_result = futures::poll!(self.ready());
            ready_result.is_pending()
        }
        .boxed()
    }

    fn is_failed(&mut self) -> BoxFuture<bool> {
        async move {
            let ready_result = futures::poll!(self.ready());
            matches!(ready_result, Poll::Ready(Err(_)))
        }
        .boxed()
    }
}