Add test utility that verifies an AsyncWrite is closed correctly

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

@@ -1,9 +1,12 @@
 //! Additional combinators for testing async writers.
 
+mod track_closed;
+
 use futures_io::AsyncWrite;
 
 pub use super::limited::Limited;
 pub use crate::interleave_pending::InterleavePending;
+pub use track_closed::TrackClosed;
 
 /// Additional combinators for testing async writers.
 pub trait AsyncWriteTestExt: AsyncWrite {
@@ -80,6 +83,45 @@ pub trait AsyncWriteTestExt: AsyncWrite {
     {
         Limited::new(self, limit)
     }
+
+    /// Track whether this stream has been closed and errors if it is used after closing.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # futures::executor::block_on(async {
+    /// use futures::io::{AsyncWriteExt, Cursor};
+    /// use futures_test::io::AsyncWriteTestExt;
+    ///
+    /// let mut writer = Cursor::new(vec![0u8; 4]).track_closed();
+    ///
+    /// writer.write_all(&[1, 2]).await?;
+    /// assert!(!writer.is_closed());
+    /// writer.close().await?;
+    /// assert!(writer.is_closed());
+    ///
+    /// # Ok::<(), std::io::Error>(()) })?;
+    /// # Ok::<(), std::io::Error>(())
+    /// ```
+    ///
+    /// ```
+    /// # futures::executor::block_on(async {
+    /// use futures::io::{AsyncWriteExt, Cursor};
+    /// use futures_test::io::AsyncWriteTestExt;
+    ///
+    /// let mut writer = Cursor::new(vec![0u8; 4]).track_closed();
+    ///
+    /// writer.close().await?;
+    /// assert!(writer.write_all(&[1, 2]).await.is_err());
+    /// # Ok::<(), std::io::Error>(()) })?;
+    /// # Ok::<(), std::io::Error>(())
+    /// ```
+    fn track_closed(self) -> TrackClosed<Self>
+    where
+        Self: Sized,
+    {
+        TrackClosed::new(self)
+    }
 }
 
 impl<W> AsyncWriteTestExt for W where W: AsyncWrite {}

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

@@ -0,0 +1,110 @@
+use futures_io::AsyncWrite;
+use std::{
+    io::{Error, ErrorKind, IoSlice, Result},
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+/// I/O wrapper that tracks whether it has been closed.
+///
+/// See the [`track_closed`] method for more details.
+///
+/// [`track_closed`]: super::AsyncWriteTestExt::limited_write
+#[pin_project::pin_project]
+#[derive(Debug)]
+pub struct TrackClosed<W> {
+    #[pin]
+    inner: W,
+    closed: bool,
+}
+
+impl<W> TrackClosed<W> {
+    pub(crate) fn new(inner: W) -> TrackClosed<W> {
+        TrackClosed {
+            inner,
+            closed: false,
+        }
+    }
+
+    /// Check whether this stream has been closed.
+    pub fn is_closed(&self) -> bool {
+        self.closed
+    }
+
+    /// Acquires a reference to the underlying I/O object that this adaptor is
+    /// wrapping.
+    pub fn get_ref(&self) -> &W {
+        &self.inner
+    }
+
+    /// Acquires a mutable reference to the underlying I/O object that this
+    /// adaptor is wrapping.
+    pub fn get_mut(&mut self) -> &mut W {
+        &mut self.inner
+    }
+
+    /// Acquires a pinned mutable reference to the underlying I/O object that
+    /// this adaptor is wrapping.
+    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut W> {
+        self.project().inner
+    }
+
+    /// Consumes this adaptor returning the underlying I/O object.
+    pub fn into_inner(self) -> W {
+        self.inner
+    }
+}
+
+impl<W: AsyncWrite> AsyncWrite for TrackClosed<W> {
+    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize>> {
+        if self.is_closed() {
+            return Poll::Ready(Err(Error::new(
+                ErrorKind::Other,
+                "Attempted to write after stream was closed",
+            )));
+        }
+        self.project().inner.poll_write(cx, buf)
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
+        if self.is_closed() {
+            return Poll::Ready(Err(Error::new(
+                ErrorKind::Other,
+                "Attempted to flush after stream was closed",
+            )));
+        }
+        assert!(!self.is_closed());
+        self.project().inner.poll_flush(cx)
+    }
+
+    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
+        if self.is_closed() {
+            return Poll::Ready(Err(Error::new(
+                ErrorKind::Other,
+                "Attempted to close after stream was closed",
+            )));
+        }
+        let this = self.project();
+        match this.inner.poll_close(cx) {
+            Poll::Ready(Ok(())) => {
+                *this.closed = true;
+                Poll::Ready(Ok(()))
+            }
+            other => other,
+        }
+    }
+
+    fn poll_write_vectored(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        bufs: &[IoSlice<'_>],
+    ) -> Poll<Result<usize>> {
+        if self.is_closed() {
+            return Poll::Ready(Err(Error::new(
+                ErrorKind::Other,
+                "Attempted to write after stream was closed",
+            )));
+        }
+        self.project().inner.poll_write_vectored(cx, bufs)
+    }
+}