1use std::{
2 io::{self, Cursor, Read, Write},
3 pin::Pin,
4 task::{Context, Poll},
5};
6use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
7
8#[derive(Clone, Debug)]
10pub struct MockStream {
11 written: Cursor<Vec<u8>>,
12 received: Cursor<Vec<u8>>,
13}
14
15impl MockStream {
16 pub fn empty() -> MockStream {
18 MockStream::new(&[])
19 }
20
21 pub fn new(initial: &[u8]) -> MockStream {
23 MockStream {
24 written: Cursor::new(vec![]),
25 received: Cursor::new(initial.to_owned()),
26 }
27 }
28
29 pub fn written(&self) -> &[u8] {
31 self.written.get_ref()
32 }
33
34 pub fn received(&self) -> &[u8] {
36 self.received.get_ref()
37 }
38}
39
40impl AsyncRead for MockStream {
41 fn poll_read(
42 mut self: Pin<&mut Self>,
43 _: &mut Context<'_>,
44 buf: &mut ReadBuf<'_>,
45 ) -> Poll<io::Result<()>> {
46 let n = self.as_mut().received.read(buf.initialize_unfilled())?;
47 buf.advance(n);
48 Poll::Ready(Ok(()))
49 }
50}
51
52impl AsyncWrite for MockStream {
53 fn poll_write(
54 mut self: Pin<&mut Self>,
55 _: &mut Context<'_>,
56 buf: &[u8],
57 ) -> Poll<Result<usize, io::Error>> {
58 Poll::Ready(self.as_mut().written.write(buf))
59 }
60
61 fn poll_flush(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
62 Poll::Ready(self.as_mut().written.flush())
63 }
64
65 fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
66 Poll::Ready(Ok(()))
67 }
68}