From 346d50ebf4e32610fbf966b96ed4a8462ee15971 Mon Sep 17 00:00:00 2001
From: LambdaClass <121504986+lambdaclass-user@users.noreply.github.com>
Date: Mon, 2 Jan 2023 14:36:50 -0300
Subject: [PATCH] feat: test trusted peers functionality (#687)

* Add test for trusted_peer prioritization

* Add test for connect_trusted_nodes_only flag
---
 crates/net/network/src/peers/manager.rs | 77 +++++++++++++++++++++++++
 1 file changed, 77 insertions(+)

diff --git a/crates/net/network/src/peers/manager.rs b/crates/net/network/src/peers/manager.rs
index 283fe1e4e..f899dedf3 100644
--- a/crates/net/network/src/peers/manager.rs
+++ b/crates/net/network/src/peers/manager.rs
@@ -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;
+    }
 }