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
use std::sync::{Arc, Mutex};
#[derive(Default, Clone)]
pub struct TestStream {
read: Arc<Mutex<Vec<u8>>>,
write: Arc<Mutex<Vec<u8>>>,
buf: Arc<Mutex<Vec<u8>>>,
}
impl TestStream {
pub fn new() -> Self {
Self::default()
}
pub fn drain_buffer(&mut self) -> Option<String> {
let mut w = self.write.lock().unwrap();
if w.is_empty() {
return None;
}
String::from_utf8(w.drain(..).collect::<Vec<_>>()).ok()
}
pub fn write_message<S: AsRef<[u8]>>(&mut self, data: S) {
let mut w = self.read.lock().unwrap();
w.extend_from_slice(&data.as_ref());
w.extend_from_slice(&[b'\r', b'\n']);
}
}
impl std::io::Read for TestStream {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
use std::io::{Cursor, Error, ErrorKind};
let mut read = self.read.lock().unwrap();
if read.ends_with(&[b'\r', b'\n']) {
let w = read.drain(..).collect::<Vec<_>>();
return Cursor::new(w).read(buf);
}
Err(Error::new(ErrorKind::WouldBlock, ""))
}
}
impl std::io::Write for TestStream {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
if buf == b"\r\n" {
return Ok(buf.len());
}
self.buf.lock().unwrap().extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
let mut buf = self.buf.lock().unwrap();
if buf.is_empty() {
return Ok(());
}
self.write.lock().unwrap().extend(buf.drain(..));
Ok(())
}
}