fedimint_core/net/peers/
fake.rs

1use async_channel::bounded;
2/// Fake (channel-based) implementation of [`super::DynP2PConnections`].
3use async_trait::async_trait;
4use fedimint_core::PeerId;
5use fedimint_core::net::peers::{DynP2PConnections, IP2PConnections};
6
7use crate::net::peers::Recipient;
8
9#[derive(Clone)]
10struct FakePeerConnections<M> {
11    tx: async_channel::Sender<M>,
12    rx: async_channel::Receiver<M>,
13    peer: PeerId,
14}
15
16#[async_trait]
17impl<M: Clone + Send + 'static> IP2PConnections<M> for FakePeerConnections<M> {
18    async fn send(&self, recipient: Recipient, msg: M) {
19        assert_eq!(recipient, Recipient::Peer(self.peer));
20
21        // If the peer is gone, just pretend we are going to resend
22        // the msg eventually, even if it will never happen.
23        self.tx.send(msg).await.ok();
24    }
25
26    fn try_send(&self, recipient: Recipient, msg: M) {
27        assert_eq!(recipient, Recipient::Peer(self.peer));
28
29        // If the peer is gone, just pretend we are going to resend
30        // the msg eventually, even if it will never happen.
31        self.tx.try_send(msg).ok();
32    }
33
34    async fn receive(&self) -> Option<(PeerId, M)> {
35        self.rx.recv().await.map(|msg| (self.peer, msg)).ok()
36    }
37
38    async fn receive_from_peer(&self, _peer: PeerId) -> Option<M> {
39        self.rx.recv().await.ok()
40    }
41}
42
43/// Create a fake link between `peer1` and `peer2` for test purposes
44///
45/// `buf_size` controls the size of the `tokio::mpsc::channel` used
46/// under the hood (both ways).
47pub fn make_fake_peer_connection<M: Clone + Send + 'static>(
48    peer_1: PeerId,
49    peer_2: PeerId,
50    buf_size: usize,
51) -> (DynP2PConnections<M>, DynP2PConnections<M>) {
52    let (tx1, rx1) = bounded(buf_size);
53    let (tx2, rx2) = bounded(buf_size);
54
55    let c_1 = FakePeerConnections {
56        tx: tx1,
57        rx: rx2,
58        peer: peer_2,
59    };
60
61    let c_2 = FakePeerConnections {
62        tx: tx2,
63        rx: rx1,
64        peer: peer_1,
65    };
66
67    (c_1.into_dyn(), c_2.into_dyn())
68}