1#![warn(rust_2018_idioms)]
2#![cfg(feature = "full")]
3
4use tokio::io::{AsyncWrite, AsyncWriteExt};
5use tokio_test::assert_ok;
6
7use bytes::BytesMut;
8use std::cmp;
9use std::io;
10use std::pin::Pin;
11use std::task::{Context, Poll};
12
13#[tokio::test]
14async 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