Add LimitedWrite test adaptor for AsyncWrite

Author: Nemo157

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

@@ -0,0 +1,72 @@
+use futures_io::{self as io, AsyncWrite};
+use pin_utils::{unsafe_pinned, unsafe_unpinned};
+use std::{
+    cmp,
+    marker::Unpin,
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+/// Writer for the [`limited_write`](super::AsyncWriteTestExt::limited_write) method.
+#[derive(Debug)]
+pub struct LimitedWrite<W: AsyncWrite> {
+    writer: W,
+    limit: usize,
+}
+
+impl<W: AsyncWrite + Unpin> Unpin for LimitedWrite<W> {}
+
+impl<W: AsyncWrite> LimitedWrite<W> {
+    unsafe_pinned!(writer: W);
+    unsafe_unpinned!(limit: usize);
+
+    pub(crate) fn new(writer: W, limit: usize) -> LimitedWrite<W> {
+        LimitedWrite { writer, limit }
+    }
+
+    /// 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.writer()
+    }
+
+    /// Consumes this adaptor returning the underlying writer.
+    pub fn into_inner(self) -> W {
+        self.writer
+    }
+}
+
+impl<W: AsyncWrite> AsyncWrite for LimitedWrite<W> {
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        let limit = *self.as_mut().limit();
+        self.writer()
+            .poll_write(cx, &buf[..cmp::min(limit, buf.len())])
+    }
+
+    fn poll_flush(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<io::Result<()>> {
+        self.writer().poll_flush(cx)
+    }
+
+    fn poll_close(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<io::Result<()>> {
+        self.writer().poll_close(cx)
+    }
+}

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

@@ -5,6 +5,9 @@ use futures_io::AsyncWrite;
 mod interleave_pending_write;
 pub use self::interleave_pending_write::InterleavePendingWrite;
 
+mod limited_write;
+pub use self::limited_write::LimitedWrite;
+
 /// Additional combinators for testing async writers.
 pub trait AsyncWriteTestExt: AsyncWrite {
     /// Introduces an extra [`Poll::Pending`](futures_core::task::Poll::Pending)
@@ -48,6 +51,40 @@ pub trait AsyncWriteTestExt: AsyncWrite {
     {
         InterleavePendingWrite::new(self)
     }
+
+    /// Limit the number of bytes allowed to be written on each call to `poll_write`.
+    ///
+    /// # 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]).limited_write(2);
+    /// pin_mut!(writer);
+    ///
+    /// let mut cx = noop_context();
+    ///
+    /// 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])?, Poll::Ready(1));
+    /// assert_eq!(writer.get_ref().get_ref(), &[1, 2, 3, 0]);
+    /// assert_eq!(writer.as_mut().poll_write(&mut cx, &[4, 5])?, Poll::Ready(1));
+    /// assert_eq!(writer.get_ref().get_ref(), &[1, 2, 3, 4]);
+    /// assert_eq!(writer.as_mut().poll_write(&mut cx, &[5])?, Poll::Ready(0));
+    ///
+    /// # Ok::<(), std::io::Error>(())
+    /// ```
+    fn limited_write(self, limit: usize) -> LimitedWrite<Self>
+    where
+        Self: Sized,
+    {
+        LimitedWrite::new(self, limit)
+    }
 }
 
 impl<W> AsyncWriteTestExt for W where W: AsyncWrite {}