feat(future_tests): add extension trait for manual future polling

[Vaiz]

I'd like to propose adding a small extension trait that lets you poll a future a fixed number of times and assert Ready/Pending directly, with no block_on or executor.

Motivation

A large fraction of good unit tests don't do any real async calls and use stubs that return a value immediately or after a certain number of polls. Most of these tests don't require a runtime at all. Such an extension can make tests more deterministic and avoid any overhead from starting a runtime or using block_on.

Proposed API

• poll_once(&mut fut) -> Poll<T> - exactly one poll of an Unpin future, using a no-op waker.
• FutureTestExt (blanket-impl'd for all Future), taking self by value and pinning internally so it works on !Unpin async fn futures with no manual pin!:
  • unwrap_ready() - poll once, return output or panic if pending.
  • unwrap_pending() - poll once, assert pending.
  • unwrap_ready_within(max_polls) - poll up to max_polls, return first ready.
  • unwrap_ready_after(n_pending) - assert pending exactly n_pending times, then ready.
  • unwrap_pending_for(n) - assert pending for n polls.

Code

pub fn poll_once<F>(future: &mut F) -> Poll<F::Output>
where
    F: Future + Unpin,
{
    let waker = std::task::Waker::noop();
    let mut cx = Context::from_waker(waker);
    let mut future = std::pin::pin!(future);
    future.as_mut().poll(&mut cx)
}

pub trait FutureTestExt: Future + Sized {
    fn unwrap_ready(self) -> Self::Output {
        let mut fut = std::pin::pin!(self);
        match poll_once(&mut fut) {
            Poll::Ready(value) => value,
            Poll::Pending => {
                panic!("expected future to be Ready after one poll, but it was Pending")
            }
        }
    }

    fn unwrap_pending(self) {
        let mut fut = std::pin::pin!(self);
        assert!(
            poll_once(&mut fut).is_pending(),
            "expected future to be Pending after one poll, but it was Ready"
        );
    }

    fn unwrap_ready_within(self, max_polls: usize) -> Self::Output {
        let mut fut = std::pin::pin!(self);
        for _ in 0..max_polls {
            if let Poll::Ready(value) = poll_once(&mut fut) {
                return value;
            }
        }
        panic!("expected future to be Ready within {max_polls} polls, but it was still Pending");
    }

    fn unwrap_ready_after(self, n_pending: usize) -> Self::Output
    where
        Self::Output: Debug,
    {
        let mut fut = std::pin::pin!(self);
        for i in 0..n_pending {
            match poll_once(&mut fut) {
                Poll::Pending => {}
                Poll::Ready(value) => panic!(
                    "expected Pending, but got Ready({value:?}) after {} polls",
                    i + 1
                ),
            }
        }

        match poll_once(&mut fut) {
            Poll::Ready(value) => value,
            Poll::Pending => {
                panic!(
                    "expected Ready after {} polls, but got Pending",
                    n_pending + 1
                )
            }
        }
    }

    fn unwrap_pending_for(self, n: usize) {
        let mut fut = std::pin::pin!(self);
        for _ in 0..n {
            assert!(
                matches!(poll_once(&mut fut), Poll::Pending),
                "expected future to stay Pending for {n} polls, but it was Ready"
            );
        }
    }
}

impl<F: Future> FutureTestExt for F {}