reth_e2e_test_utils/
network.rs

Help
1use futures_util::StreamExt;
2use reth_network_api::{
3    events::PeerEvent, test_utils::PeersHandleProvider, NetworkEvent, NetworkEventListenerProvider,
4    PeerRequest, PeersInfo,
5};
6use reth_network_peers::{NodeRecord, PeerId};
7use reth_tokio_util::EventStream;
8use reth_tracing::tracing::info;
9
10/// Helper for network operations
11#[derive(Debug)]
12pub struct NetworkTestContext<Network: NetworkEventListenerProvider> {
13    network_events: EventStream<NetworkEvent<PeerRequest<Network::Primitives>>>,
14    network: Network,
15}
16
17impl<Network> NetworkTestContext<Network>
18where
19    Network: NetworkEventListenerProvider + PeersInfo + PeersHandleProvider,
20{
21    /// Creates a new network helper
22    pub fn new(network: Network) -> Self {
23        let network_events = network.event_listener();
24        Self { network_events, network }
25    }
26
27    /// Adds a peer to the network node via network handle
28    pub async fn add_peer(&mut self, node_record: NodeRecord) {
29        self.network.peers_handle().add_peer(node_record.id, node_record.tcp_addr());
30
31        match self.network_events.next().await {
32            Some(NetworkEvent::Peer(PeerEvent::PeerAdded(_))) => (),
33            ev => panic!("Expected a peer added event, got: {ev:?}"),
34        }
35    }
36
37    /// Returns the network node record
38    pub fn record(&self) -> NodeRecord {
39        self.network.local_node_record()
40    }
41
42    /// Awaits the next event for an established session.
43    pub async fn next_session_established(&mut self) -> Option<PeerId> {
44        while let Some(ev) = self.network_events.next().await {
45            match ev {
46                NetworkEvent::ActivePeerSession { info, .. } |
47                NetworkEvent::Peer(PeerEvent::SessionEstablished(info)) => {
48                    let peer_id = info.peer_id;
49                    info!("Session established with peer: {:?}", peer_id);
50                    return Some(peer_id)
51                }
52                _ => {}
53            }
54        }
55        None
56    }
57}
reth_e2e_test_utils/network.rs

reth_e2e_test_utils/
network.rs
