1 use futures::executor::block_on;
2 use futures::future::{Future, FutureExt};
3 use futures::io::{
4     AsyncSeek, AsyncSeekExt, AsyncWrite, AsyncWriteExt, BufWriter, Cursor, SeekFrom,
5 };
6 use futures::task::{Context, Poll};
7 use futures_test::task::noop_context;
8 use std::io;
9 use std::pin::Pin;
10 
11 struct MaybePending {
12     inner: Vec<u8>,
13     ready: bool,
14 }
15 
16 impl MaybePending {
new(inner: Vec<u8>) -> Self17     fn new(inner: Vec<u8>) -> Self {
18         Self { inner, ready: false }
19     }
20 }
21 
22 impl AsyncWrite for MaybePending {
poll_write( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll<io::Result<usize>>23     fn poll_write(
24         mut self: Pin<&mut Self>,
25         cx: &mut Context<'_>,
26         buf: &[u8],
27     ) -> Poll<io::Result<usize>> {
28         if self.ready {
29             self.ready = false;
30             Pin::new(&mut self.inner).poll_write(cx, buf)
31         } else {
32             self.ready = true;
33             Poll::Pending
34         }
35     }
36 
poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>>37     fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
38         Pin::new(&mut self.inner).poll_flush(cx)
39     }
40 
poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>>41     fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
42         Pin::new(&mut self.inner).poll_close(cx)
43     }
44 }
45 
run<F: Future + Unpin>(mut f: F) -> F::Output46 fn run<F: Future + Unpin>(mut f: F) -> F::Output {
47     let mut cx = noop_context();
48     loop {
49         if let Poll::Ready(x) = f.poll_unpin(&mut cx) {
50             return x;
51         }
52     }
53 }
54 
55 #[test]
buf_writer()56 fn buf_writer() {
57     let mut writer = BufWriter::with_capacity(2, Vec::new());
58 
59     block_on(writer.write(&[0, 1])).unwrap();
60     assert_eq!(writer.buffer(), []);
61     assert_eq!(*writer.get_ref(), [0, 1]);
62 
63     block_on(writer.write(&[2])).unwrap();
64     assert_eq!(writer.buffer(), [2]);
65     assert_eq!(*writer.get_ref(), [0, 1]);
66 
67     block_on(writer.write(&[3])).unwrap();
68     assert_eq!(writer.buffer(), [2, 3]);
69     assert_eq!(*writer.get_ref(), [0, 1]);
70 
71     block_on(writer.flush()).unwrap();
72     assert_eq!(writer.buffer(), []);
73     assert_eq!(*writer.get_ref(), [0, 1, 2, 3]);
74 
75     block_on(writer.write(&[4])).unwrap();
76     block_on(writer.write(&[5])).unwrap();
77     assert_eq!(writer.buffer(), [4, 5]);
78     assert_eq!(*writer.get_ref(), [0, 1, 2, 3]);
79 
80     block_on(writer.write(&[6])).unwrap();
81     assert_eq!(writer.buffer(), [6]);
82     assert_eq!(*writer.get_ref(), [0, 1, 2, 3, 4, 5]);
83 
84     block_on(writer.write(&[7, 8])).unwrap();
85     assert_eq!(writer.buffer(), []);
86     assert_eq!(*writer.get_ref(), [0, 1, 2, 3, 4, 5, 6, 7, 8]);
87 
88     block_on(writer.write(&[9, 10, 11])).unwrap();
89     assert_eq!(writer.buffer(), []);
90     assert_eq!(*writer.get_ref(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]);
91 
92     block_on(writer.flush()).unwrap();
93     assert_eq!(writer.buffer(), []);
94     assert_eq!(*writer.get_ref(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]);
95 }
96 
97 #[test]
buf_writer_inner_flushes()98 fn buf_writer_inner_flushes() {
99     let mut w = BufWriter::with_capacity(3, Vec::new());
100     block_on(w.write(&[0, 1])).unwrap();
101     assert_eq!(*w.get_ref(), []);
102     block_on(w.flush()).unwrap();
103     let w = w.into_inner();
104     assert_eq!(w, [0, 1]);
105 }
106 
107 #[test]
buf_writer_seek()108 fn buf_writer_seek() {
109     // FIXME: when https://github.com/rust-lang/futures-rs/issues/1510 fixed,
110     // use `Vec::new` instead of `vec![0; 8]`.
111     let mut w = BufWriter::with_capacity(3, Cursor::new(vec![0; 8]));
112     block_on(w.write_all(&[0, 1, 2, 3, 4, 5])).unwrap();
113     block_on(w.write_all(&[6, 7])).unwrap();
114     assert_eq!(block_on(w.seek(SeekFrom::Current(0))).ok(), Some(8));
115     assert_eq!(&w.get_ref().get_ref()[..], &[0, 1, 2, 3, 4, 5, 6, 7][..]);
116     assert_eq!(block_on(w.seek(SeekFrom::Start(2))).ok(), Some(2));
117     block_on(w.write_all(&[8, 9])).unwrap();
118     block_on(w.flush()).unwrap();
119     assert_eq!(&w.into_inner().into_inner()[..], &[0, 1, 8, 9, 4, 5, 6, 7]);
120 }
121 
122 #[test]
maybe_pending_buf_writer()123 fn maybe_pending_buf_writer() {
124     let mut writer = BufWriter::with_capacity(2, MaybePending::new(Vec::new()));
125 
126     run(writer.write(&[0, 1])).unwrap();
127     assert_eq!(writer.buffer(), []);
128     assert_eq!(&writer.get_ref().inner, &[0, 1]);
129 
130     run(writer.write(&[2])).unwrap();
131     assert_eq!(writer.buffer(), [2]);
132     assert_eq!(&writer.get_ref().inner, &[0, 1]);
133 
134     run(writer.write(&[3])).unwrap();
135     assert_eq!(writer.buffer(), [2, 3]);
136     assert_eq!(&writer.get_ref().inner, &[0, 1]);
137 
138     run(writer.flush()).unwrap();
139     assert_eq!(writer.buffer(), []);
140     assert_eq!(&writer.get_ref().inner, &[0, 1, 2, 3]);
141 
142     run(writer.write(&[4])).unwrap();
143     run(writer.write(&[5])).unwrap();
144     assert_eq!(writer.buffer(), [4, 5]);
145     assert_eq!(&writer.get_ref().inner, &[0, 1, 2, 3]);
146 
147     run(writer.write(&[6])).unwrap();
148     assert_eq!(writer.buffer(), [6]);
149     assert_eq!(writer.get_ref().inner, &[0, 1, 2, 3, 4, 5]);
150 
151     run(writer.write(&[7, 8])).unwrap();
152     assert_eq!(writer.buffer(), []);
153     assert_eq!(writer.get_ref().inner, &[0, 1, 2, 3, 4, 5, 6, 7, 8]);
154 
155     run(writer.write(&[9, 10, 11])).unwrap();
156     assert_eq!(writer.buffer(), []);
157     assert_eq!(writer.get_ref().inner, &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]);
158 
159     run(writer.flush()).unwrap();
160     assert_eq!(writer.buffer(), []);
161     assert_eq!(&writer.get_ref().inner, &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]);
162 }
163 
164 #[test]
maybe_pending_buf_writer_inner_flushes()165 fn maybe_pending_buf_writer_inner_flushes() {
166     let mut w = BufWriter::with_capacity(3, MaybePending::new(Vec::new()));
167     run(w.write(&[0, 1])).unwrap();
168     assert_eq!(&w.get_ref().inner, &[]);
169     run(w.flush()).unwrap();
170     let w = w.into_inner().inner;
171     assert_eq!(w, [0, 1]);
172 }
173 
174 #[test]
maybe_pending_buf_writer_seek()175 fn maybe_pending_buf_writer_seek() {
176     struct MaybePendingSeek {
177         inner: Cursor<Vec<u8>>,
178         ready_write: bool,
179         ready_seek: bool,
180     }
181 
182     impl MaybePendingSeek {
183         fn new(inner: Vec<u8>) -> Self {
184             Self { inner: Cursor::new(inner), ready_write: false, ready_seek: false }
185         }
186     }
187 
188     impl AsyncWrite for MaybePendingSeek {
189         fn poll_write(
190             mut self: Pin<&mut Self>,
191             cx: &mut Context<'_>,
192             buf: &[u8],
193         ) -> Poll<io::Result<usize>> {
194             if self.ready_write {
195                 self.ready_write = false;
196                 Pin::new(&mut self.inner).poll_write(cx, buf)
197             } else {
198                 self.ready_write = true;
199                 Poll::Pending
200             }
201         }
202 
203         fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
204             Pin::new(&mut self.inner).poll_flush(cx)
205         }
206 
207         fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
208             Pin::new(&mut self.inner).poll_close(cx)
209         }
210     }
211 
212     impl AsyncSeek for MaybePendingSeek {
213         fn poll_seek(
214             mut self: Pin<&mut Self>,
215             cx: &mut Context<'_>,
216             pos: SeekFrom,
217         ) -> Poll<io::Result<u64>> {
218             if self.ready_seek {
219                 self.ready_seek = false;
220                 Pin::new(&mut self.inner).poll_seek(cx, pos)
221             } else {
222                 self.ready_seek = true;
223                 Poll::Pending
224             }
225         }
226     }
227 
228     // FIXME: when https://github.com/rust-lang/futures-rs/issues/1510 fixed,
229     // use `Vec::new` instead of `vec![0; 8]`.
230     let mut w = BufWriter::with_capacity(3, MaybePendingSeek::new(vec![0; 8]));
231     run(w.write_all(&[0, 1, 2, 3, 4, 5])).unwrap();
232     run(w.write_all(&[6, 7])).unwrap();
233     assert_eq!(run(w.seek(SeekFrom::Current(0))).ok(), Some(8));
234     assert_eq!(&w.get_ref().inner.get_ref()[..], &[0, 1, 2, 3, 4, 5, 6, 7][..]);
235     assert_eq!(run(w.seek(SeekFrom::Start(2))).ok(), Some(2));
236     run(w.write_all(&[8, 9])).unwrap();
237     run(w.flush()).unwrap();
238     assert_eq!(&w.into_inner().inner.into_inner()[..], &[0, 1, 8, 9, 4, 5, 6, 7]);
239 }
240