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    pub fn add_peer(&self, peer_id: PeerId, addr: SocketAddr) {
36        self.send(PeerCommand::Add(peer_id, addr));
37    }
38
39    /// Removes a peer from the set.
40    pub fn remove_peer(&self, peer_id: PeerId) {
41        self.send(PeerCommand::Remove(peer_id));
42    }
43
44    /// Send a reputation change for the given peer.
45    pub fn reputation_change(&self, peer_id: PeerId, kind: ReputationChangeKind) {
46        self.send(PeerCommand::ReputationChange(peer_id, kind));
47    }
48
49    /// Returns a peer by its [`PeerId`], or `None` if the peer is not in the peer set.
50    pub async fn peer_by_id(&self, peer_id: PeerId) -> Option<Peer> {
51        let (tx, rx) = oneshot::channel();
52        self.send(PeerCommand::GetPeer(peer_id, tx));
53
54        rx.await.unwrap_or(None)
55    }
56
57    /// Returns all peers in the peerset.
58    pub async fn all_peers(&self) -> Vec<NodeRecord> {
59        let (tx, rx) = oneshot::channel();
60        self.send(PeerCommand::GetPeers(tx));
61
62        rx.await.unwrap_or_default()
63    }
64}
65
66/// Commands the `PeersManager` listens for.
67#[derive(Debug)]
68pub enum PeerCommand {
69    /// Command for manually add
70    Add(PeerId, SocketAddr),
71    /// Remove a peer from the set
72    ///
73    /// If currently connected this will disconnect the session
74    Remove(PeerId),
75    /// Apply a reputation change to the given peer.
76    ReputationChange(PeerId, ReputationChangeKind),
77    /// Get information about a peer
78    GetPeer(PeerId, oneshot::Sender<Option<Peer>>),
79    /// Get node information on all peers
80    GetPeers(oneshot::Sender<Vec<NodeRecord>>),
81}
reth_network_api/test_utils/peers_manager.rs

reth_network_api/test_utils/
peers_manager.rs
