feat: test trusted peers functionality (#687)

* Add test for trusted_peer prioritization

* Add test for connect_trusted_nodes_only flag
This commit is contained in:
LambdaClass
2023-01-02 14:36:50 -03:00
committed by GitHub
parent 13441ff299
commit 346d50ebf4

View File

@ -940,6 +940,7 @@ mod test {
session::PendingSessionHandshakeError,
PeersConfig,
};
use reth_discv4::NodeRecord;
use reth_eth_wire::{
errors::{EthHandshakeError, EthStreamError, P2PHandshakeError, P2PStreamError},
DisconnectReason,
@ -1489,4 +1490,80 @@ mod test {
assert_eq!(info.num_inbound, 0);
assert_eq!(info.num_outbound, 0);
}
#[tokio::test]
async fn test_trusted_peers_are_prioritized() {
let trusted_peer = PeerId::random();
let trusted_sock = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
let config = PeersConfig::default().with_trusted_nodes(HashSet::from([NodeRecord {
address: IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)),
tcp_port: 8008,
udp_port: 8008,
id: trusted_peer,
}]));
let mut peers = PeersManager::new(config);
let basic_peer = PeerId::random();
let basic_sock = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8009);
peers.add_discovered_node(basic_peer, basic_sock);
match event!(peers) {
PeerAction::PeerAdded(peer_id) => {
assert_eq!(peer_id, basic_peer);
}
_ => unreachable!(),
}
match event!(peers) {
PeerAction::Connect { peer_id, remote_addr } => {
assert_eq!(peer_id, trusted_peer);
assert_eq!(remote_addr, trusted_sock);
}
_ => unreachable!(),
}
match event!(peers) {
PeerAction::Connect { peer_id, remote_addr } => {
assert_eq!(peer_id, basic_peer);
assert_eq!(remote_addr, basic_sock);
}
_ => unreachable!(),
}
}
#[tokio::test]
async fn test_connect_trusted_nodes_only() {
let trusted_peer = PeerId::random();
let trusted_sock = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8008);
let config = PeersConfig::default()
.with_trusted_nodes(HashSet::from([NodeRecord {
address: IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)),
tcp_port: 8008,
udp_port: 8008,
id: trusted_peer,
}]))
.with_connect_trusted_nodes_only(true);
let mut peers = PeersManager::new(config);
let basic_peer = PeerId::random();
let basic_sock = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 2)), 8009);
peers.add_discovered_node(basic_peer, basic_sock);
match event!(peers) {
PeerAction::PeerAdded(peer_id) => {
assert_eq!(peer_id, basic_peer);
}
_ => unreachable!(),
}
match event!(peers) {
PeerAction::Connect { peer_id, remote_addr } => {
assert_eq!(peer_id, trusted_peer);
assert_eq!(remote_addr, trusted_sock);
}
_ => unreachable!(),
}
poll_fn(|cx| {
assert!(peers.poll(cx).is_pending());
Poll::Ready(())
})
.await;
}
}