1 #![warn(rust_2018_idioms)]
2 #![cfg(feature = "full")]
3 
4 use tokio::io::{AsyncWrite, AsyncWriteExt};
5 use tokio_test::assert_ok;
6 
7 use bytes::BytesMut;
8 use std::cmp;
9 use std::io;
10 use std::pin::Pin;
11 use std::task::{Context, Poll};
12 
13 #[tokio::test]
write_all()14 async fn write_all() {
15     struct Wr {
16         buf: BytesMut,
17         cnt: usize,
18     }
19 
20     impl AsyncWrite for Wr {
21         fn poll_write(
22             mut self: Pin<&mut Self>,
23             _cx: &mut Context<'_>,
24             buf: &[u8],
25         ) -> Poll<io::Result<usize>> {
26             let n = cmp::min(4, buf.len());
27             let buf = &buf[0..n];
28 
29             self.cnt += 1;
30             self.buf.extend(buf);
31             Ok(buf.len()).into()
32         }
33 
34         fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
35             Ok(()).into()
36         }
37 
38         fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
39             Ok(()).into()
40         }
41     }
42 
43     let mut wr = Wr {
44         buf: BytesMut::with_capacity(64),
45         cnt: 0,
46     };
47 
48     assert_ok!(wr.write_all(b"hello world").await);
49     assert_eq!(wr.buf, b"hello world"[..]);
50     assert_eq!(wr.cnt, 3);
51 }
52