commonware_runtime/
mocks.rs

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
//! Mock implementations of `commonware-runtime` traits for focused, upstream testing.

use crate::{Error, Sink, Stream};
use bytes::Bytes;
use futures::{channel::mpsc, SinkExt, StreamExt};

/// A mock sink that stores sent messages in a channel.
pub struct MockSink {
    pub sender: mpsc::UnboundedSender<Bytes>,
}

impl MockSink {
    /// Create a new `MockSink` and a corresponding receiver.
    pub fn new() -> (Self, mpsc::UnboundedReceiver<Bytes>) {
        let (sender, receiver) = mpsc::unbounded();
        (Self { sender }, receiver)
    }
}

impl Sink for MockSink {
    async fn send(&mut self, msg: Bytes) -> Result<(), Error> {
        self.sender.send(msg).await.map_err(|_| Error::WriteFailed)
    }
}

/// A mock stream that reads messages from a channel.
pub struct MockStream {
    pub receiver: mpsc::UnboundedReceiver<Bytes>,
}

impl MockStream {
    /// Create a new `MockStream` and a corresponding sender.
    pub fn new() -> (Self, mpsc::UnboundedSender<Bytes>) {
        let (sender, receiver) = mpsc::unbounded();
        (Self { receiver }, sender)
    }
}

impl Stream for MockStream {
    async fn recv(&mut self) -> Result<Bytes, Error> {
        self.receiver.next().await.ok_or(Error::ReadFailed)
    }
}