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