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::io; |
9 | use std::pin::Pin; |
10 | use std::task::{Context, Poll}; |
11 | |
12 | #[tokio::test ] |
13 | async fn write() { |
14 | struct Wr { |
15 | buf: BytesMut, |
16 | cnt: usize, |
17 | } |
18 | |
19 | impl AsyncWrite for Wr { |
20 | fn poll_write( |
21 | mut self: Pin<&mut Self>, |
22 | _cx: &mut Context<'_>, |
23 | buf: &[u8], |
24 | ) -> Poll<io::Result<usize>> { |
25 | assert_eq!(self.cnt, 0); |
26 | self.buf.extend(&buf[0..4]); |
27 | Ok(4).into() |
28 | } |
29 | |
30 | fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> { |
31 | Ok(()).into() |
32 | } |
33 | |
34 | fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> { |
35 | Ok(()).into() |
36 | } |
37 | } |
38 | |
39 | let mut wr = Wr { |
40 | buf: BytesMut::with_capacity(64), |
41 | cnt: 0, |
42 | }; |
43 | |
44 | let n = assert_ok!(wr.write(b"hello world" ).await); |
45 | assert_eq!(n, 4); |
46 | assert_eq!(wr.buf, b"hell" [..]); |
47 | } |
48 | |
49 | #[tokio::test ] |
50 | async fn write_cursor() { |
51 | use std::io::Cursor; |
52 | |
53 | let mut wr = Cursor::new(Vec::new()); |
54 | |
55 | let n = assert_ok!(wr.write(b"hello world" ).await); |
56 | assert_eq!(n, 11); |
57 | assert_eq!(wr.get_ref().as_slice(), &b"hello world" [..]); |
58 | } |
59 | |