Add AsyncReadTestExt::limited

As a partner to AsyncWriteTestExt::limited_write.

Author: Thomasdezeeuw

futures-test/src/io/limited.rs

@@ -76,3 +76,27 @@ impl<W: AsyncWrite> AsyncWrite for Limited<W> {
         self.io().poll_close(cx)
     }
 }
+
+impl<R: AsyncRead> AsyncRead for Limited<R> {
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut [u8],
+    ) -> Poll<io::Result<usize>> {
+        let limit = cmp::min(*self.as_mut().limit(), buf.len());
+        self.io().poll_read(cx, &mut buf[..limit])
+    }
+}
+
+impl<R: AsyncBufRead> AsyncBufRead for Limited<R> {
+    fn poll_fill_buf<'a>(
+        self: Pin<&'a mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<io::Result<&'a [u8]>> {
+        self.io().poll_fill_buf(cx)
+    }
+
+    fn consume(self: Pin<&mut Self>, amount: usize) {
+        self.io().consume(amount)
+    }
+}

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

@@ -3,6 +3,7 @@
 use futures_io::AsyncRead;
 
 pub use super::interleave_pending::InterleavePending;
+pub use super::limited::Limited;
 
 /// Additional combinators for testing async readers.
 pub trait AsyncReadTestExt: AsyncRead {
@@ -72,6 +73,41 @@ pub trait AsyncReadTestExt: AsyncRead {
     {
         InterleavePending::new(self)
     }
+
+    /// Limit the number of bytes allowed to be read on each call to `poll_read`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(async_await)]
+    /// use futures::task::Poll;
+    /// use futures::io::AsyncRead;
+    /// use futures_test::task::noop_context;
+    /// use futures_test::io::AsyncReadTestExt;
+    /// use pin_utils::pin_mut;
+    ///
+    /// let reader = std::io::Cursor::new(&[1, 2, 3, 4, 5]).limited(2);
+    /// pin_mut!(reader);
+    ///
+    /// let mut cx = noop_context();
+    ///
+    /// let mut buf = [0; 10];
+    ///
+    /// assert_eq!(reader.as_mut().poll_read(&mut cx, &mut buf)?, Poll::Ready(2));
+    /// assert_eq!(&buf[..2], &[1, 2]);
+    /// assert_eq!(reader.as_mut().poll_read(&mut cx, &mut buf)?, Poll::Ready(2));
+    /// assert_eq!(&buf[..2], &[3, 4]);
+    /// assert_eq!(reader.as_mut().poll_read(&mut cx, &mut buf)?, Poll::Ready(1));
+    /// assert_eq!(&buf[..1], &[5]);
+    ///
+    /// # Ok::<(), std::io::Error>(())
+    /// ```
+    fn limited(self, limit: usize) -> Limited<Self>
+    where
+        Self: Sized,
+    {
+        Limited::new(self, limit)
+    }
 }
 
 impl<R> AsyncReadTestExt for R where R: AsyncRead {}