irc/client/
mock.rs

1use std::{
2    io::{self, Cursor, Read, Write},
3    pin::Pin,
4    task::{Context, Poll},
5};
6use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
7
8/// A fake stream for testing network applications backed by buffers.
9#[derive(Clone, Debug)]
10pub struct MockStream {
11    written: Cursor<Vec<u8>>,
12    received: Cursor<Vec<u8>>,
13}
14
15impl MockStream {
16    /// Creates a new mock stream with nothing to read.
17    pub fn empty() -> MockStream {
18        MockStream::new(&[])
19    }
20
21    /// Creates a new mock stream with the specified bytes to read.
22    pub fn new(initial: &[u8]) -> MockStream {
23        MockStream {
24            written: Cursor::new(vec![]),
25            received: Cursor::new(initial.to_owned()),
26        }
27    }
28
29    /// Gets a slice of bytes representing the data that has been written.
30    pub fn written(&self) -> &[u8] {
31        self.written.get_ref()
32    }
33
34    /// Gets a slice of bytes representing the data that has been received.
35    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}