1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

#![allow(missing_docs)]

use async_std::io::{self, Cursor, Read, Write};
use async_std::pin::Pin;
use async_std::task::{Context, Poll};
use pin_project::pin_project;

pub type MockCursor = Cursor<Vec<u8>>;

#[pin_project]
#[derive(Clone, Debug)]
pub struct MockStream {
    #[pin]
    reader: MockCursor,
    #[pin]
    writer: MockCursor,
}

impl Default for MockStream {
    fn default() -> Self {
        Self::new()
    }
}

impl MockStream {
    pub fn new() -> MockStream {
        MockStream {
            reader: MockCursor::new(Vec::new()),
            writer: MockCursor::new(Vec::new()),
        }
    }

    pub fn with_vec(vec: Vec<u8>) -> MockStream {
        MockStream {
            reader: MockCursor::new(vec),
            writer: MockCursor::new(Vec::new()),
        }
    }

    pub fn take_vec(&mut self) -> Vec<u8> {
        let vec = self.writer.get_ref().to_vec();
        self.writer.set_position(0);
        self.writer.get_mut().clear();
        vec
    }

    pub fn next_vec(&mut self, vec: &[u8]) {
        let cursor = &mut self.reader;
        cursor.set_position(0);
        cursor.get_mut().clear();
        cursor.get_mut().extend_from_slice(vec);
    }

    pub fn swap(&mut self) {
        let cur_write = &mut self.writer;
        let cur_read = &mut self.reader;
        let vec_write = cur_write.get_ref().to_vec();
        let vec_read = cur_read.get_ref().to_vec();
        cur_write.set_position(0);
        cur_read.set_position(0);
        cur_write.get_mut().clear();
        cur_read.get_mut().clear();
        // swap cursors
        cur_read.get_mut().extend_from_slice(vec_write.as_slice());
        cur_write.get_mut().extend_from_slice(vec_read.as_slice());
    }
}

impl Write for MockStream {
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<io::Result<usize>> {
        let this = self.project();
        let _: Pin<&mut _> = this.writer;
        this.writer.poll_write(cx, buf)
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
        let this = self.project();
        let _: Pin<&mut _> = this.writer;
        this.writer.poll_flush(cx)
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
        let this = self.project();
        let _: Pin<&mut _> = this.writer;
        this.writer.poll_close(cx)
    }
}

impl Read for MockStream {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context,
        buf: &mut [u8],
    ) -> Poll<io::Result<usize>> {
        let this = self.project();
        let _: Pin<&mut _> = this.reader;
        this.reader.poll_read(cx, buf)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use async_std::prelude::*;

    #[async_attributes::test]
    async fn write_take_test() {
        let mut mock = MockStream::new();
        // write to mock stream
        mock.write_all(&[1, 2, 3]).await.unwrap();
        assert_eq!(mock.take_vec(), vec![1, 2, 3]);
    }

    #[async_attributes::test]
    async fn read_with_vec_test() {
        let mut mock = MockStream::with_vec(vec![4, 5]);
        let mut vec = Vec::new();
        mock.read_to_end(&mut vec).await.unwrap();
        assert_eq!(vec, vec![4, 5]);
    }

    #[async_attributes::test]
    async fn swap_test() {
        let mut mock = MockStream::new();
        let mut vec = Vec::new();
        mock.write_all(&[8, 9, 10]).await.unwrap();
        mock.swap();
        mock.read_to_end(&mut vec).await.unwrap();
        assert_eq!(vec, vec![8, 9, 10]);
    }
}