1 use std::task::{Context, Poll};
2 use std::{future::Future, pin::Pin};
3 use tower::util::ServiceExt;
4 use tower_service::Service;
5 
6 #[tokio::test(flavor = "current_thread")]
service_driven_to_readiness()7 async fn service_driven_to_readiness() {
8     // This test ensures that `oneshot` will repeatedly call `poll_ready` until
9     // the service is ready.
10     let _t = super::support::trace_init();
11 
12     struct PollMeTwice {
13         ready: bool,
14     }
15     impl Service<()> for PollMeTwice {
16         type Error = ();
17         type Response = ();
18         type Future = Pin<
19             Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + Sync + 'static>,
20         >;
21 
22         fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), ()>> {
23             if self.ready {
24                 Poll::Ready(Ok(()))
25             } else {
26                 self.ready = true;
27                 cx.waker().wake_by_ref();
28                 Poll::Pending
29             }
30         }
31 
32         fn call(&mut self, _: ()) -> Self::Future {
33             assert!(self.ready, "service not driven to readiness!");
34             Box::pin(async { Ok(()) })
35         }
36     }
37 
38     let svc = PollMeTwice { ready: false };
39     svc.oneshot(()).await.unwrap();
40 }
41