Add InterleavePending adaptor for AsyncWrite

Author: Nemo157

futures-test/src/io/mod.rs

@@ -3,3 +3,5 @@
 pub mod read;
 pub use read::AsyncReadTestExt;
 
+pub mod write;
+pub use write::AsyncWriteTestExt;

futures-test/src/io/write/interleave_pending_write.rs

@@ -0,0 +1,113 @@
+use futures_io::{self as io, AsyncWrite};
+use pin_utils::{unsafe_pinned, unsafe_unpinned};
+use std::{
+    marker::Unpin,
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+/// Writer for the [`interleave_pending_write`](super::AsyncWriteTestExt::interleave_pending_write)
+/// method.
+#[derive(Debug)]
+pub struct InterleavePendingWrite<W: AsyncWrite> {
+    writer: W,
+    pended: bool,
+}
+
+impl<W: AsyncWrite + Unpin> Unpin for InterleavePendingWrite<W> {}
+
+impl<W: AsyncWrite> InterleavePendingWrite<W> {
+    unsafe_pinned!(writer: W);
+    unsafe_unpinned!(pended: bool);
+
+    pub(crate) fn new(writer: W) -> Self {
+        Self {
+            writer,
+            pended: false,
+        }
+    }
+
+    /// Acquires a reference to the underlying writer that this adaptor is wrapping.
+    pub fn get_ref(&self) -> &W {
+        &self.writer
+    }
+
+    /// Acquires a mutable reference to the underlying writer that this adaptor is wrapping.
+    pub fn get_mut(&mut self) -> &mut W {
+        &mut self.writer
+    }
+
+    /// Acquires a pinned mutable reference to the underlying writer that this adaptor is wrapping.
+    pub fn get_pin_mut<'a>(self: Pin<&'a mut Self>) -> Pin<&'a mut W> {
+        self.project().0
+    }
+
+    /// Consumes this adaptor returning the underlying writer.
+    pub fn into_inner(self) -> W {
+        self.writer
+    }
+
+    fn project<'a>(self: Pin<&'a mut Self>) -> (Pin<&'a mut W>, &'a mut bool) {
+        unsafe {
+            let this = self.get_unchecked_mut();
+            (Pin::new_unchecked(&mut this.writer), &mut this.pended)
+        }
+    }
+}
+
+impl<W: AsyncWrite> AsyncWrite for InterleavePendingWrite<W> {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        let (writer, pended) = self.project();
+        if *pended {
+            let next = writer.poll_write(cx, buf);
+            if next.is_ready() {
+                *pended = false;
+            }
+            next
+        } else {
+            cx.waker().wake_by_ref();
+            *pended = true;
+            Poll::Pending
+        }
+    }
+
+    fn poll_flush(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<io::Result<()>> {
+        let (writer, pended) = self.project();
+        if *pended {
+            let next = writer.poll_flush(cx);
+            if next.is_ready() {
+                *pended = false;
+            }
+            next
+        } else {
+            cx.waker().wake_by_ref();
+            *pended = true;
+            Poll::Pending
+        }
+    }
+
+    fn poll_close(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<io::Result<()>> {
+        let (writer, pended) = self.project();
+        if *pended {
+            let next = writer.poll_close(cx);
+            if next.is_ready() {
+                *pended = false;
+            }
+            next
+        } else {
+            cx.waker().wake_by_ref();
+            *pended = true;
+            Poll::Pending
+        }
+    }
+}

futures-test/src/io/write/mod.rs

@@ -0,0 +1,53 @@
+//! Additional combinators for testing async writers.
+
+use futures_io::AsyncWrite;
+
+mod interleave_pending_write;
+pub use self::interleave_pending_write::InterleavePendingWrite;
+
+/// Additional combinators for testing async writers.
+pub trait AsyncWriteTestExt: AsyncWrite {
+    /// Introduces an extra [`Poll::Pending`](futures_core::task::Poll::Pending)
+    /// in between each operation on the writer.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(async_await)]
+    /// use futures::task::Poll;
+    /// use futures::io::AsyncWrite;
+    /// use futures_test::task::noop_context;
+    /// use futures_test::io::AsyncWriteTestExt;
+    /// use pin_utils::pin_mut;
+    ///
+    /// let writer = std::io::Cursor::new([0u8; 4]).interleave_pending_write();
+    /// pin_mut!(writer);
+    ///
+    /// let mut cx = noop_context();
+    ///
+    /// assert_eq!(writer.as_mut().poll_write(&mut cx, &[1, 2])?, Poll::Pending);
+    /// assert_eq!(writer.as_mut().poll_write(&mut cx, &[1, 2])?, Poll::Ready(2));
+    /// assert_eq!(writer.get_ref().get_ref(), &[1, 2, 0, 0]);
+    /// assert_eq!(writer.as_mut().poll_write(&mut cx, &[3, 4])?, Poll::Pending);
+    /// assert_eq!(writer.as_mut().poll_write(&mut cx, &[3, 4])?, Poll::Ready(2));
+    /// assert_eq!(writer.get_ref().get_ref(), &[1, 2, 3, 4]);
+    /// assert_eq!(writer.as_mut().poll_write(&mut cx, &[5, 6])?, Poll::Pending);
+    /// assert_eq!(writer.as_mut().poll_write(&mut cx, &[5, 6])?, Poll::Ready(0));
+    ///
+    /// assert_eq!(writer.as_mut().poll_flush(&mut cx)?, Poll::Pending);
+    /// assert_eq!(writer.as_mut().poll_flush(&mut cx)?, Poll::Ready(()));
+    ///
+    /// assert_eq!(writer.as_mut().poll_close(&mut cx)?, Poll::Pending);
+    /// assert_eq!(writer.as_mut().poll_close(&mut cx)?, Poll::Ready(()));
+    ///
+    /// # Ok::<(), std::io::Error>(())
+    /// ```
+    fn interleave_pending_write(self) -> InterleavePendingWrite<Self>
+    where
+        Self: Sized,
+    {
+        InterleavePendingWrite::new(self)
+    }
+}
+
+impl<W> AsyncWriteTestExt for W where W: AsyncWrite {}