diff options
author | Carl Lerche <me@carllerche.com> | 2019-07-09 12:37:14 -0700 |
---|---|---|
committer | GitHub <noreply@github.com> | 2019-07-09 12:37:14 -0700 |
commit | 64343f1b786b386bb4fbd6169fc4850cce9e245a (patch) | |
tree | cc541e2ad896afdd44cd7c3d43000deab4b41009 /tokio/tests | |
parent | 82795184c1a8ce136c4f0fee2a6c9127ff495565 (diff) |
tokio: add AsyncWriteExt::write_all (#1277)
Diffstat (limited to 'tokio/tests')
-rw-r--r-- | tokio/tests/io_write.rs | 18 | ||||
-rw-r--r-- | tokio/tests/io_write_all.rs | 51 |
2 files changed, 64 insertions, 5 deletions
diff --git a/tokio/tests/io_write.rs b/tokio/tests/io_write.rs index 990abdf2..482f21ba 100644 --- a/tokio/tests/io_write.rs +++ b/tokio/tests/io_write.rs @@ -11,7 +11,10 @@ use std::task::{Context, Poll}; #[tokio::test] async fn write() { - struct Wr(BytesMut); + struct Wr { + buf: BytesMut, + cnt: usize, + } impl AsyncWrite for Wr { fn poll_write( @@ -19,8 +22,9 @@ async fn write() { _cx: &mut Context<'_>, buf: &[u8], ) -> Poll<io::Result<usize>> { - self.0.extend(buf); - Ok(buf.len()).into() + assert_eq!(self.cnt, 0); + self.buf.extend(&buf[0..4]); + Ok(4).into() } fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> { @@ -32,8 +36,12 @@ async fn write() { } } - let mut wr = Wr(BytesMut::with_capacity(64)); + let mut wr = Wr { + buf: BytesMut::with_capacity(64), + cnt: 0, + }; let n = assert_ok!(wr.write(b"hello world").await); - assert_eq!(n, 11); + assert_eq!(n, 4); + assert_eq!(wr.buf, b"hell"[..]); } diff --git a/tokio/tests/io_write_all.rs b/tokio/tests/io_write_all.rs new file mode 100644 index 00000000..b12e0f81 --- /dev/null +++ b/tokio/tests/io_write_all.rs @@ -0,0 +1,51 @@ +#![deny(warnings, rust_2018_idioms)] +#![feature(async_await)] + +use tokio::io::{AsyncWrite, AsyncWriteExt}; +use tokio_test::assert_ok; + +use bytes::BytesMut; +use std::cmp; +use std::io; +use std::pin::Pin; +use std::task::{Context, Poll}; + +#[tokio::test] +async fn write_all() { + struct Wr { + buf: BytesMut, + cnt: usize, + } + + impl AsyncWrite for Wr { + fn poll_write( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + let n = cmp::min(4, buf.len()); + let buf = &buf[0..n]; + + self.cnt += 1; + self.buf.extend(buf); + Ok(buf.len()).into() + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> { + Ok(()).into() + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> { + Ok(()).into() + } + } + + let mut wr = Wr { + buf: BytesMut::with_capacity(64), + cnt: 0, + }; + + assert_ok!(wr.write_all(b"hello world").await); + assert_eq!(wr.buf, b"hello world"[..]); + assert_eq!(wr.cnt, 3); +} |