reth_network_api/test_utils/
peers_manager.rs

1//! Interaction with `reth_network::PeersManager`, for integration testing. Otherwise
2//! `reth_network::NetworkManager` manages `reth_network::PeersManager`.
3
4use std::net::SocketAddr;
5
6use derive_more::Constructor;
7use reth_network_peers::{NodeRecord, PeerId};
8use reth_network_types::{Peer, ReputationChangeKind};
9use tokio::sync::{mpsc, oneshot};
10
11/// Provides an API for managing the peers of the network.
12#[auto_impl::auto_impl(&, Arc)]
13pub trait PeersHandleProvider {
14    /// Returns the [`PeersHandle`] that can be cloned and shared.
15    ///
16    /// The [`PeersHandle`] can be used to interact with the network's peer set.
17    fn peers_handle(&self) -> &PeersHandle;
18}
19
20/// A communication channel to the `PeersManager` to apply manual changes to the peer set.
21#[derive(Clone, Debug, Constructor)]
22pub struct PeersHandle {
23    /// Sender half of command channel back to the `PeersManager`
24    manager_tx: mpsc::UnboundedSender<PeerCommand>,
25}
26
27// === impl PeersHandle ===
28
29impl PeersHandle {
30    fn send(&self, cmd: PeerCommand) {
31        let _ = self.manager_tx.send(cmd);
32    }
33
34    /// Adds a peer to the set.
35    ///
36    /// If the peer already exists, then this will update only the provided address, this is
37    /// equivalent to discovering a peer.
38    pub fn add_peer(&self, peer_id: PeerId, addr: SocketAddr) {
39        self.send(PeerCommand::Add(peer_id, addr));
40    }
41
42    /// Removes a peer from the set.
43    pub fn remove_peer(&self, peer_id: PeerId) {
44        self.send(PeerCommand::Remove(peer_id));
45    }
46
47    /// Send a reputation change for the given peer.
48    pub fn reputation_change(&self, peer_id: PeerId, kind: ReputationChangeKind) {
49        self.send(PeerCommand::ReputationChange(peer_id, kind));
50    }
51
52    /// Returns a peer by its [`PeerId`], or `None` if the peer is not in the peer set.
53    pub async fn peer_by_id(&self, peer_id: PeerId) -> Option<Peer> {
54        let (tx, rx) = oneshot::channel();
55        self.send(PeerCommand::GetPeer(peer_id, tx));
56
57        rx.await.unwrap_or(None)
58    }
59
60    /// Returns all peers in the peerset.
61    pub async fn all_peers(&self) -> Vec<NodeRecord> {
62        let (tx, rx) = oneshot::channel();
63        self.send(PeerCommand::GetPeers(tx));
64
65        rx.await.unwrap_or_default()
66    }
67}
68
69/// Commands the `PeersManager` listens for.
70#[derive(Debug)]
71pub enum PeerCommand {
72    /// Command for manually add
73    Add(PeerId, SocketAddr),
74    /// Remove a peer from the set
75    ///
76    /// If currently connected this will disconnect the session
77    Remove(PeerId),
78    /// Apply a reputation change to the given peer.
79    ReputationChange(PeerId, ReputationChangeKind),
80    /// Get information about a peer
81    GetPeer(PeerId, oneshot::Sender<Option<Peer>>),
82    /// Get node information on all peers
83    GetPeers(oneshot::Sender<Vec<NodeRecord>>),
84}
reth_network_api/test_utils/peers_manager.rs

reth_network_api/test_utils/
peers_manager.rs
