io: add shutdown() and shutdown_error() on tokio_test::io::Builder

Signed-off-by: ADD-SP <[email protected]>

Author: ADD-SP

tokio-test/src/io.rs

@@ -60,11 +60,13 @@ pub struct Builder {
 enum Action {
     Read(Vec<u8>),
     Write(Vec<u8>),
+    Shutdown,
     Wait(Duration),
     // Wrapped in Arc so that Builder can be cloned and Send.
     // Mock is not cloned as does not need to check Rc for ref counts.
     ReadError(Option<Arc<io::Error>>),
     WriteError(Option<Arc<io::Error>>),
+    ShutdownError(Option<Arc<io::Error>>),
 }
 
 struct Inner {
@@ -76,6 +78,13 @@ struct Inner {
     name: String,
 }
 
+enum Shutdown {
+    ShouldSuccess,
+    ShouldError(io::Error),
+    NeedWait,
+    NoActions,
+}
+
 impl Builder {
     /// Return a new, empty `Builder`.
     pub fn new() -> Self {
@@ -130,6 +139,25 @@ impl Builder {
         self
     }
 
+    /// Sequence a shutdown operation.
+    ///
+    /// The next operation in the mock's script will be to expect a
+    /// [`AsyncWrite::poll_shutdown`] call.
+    pub fn shutdown(&mut self) -> &mut Self {
+        self.actions.push_back(Action::Shutdown);
+        self
+    }
+
+    /// Sequence a shutdown operation that produces an error.
+    ///
+    /// The next operation in the mock's script will be to expect a
+    /// [`AsyncWrite::poll_shutdown`] call that returns `error`.
+    pub fn shutdown_error(&mut self, error: io::Error) -> &mut Self {
+        let error = Some(error.into());
+        self.actions.push_back(Action::ShutdownError(error));
+        self
+    }
+
     /// Set name of the mock IO object to include in panic messages and debug output
     pub fn name(&mut self, name: impl Into<String>) -> &mut Self {
         self.name = name.into();
@@ -198,6 +226,10 @@ impl Inner {
 
         let rx = UnboundedReceiverStream::new(rx);
 
+        // Actually, we should deny any write action after the shutdown action.
+        // However, since we currently doesn't check the write action after the error
+        // like BrokenPipe error, we ignore this case to keep the behavior consistent.
+
         let inner = Inner {
             actions,
             sleep: None,
@@ -242,6 +274,9 @@ impl Inner {
                 // Either waiting or expecting a write
                 Err(io::ErrorKind::WouldBlock.into())
             }
+            Some(&mut Action::Shutdown) | Some(&mut Action::ShutdownError(_)) => {
+                panic!("unexpected read, expect poll_shutdown");
+            }
             None => Ok(()),
         }
     }
@@ -284,6 +319,9 @@ impl Inner {
                     break;
                 }
                 Action::Read(_) | Action::ReadError(_) => (),
+                Action::Shutdown | Action::ShutdownError(_) => {
+                    panic!("unexpected write, expect poll_shutdown");
+                }
             }
         }
 
@@ -297,6 +335,25 @@ impl Inner {
         }
     }
 
+    fn shutdown(&mut self) -> Shutdown {
+        match self.action() {
+            Some(&mut Action::Shutdown) => Shutdown::ShouldSuccess,
+            Some(&mut Action::ShutdownError(ref mut err)) => {
+                let err = err.take().expect("Should have been removed from actions.");
+                let err = Arc::try_unwrap(err).expect("There are no other references.");
+                Shutdown::ShouldError(err)
+            }
+            Some(&mut Action::Read(_)) | Some(&mut Action::ReadError(_)) => {
+                panic!("unexpected poll_shutdown, expect read");
+            }
+            Some(&mut Action::Write(_)) | Some(&mut Action::WriteError(_)) => {
+                panic!("unexpected poll_shutdown, expect write");
+            }
+            Some(&mut Action::Wait(_)) => Shutdown::NeedWait,
+            None => Shutdown::NoActions,
+        }
+    }
+
     fn action(&mut self) -> Option<&mut Action> {
         loop {
             if self.actions.is_empty() {
@@ -333,6 +390,12 @@ impl Inner {
                         break;
                     }
                 }
+                Action::Shutdown => break,
+                Action::ShutdownError(ref mut error) => {
+                    if error.is_some() {
+                        break;
+                    }
+                }
             }
 
             let _action = self.actions.pop_front();
@@ -472,8 +535,55 @@ impl AsyncWrite for Mock {
         Poll::Ready(Ok(()))
     }
 
-    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
-        Poll::Ready(Ok(()))
+    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
+        loop {
+            if let Some(ref mut sleep) = self.inner.sleep {
+                ready!(Pin::new(sleep).poll(cx));
+            }
+
+            // If a sleep is set, it has already fired
+            self.inner.sleep = None;
+
+            match self.inner.shutdown() {
+                Shutdown::ShouldSuccess => {
+                    assert!(matches!(
+                        self.inner.actions.pop_front(),
+                        Some(Action::Shutdown)
+                    ));
+                    self.maybe_wakeup_reader();
+                    return Poll::Ready(Ok(()));
+                }
+                Shutdown::ShouldError(e) => {
+                    assert!(matches!(
+                        self.inner.actions.pop_front(),
+                        Some(Action::ShutdownError(_))
+                    ));
+                    self.maybe_wakeup_reader();
+                    return Poll::Ready(Err(e));
+                }
+                Shutdown::NeedWait => {
+                    if let Some(rem) = self.inner.remaining_wait() {
+                        let until = Instant::now() + rem;
+                        self.inner.sleep = Some(Box::pin(time::sleep_until(until)));
+                    } else {
+                        panic!(
+                            "unexpected poll_shutdown, expect read or write {}",
+                            self.pmsg()
+                        );
+                    }
+                }
+                Shutdown::NoActions => {
+                    if let Some(action) = ready!(self.inner.poll_action(cx)) {
+                        self.inner.actions.push_back(action);
+                    } else {
+                        panic!(
+                            "unexpected poll_shutdown, but actually no more sequenced actions {}",
+                            self.pmsg()
+                        );
+                    }
+                }
+            }
+        }
     }
 }
 
@@ -496,7 +606,11 @@ impl Drop for Mock {
                 "There is still data left to write. {}",
                 self.pmsg()
             ),
+            Action::Shutdown => panic!("AsyncWrite::poll_shutdown was not called. {}", self.pmsg()),
             Action::ReadError(_) | Action::WriteError(_) | Action::Wait(_) => (),
+            // Since the existing implementation ignores the read/write error, so we also ignore the
+            // shutdown error here to keep the behavior consistent.
+            Action::ShutdownError(_) => (),
         });
     }
 }

tokio-test/tests/io.rs

@@ -8,8 +8,8 @@ use std::task::{Context, Poll};
 use tokio::io::AsyncWrite;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::time::{Duration, Instant};
-use tokio_test::assert_pending;
 use tokio_test::io::Builder;
+use tokio_test::{assert_pending, assert_ready};
 
 #[tokio::test]
 async fn read() {
@@ -232,7 +232,7 @@ fn do_not_panic_unconsumed_error() {
 // TODO: fix this bug in the next major release
 #[tokio::test]
 #[should_panic = "There are no other references.: Custom { kind: Other, error: \"cruel\" }"]
-async fn panic_if_clone_the_build_with_error_action() {
+async fn should_panic_if_clone_the_builder_with_error_action() {
     let mut builder = Builder::new();
     builder.write_error(io::Error::new(io::ErrorKind::Other, "cruel"));
     let mut builder2 = builder.clone();
@@ -257,3 +257,44 @@ async fn should_not_hang_forever_on_zero_length_write() {
         Poll::Pending => panic!("expected to write 0 bytes immediately, but pending instead"),
     }
 }
+
+#[tokio::test]
+async fn shutdown() {
+    let mut mock = Builder::new().shutdown().build();
+    mock.shutdown().await.unwrap();
+}
+
+#[tokio::test]
+async fn shutdown_error() {
+    let error = io::Error::new(io::ErrorKind::Other, "cruel");
+    let mut mock = Builder::new().shutdown_error(error).build();
+    let err = mock.shutdown().await.unwrap_err();
+    assert_eq!(err.kind(), io::ErrorKind::Other);
+    assert_eq!("cruel", format!("{err}"));
+}
+
+#[tokio::test(start_paused = true)]
+async fn shutdown_wait() {
+    const WAIT: Duration = Duration::from_secs(1);
+
+    let mock = Builder::new().wait(WAIT).shutdown().build();
+    pin_mut!(mock);
+
+    assert_pending!(mock.as_mut().poll_shutdown(&mut noop_context()));
+
+    tokio::time::advance(WAIT).await;
+    let _ = assert_ready!(mock.as_mut().poll_shutdown(&mut noop_context()));
+}
+
+#[test]
+#[should_panic = "AsyncWrite::poll_shutdown was not called. (1 actions remain)"]
+fn should_panic_on_leftover_shutdown_action() {
+    let _mock = Builder::new().shutdown().build();
+}
+
+#[test]
+fn should_not_panic_on_leftover_shutdown_error_action() {
+    let _mock = Builder::new()
+        .shutdown_error(io::Error::new(io::ErrorKind::Other, "cruel"))
+        .build();
+}