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::{self, Cursor};
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             assert_eq!(self.cnt, 0);
27 
28             let n = cmp::min(4, buf.len());
29             let buf = &buf[0..n];
30 
31             self.cnt += 1;
32             self.buf.extend(buf);
33             Ok(buf.len()).into()
34         }
35 
36         fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
37             Ok(()).into()
38         }
39 
40         fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
41             Ok(()).into()
42         }
43     }
44 
45     let mut wr = Wr {
46         buf: BytesMut::with_capacity(64),
47         cnt: 0,
48     };
49 
50     let mut buf = Cursor::new(&b"hello world"[..]);
51 
52     assert_ok!(wr.write_buf(&mut buf).await);
53     assert_eq!(wr.buf, b"hell"[..]);
54     assert_eq!(wr.cnt, 1);
55     assert_eq!(buf.position(), 4);
56 }
57