tk_bufstream/
mock.rs

1use std::cmp::min;
2use std::io::{self, Read, Write};
3use std::sync::{Arc, Mutex};
4
5use netbuf::RangeArgument;
6use futures::{Poll, Async};
7use tokio_io::{AsyncRead, AsyncWrite};
8
9/// A thing that implements tokio_core::io::Io but never ready
10///
11/// Useful for tests on codecs, where tests owns actual buffered stream.
12pub struct Mock;
13
14impl Read for Mock {
15    fn read(&mut self, _buf: &mut [u8]) -> io::Result<usize> {
16        Err(io::Error::new(io::ErrorKind::WouldBlock, "No read"))
17    }
18}
19
20impl Write for Mock {
21    fn write(&mut self, _buf: &[u8]) -> io::Result<usize> {
22        Err(io::Error::new(io::ErrorKind::WouldBlock, "No write"))
23    }
24    fn flush(&mut self) -> io::Result<()> {
25        Err(io::Error::new(io::ErrorKind::WouldBlock, "No flush"))
26    }
27}
28
29impl AsyncRead for Mock {}
30impl AsyncWrite for Mock {
31    fn shutdown(&mut self) -> Poll<(), io::Error> {
32        Ok(Async::NotReady)
33    }
34}
35
36
37/// A mock stream where you can push data to/from
38///
39/// Useful for more complex tests on streams than are possible with `Mock`
40#[derive(Clone)]
41pub struct MockData {
42    input: Arc<Mutex<Vec<u8>>>,
43    output: Arc<Mutex<Vec<u8>>>,
44}
45
46impl MockData {
47    /// New, empty Mockdata
48    pub fn new() -> MockData {
49        MockData {
50            input: Arc::new(Mutex::new(Vec::new())),
51            output: Arc::new(Mutex::new(Vec::new())),
52        }
53    }
54    /// Add some bytes to the next input. This data will be read by IoBuf
55    /// on thext `read()` call
56    pub fn add_input<D: AsRef<[u8]>>(&self, data: D) {
57        self.input.lock().unwrap().extend(data.as_ref())
58    }
59
60    /// Get slice of output
61    ///
62    /// Note: we only read bytes that has been `write`'en here, not
63    /// the bytes buffered inside `IoBuf`.
64    ///
65    /// Note 2: we use RangeArgument from `netbuf` here, as we already depend
66    /// on netbuf anyway. Eventually we will switch to `RangeArgument` from
67    /// `std` (both in `netbuf` and here) when latter is stable.
68    pub fn output<T: Into<RangeArgument>>(&self, range: T) -> Vec<u8> {
69        let buf = self.output.lock().unwrap();
70        use netbuf::RangeArgument::*;
71        match range.into() {
72            RangeTo(x) => buf[..x].to_vec(),
73            RangeFrom(x) => buf[x..].to_vec(),
74            Range(x, y) => buf[x..y].to_vec(),
75        }
76    }
77
78    /// Get first bytes of output and remove them from output buffer
79    ///
80    /// Note: we only read bytes that has been `write`'en here, not
81    /// the bytes buffered inside `IoBuf`.
82    pub fn get_output(&self, num: usize) -> Vec<u8> {
83        let mut buf = self.output.lock().unwrap();
84        return buf.drain(..num).collect();
85    }
86}
87
88impl Read for MockData {
89    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
90        let mut inp = self.input.lock().unwrap();
91        let bytes = min(buf.len(), inp.len());
92        if bytes == 0 {
93            return Err(io::ErrorKind::WouldBlock.into());
94        }
95        buf[..bytes].copy_from_slice(&inp[..bytes]);
96        inp.drain(..bytes);
97        return Ok(bytes);
98    }
99}
100
101impl Write for MockData {
102    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
103        let mut out = self.output.lock().unwrap();
104        out.extend(buf);
105        return Ok(buf.len());
106    }
107    fn flush(&mut self) -> io::Result<()> {
108        Ok(())
109    }
110}
111
112impl AsyncRead for MockData {}
113impl AsyncWrite for MockData {
114    fn shutdown(&mut self) -> Poll<(), io::Error> {
115        if self.output.lock().unwrap().len() > 0 {
116            Ok(Async::NotReady)
117        } else {
118            Ok(Async::Ready(()))
119        }
120    }
121}
122
123
124#[cfg(test)]
125mod test {
126
127    use {IoBuf, Mock, MockData};
128
129    #[test]
130    fn mock() {
131        let mut buf = IoBuf::new(Mock);
132        buf.read().ok();
133        assert_eq!(&buf.in_buf[..], b"");
134        buf.out_buf.extend(b"hello");
135        assert_eq!(&buf.out_buf[..], b"hello");
136        buf.flush().ok(); // never ready
137        assert_eq!(&buf.out_buf[..], b"hello");
138    }
139
140    #[test]
141    fn mock_data() {
142        let data = MockData::new();
143        let mut buf = IoBuf::new(data.clone());
144        buf.read().ok();
145        assert_eq!(&buf.in_buf[..], b"");
146        data.add_input("test me");
147        buf.read().ok();
148        assert_eq!(&buf.in_buf[..], b"test me");
149
150        buf.out_buf.extend(b"hello");
151        assert_eq!(&buf.out_buf[..], b"hello");
152        buf.flush().ok();
153        assert_eq!(&buf.out_buf[..], b"");
154        assert_eq!(&data.output(..), b"hello");
155    }
156
157}
tk_bufstream/mock.rs

tk_bufstream/
mock.rs
