test: add Spawn::poll_until_idle (#8213)

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();
+}