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()7async 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