tracing_test/
lib.rs

Help
1use std::{
2    pin::Pin,
3    task::{Context, Poll},
4};
5
6#[allow(missing_docs)]
7pub struct PollN<T, E> {
8    and_return: Option<Result<T, E>>,
9    finish_at: usize,
10    polls: usize,
11}
12
13impl<T, E> std::future::Future for PollN<T, E>
14where
15    T: Unpin,
16    E: Unpin,
17{
18    type Output = Result<T, E>;
19    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
20        let this = self.get_mut();
21
22        this.polls += 1;
23        if this.polls == this.finish_at {
24            let value = this.and_return.take().expect("polled after ready");
25
26            Poll::Ready(value)
27        } else {
28            cx.waker().wake_by_ref();
29            Poll::Pending
30        }
31    }
32}
33
34impl PollN<(), ()> {
35    pub fn new_ok(finish_at: usize) -> Self {
36        Self {
37            and_return: Some(Ok(())),
38            finish_at,
39            polls: 0,
40        }
41    }
42
43    pub fn new_err(finish_at: usize) -> Self {
44        Self {
45            and_return: Some(Err(())),
46            finish_at,
47            polls: 0,
48        }
49    }
50}
51
52pub fn block_on_future<F>(future: F) -> F::Output
53where
54    F: std::future::Future,
55{
56    use tokio_test::task;
57
58    let mut task = task::spawn(future);
59    loop {
60        if let Poll::Ready(v) = task.poll() {
61            break v;
62        }
63    }
64}
tracing_test/lib.rs

tracing_test/
lib.rs
