// this is the only way to make the path attribute play nice with `in // crate::support`... #[allow(clippy::module_inception)] #[path = "../../tracing/tests/support/mod.rs"] mod support; use std::{ future::Future, pin::Pin, task::{Context, Poll}, }; pub use support::*; use tokio_test::task; pub struct PollN { and_return: Option>, finish_at: usize, polls: usize, } impl std::future::Future for PollN where T: Unpin, E: Unpin, { type Output = Result; fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll { let this = self.get_mut(); this.polls += 1; if this.polls == this.finish_at { let value = this.and_return.take().expect("polled after ready"); Poll::Ready(value) } else { cx.waker().wake_by_ref(); Poll::Pending } } } impl PollN<(), ()> { pub fn new_ok(finish_at: usize) -> Self { Self { and_return: Some(Ok(())), finish_at, polls: 0, } } pub fn new_err(finish_at: usize) -> Self { Self { and_return: Some(Err(())), finish_at, polls: 0, } } } pub fn block_on_future(future: F) -> F::Output where F: Future, { let mut task = task::spawn(future); loop { if let Poll::Ready(v) = task.poll() { break v; } } }