diff --git a/tokio-test/src/io.rs b/tokio-test/src/io.rs index 2e16f1fe0..1fc2f41aa 100644 --- a/tokio-test/src/io.rs +++ b/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() { diff --git a/tokio-test/tests/io.rs b/tokio-test/tests/io.rs index f182e37ee..5f2ed427c 100644 --- a/tokio-test/tests/io.rs +++ b/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");