test: fetch actions from mock handle before write (#5814)

tokio-test/src/io.rs

@@ -409,6 +409,20 @@ impl AsyncWrite for Mock {
             // If a sleep is set, it has already fired
             self.inner.sleep = None;
 
+            if self.inner.actions.is_empty() {
+                match self.inner.poll_action(cx) {
+                    Poll::Pending => {
+                        // do not propagate pending
+                    }
+                    Poll::Ready(Some(action)) => {
+                        self.inner.actions.push_back(action);
+                    }
+                    Poll::Ready(None) => {
+                        panic!("unexpected write");
+                    }
+                }
+            }
+
             match self.inner.write(buf) {
                 Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                     if let Some(rem) = self.inner.remaining_wait() {

tokio-test/tests/io.rs

@@ -51,6 +51,29 @@ async fn write() {
     mock.write_all(b"world!").await.expect("write 2");
 }
 
+#[tokio::test]
+async fn write_with_handle() {
+    let (mut mock, mut handle) = Builder::new().build_with_handle();
+    handle.write(b"hello ");
+    handle.write(b"world!");
+
+    mock.write_all(b"hello ").await.expect("write 1");
+    mock.write_all(b"world!").await.expect("write 2");
+}
+
+#[tokio::test]
+async fn read_with_handle() {
+    let (mut mock, mut handle) = Builder::new().build_with_handle();
+    handle.read(b"hello ");
+    handle.read(b"world!");
+
+    let mut buf = vec![0; 6];
+    mock.read_exact(&mut buf).await.expect("read 1");
+    assert_eq!(&buf[..], b"hello ");
+    mock.read_exact(&mut buf).await.expect("read 2");
+    assert_eq!(&buf[..], b"world!");
+}
+
 #[tokio::test]
 async fn write_error() {
     let error = io::Error::new(io::ErrorKind::Other, "cruel");