-
Notifications
You must be signed in to change notification settings - Fork 626
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add test utility that verifies an AsyncWrite is closed correctly (#2159)
* Add test utility that verifies an AsyncWrite is closed correctly * Add track_closed for sinks too
- Loading branch information
Showing
5 changed files
with
247 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,56 @@ | ||
| //! Additional combinators for testing sinks. | ||
|
|
||
| use futures_sink::Sink; | ||
|
|
||
| pub use crate::track_closed::TrackClosed; | ||
|
|
||
| /// Additional combinators for testing sinks. | ||
| pub trait SinkTestExt<Item>: Sink<Item> { | ||
| /// Track whether this sink has been closed and panics if it is used after closing. | ||
| /// | ||
| /// # Examples | ||
| /// | ||
| /// ``` | ||
| /// # futures::executor::block_on(async { | ||
| /// use futures::sink::{SinkExt, drain}; | ||
| /// use futures_test::sink::SinkTestExt; | ||
| /// | ||
| /// let mut sink = drain::<i32>().track_closed(); | ||
| /// | ||
| /// sink.send(1).await?; | ||
| /// assert!(!sink.is_closed()); | ||
| /// sink.close().await?; | ||
| /// assert!(sink.is_closed()); | ||
| /// | ||
| /// # Ok::<(), std::convert::Infallible>(()) })?; | ||
| /// # Ok::<(), std::convert::Infallible>(()) | ||
| /// ``` | ||
| /// | ||
| /// Note: Unlike [`AsyncWriteTestExt::track_closed`] when | ||
| /// used as a sink the adaptor will panic if closed too early as there's no easy way to | ||
| /// integrate as an error. | ||
| /// | ||
| /// [`AsyncWriteTestExt::track_closed`]: crate::io::AsyncWriteTestExt::track_closed | ||
| /// | ||
| /// ``` | ||
| /// # futures::executor::block_on(async { | ||
| /// use std::panic::AssertUnwindSafe; | ||
| /// use futures::{sink::{SinkExt, drain}, future::FutureExt}; | ||
| /// use futures_test::sink::SinkTestExt; | ||
| /// | ||
| /// let mut sink = drain::<i32>().track_closed(); | ||
| /// | ||
| /// sink.close().await?; | ||
| /// assert!(AssertUnwindSafe(sink.send(1)).catch_unwind().await.is_err()); | ||
| /// # Ok::<(), std::convert::Infallible>(()) })?; | ||
| /// # Ok::<(), std::convert::Infallible>(()) | ||
| /// ``` | ||
| fn track_closed(self) -> TrackClosed<Self> | ||
| where | ||
| Self: Sized, | ||
| { | ||
| TrackClosed::new(self) | ||
| } | ||
| } | ||
|
|
||
| impl<Item, W> SinkTestExt<Item> for W where W: Sink<Item> {} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,146 @@ | ||
| use futures_io::AsyncWrite; | ||
| use futures_sink::Sink; | ||
| use std::{ | ||
| io::{self, IoSlice}, | ||
| pin::Pin, | ||
| task::{Context, Poll}, | ||
| }; | ||
|
|
||
| /// Async wrapper that tracks whether it has been closed. | ||
| /// | ||
| /// See the `track_closed` methods on: | ||
| /// * [`SinkTestExt`](crate::sink::SinkTestExt::track_closed) | ||
| /// * [`AsyncWriteTestExt`](crate::io::AsyncWriteTestExt::track_closed) | ||
| #[pin_project::pin_project] | ||
| #[derive(Debug)] | ||
| pub struct TrackClosed<T> { | ||
| #[pin] | ||
| inner: T, | ||
| closed: bool, | ||
| } | ||
|
|
||
| impl<T> TrackClosed<T> { | ||
| pub(crate) fn new(inner: T) -> TrackClosed<T> { | ||
| TrackClosed { | ||
| inner, | ||
| closed: false, | ||
| } | ||
| } | ||
|
|
||
| /// Check whether this object has been closed. | ||
| pub fn is_closed(&self) -> bool { | ||
| self.closed | ||
| } | ||
|
|
||
| /// Acquires a reference to the underlying object that this adaptor is | ||
| /// wrapping. | ||
| pub fn get_ref(&self) -> &T { | ||
| &self.inner | ||
| } | ||
|
|
||
| /// Acquires a mutable reference to the underlying object that this | ||
| /// adaptor is wrapping. | ||
| pub fn get_mut(&mut self) -> &mut T { | ||
| &mut self.inner | ||
| } | ||
|
|
||
| /// Acquires a pinned mutable reference to the underlying object that | ||
| /// this adaptor is wrapping. | ||
| pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut T> { | ||
| self.project().inner | ||
| } | ||
|
|
||
| /// Consumes this adaptor returning the underlying object. | ||
| pub fn into_inner(self) -> T { | ||
| self.inner | ||
| } | ||
| } | ||
|
|
||
| impl<T: AsyncWrite> AsyncWrite for TrackClosed<T> { | ||
| fn poll_write( | ||
| self: Pin<&mut Self>, | ||
| cx: &mut Context<'_>, | ||
| buf: &[u8], | ||
| ) -> Poll<io::Result<usize>> { | ||
| if self.is_closed() { | ||
| return Poll::Ready(Err(io::Error::new( | ||
| io::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<io::Result<()>> { | ||
| if self.is_closed() { | ||
| return Poll::Ready(Err(io::Error::new( | ||
| io::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<io::Result<()>> { | ||
| if self.is_closed() { | ||
| return Poll::Ready(Err(io::Error::new( | ||
| io::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<io::Result<usize>> { | ||
| if self.is_closed() { | ||
| return Poll::Ready(Err(io::Error::new( | ||
| io::ErrorKind::Other, | ||
| "Attempted to write after stream was closed", | ||
| ))); | ||
| } | ||
| self.project().inner.poll_write_vectored(cx, bufs) | ||
| } | ||
| } | ||
|
|
||
| impl<Item, T: Sink<Item>> Sink<Item> for TrackClosed<T> { | ||
| type Error = T::Error; | ||
|
|
||
| fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { | ||
| assert!(!self.is_closed()); | ||
| self.project().inner.poll_ready(cx) | ||
| } | ||
|
|
||
| fn start_send(self: Pin<&mut Self>, item: Item) -> Result<(), Self::Error> { | ||
| assert!(!self.is_closed()); | ||
| self.project().inner.start_send(item) | ||
| } | ||
|
|
||
| fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { | ||
| assert!(!self.is_closed()); | ||
| self.project().inner.poll_flush(cx) | ||
| } | ||
|
|
||
| fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { | ||
| assert!(!self.is_closed()); | ||
| let this = self.project(); | ||
| match this.inner.poll_close(cx) { | ||
| Poll::Ready(Ok(())) => { | ||
| *this.closed = true; | ||
| Poll::Ready(Ok(())) | ||
| } | ||
| other => other, | ||
| } | ||
| } | ||
| } |