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 ] |
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 | |