test: add Spawn::poll_until_idle (#8213)

tokio-test/src/task.rs

@@ -70,6 +70,9 @@ const IDLE: usize = 0;
 const WAKE: usize = 1;
 const SLEEP: usize = 2;
 
+/// Default maximum number of poll iterations in [`Spawn::poll_until_idle`].
+const POLL_UNTIL_IDLE_MAX_ITERATIONS: usize = 150;
+
 impl<T> Spawn<T> {
     /// Consumes `self` returning the inner value
     pub fn into_inner(self) -> T
@@ -123,6 +126,43 @@ impl<T: Future> Spawn<T> {
         let fut = self.future.as_mut();
         self.task.enter(|cx| fut.poll(cx))
     }
+
+    /// Polls the future until it is idle.
+    ///
+    /// A future is considered idle when it either completes, or returns
+    /// [`Poll::Pending`] without a pending wake notification.
+    ///
+    /// Unlike [`poll`](Self::poll), this method keeps polling while the future
+    /// returns [`Poll::Pending`] but has received a wake notification, advancing
+    /// the future as far as possible without waiting for external events.
+    ///
+    /// Polling is bounded to avoid infinite loops when a future wakes without
+    /// making progress.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the iteration limit is exceeded.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use tokio_test::task;
+    ///
+    /// let mut task = task::spawn(async { 42 });
+    ///
+    /// assert!(task.poll_until_idle().is_ready());
+    /// ```
+    pub fn poll_until_idle(&mut self) -> Poll<T::Output> {
+        for _ in 0..POLL_UNTIL_IDLE_MAX_ITERATIONS {
+            let result = self.poll();
+            if result.is_ready() || !self.is_woken() {
+                return result;
+            }
+        }
+        panic!(
+            "poll_until_idle exceeded {POLL_UNTIL_IDLE_MAX_ITERATIONS} iterations; future may be waking without making progress"
+        );
+    }
 }
 
 impl<T: Stream> Spawn<T> {

tokio-test/tests/task.rs

@@ -1,3 +1,4 @@
+use std::future::{pending, Future};
 use std::pin::Pin;
 use std::task::{Context, Poll};
 use tokio_stream::Stream;
@@ -23,3 +24,81 @@ fn test_spawn_stream_size_hint() {
     let spawn = task::spawn(SizedStream);
     assert_eq!(spawn.size_hint(), (100, Some(200)));
 }
+
+#[test]
+fn poll_until_idle_ready() {
+    let mut task = task::spawn(async { 42 });
+    assert_eq!(task.poll_until_idle(), Poll::Ready(42));
+}
+
+#[test]
+fn poll_until_idle_pending_not_woken() {
+    let mut task = task::spawn(pending::<()>());
+    assert!(task.poll_until_idle().is_pending());
+}
+
+struct WakeThenReady {
+    step: u8,
+}
+
+impl Future for WakeThenReady {
+    type Output = ();
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
+        match self.step {
+            0 => {
+                self.step = 1;
+                cx.waker().wake_by_ref();
+                Poll::Pending
+            }
+            _ => Poll::Ready(()),
+        }
+    }
+}
+
+#[test]
+fn poll_until_idle_advances_on_wake() {
+    let mut task = task::spawn(WakeThenReady { step: 0 });
+    assert!(task.poll_until_idle().is_ready());
+}
+
+struct WakeNTimes {
+    remaining: u8,
+}
+
+impl Future for WakeNTimes {
+    type Output = u8;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<u8> {
+        if self.remaining == 0 {
+            return Poll::Ready(0);
+        }
+        self.remaining -= 1;
+        cx.waker().wake_by_ref();
+        Poll::Pending
+    }
+}
+
+#[test]
+fn poll_until_idle_multiple_wakes() {
+    let mut task = task::spawn(WakeNTimes { remaining: 3 });
+    assert_eq!(task.poll_until_idle(), Poll::Ready(0));
+}
+
+struct WakeForever;
+
+impl Future for WakeForever {
+    type Output = ();
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
+        cx.waker().wake_by_ref();
+        Poll::Pending
+    }
+}
+
+#[test]
+#[should_panic(expected = "poll_until_idle exceeded 150 iterations")]
+fn poll_until_idle_panics_on_infinite_wake() {
+    let mut task = task::spawn(WakeForever);
+    let _ = task.poll_until_idle();
+}