Add Future::interleave_pending

Author: Thomasdezeeuw

futures-test/src/future/mod.rs

@@ -10,6 +10,8 @@ use futures_core::future::Future;
 use futures_executor;
 use std::thread;
 
+pub use crate::interleave_pending::InterleavePending;
+
 /// Additional combinators for testing futures.
 pub trait FutureTestExt: Future {
     /// Asserts that the given is not moved after being polled.
@@ -79,6 +81,34 @@ pub trait FutureTestExt: Future {
     {
         thread::spawn(|| futures_executor::block_on(self));
     }
+
+    /// Introduces an extra [`Poll::Pending`](futures_core::task::Poll::Pending)
+    /// in between each call to poll.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(async_await)]
+    /// use futures::task::Poll;
+    /// use futures::future::{self, Future};
+    /// use futures_test::task::noop_context;
+    /// use futures_test::future::FutureTestExt;
+    /// use pin_utils::pin_mut;
+    ///
+    /// let future = future::ready(1).interleave_pending();
+    /// pin_mut!(future);
+    ///
+    /// let mut cx = noop_context();
+    ///
+    /// assert_eq!(future.as_mut().poll(&mut cx), Poll::Pending);
+    /// assert_eq!(future.as_mut().poll(&mut cx), Poll::Ready(1));
+    /// ```
+    fn interleave_pending(self) -> InterleavePending<Self>
+    where
+        Self: Sized,
+    {
+        InterleavePending::new(self)
+    }
 }
 
 impl<Fut> FutureTestExt for Fut where Fut: Future {}

futures-test/src/interleave_pending.rs

@@ -1,3 +1,4 @@
+use futures_core::future::Future;
 use futures_core::stream::Stream;
 use futures_io::{self as io, AsyncBufRead, AsyncRead, AsyncWrite};
 use pin_utils::{unsafe_pinned, unsafe_unpinned};
@@ -63,6 +64,27 @@ impl<T> InterleavePending<T> {
     }
 }
 
+impl<Fut: Future> Future for InterleavePending<Fut> {
+    type Output = Fut::Output;
+
+    fn poll(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Self::Output> {
+        if *self.as_mut().pended() {
+            let next = self.as_mut().inner().poll(cx);
+            if next.is_ready() {
+                *self.pended() = false;
+            }
+            next
+        } else {
+            cx.waker().wake_by_ref();
+            *self.pended() = true;
+            Poll::Pending
+        }
+    }
+}
+
 impl<St: Stream> Stream for InterleavePending<St> {
     type Item = St::Item;