diff --git a/node/tests/common/mod.rs b/node/tests/common/mod.rs index 2e94b79196..1aa73cb7c2 100644 --- a/node/tests/common/mod.rs +++ b/node/tests/common/mod.rs @@ -14,8 +14,8 @@ // You should have received a copy of the GNU General Public License // along with the snarkOS library. If not, see . +use pea2pea::{protocols::Handshake, Config, Connection, Node, Pea2Pea}; use snarkos_node_messages::{ChallengeRequest, ChallengeResponse, Data, Message, MessageCodec, NodeType, Status}; -use snarkos_node_tcp::{protocols::Handshake, Config, Connection, Tcp, P2P}; use snarkvm::prelude::{Block, FromBytes, Network, Testnet3 as CurrentNetwork}; use futures_util::{sink::SinkExt, TryStreamExt}; @@ -29,14 +29,13 @@ const ALEO_MAXIMUM_FORK_DEPTH: u32 = 4096; #[derive(Clone)] pub struct TestPeer { - // TODO: should be using pea2pea directly (to keep impls separate) - tcp: Tcp, + node: Node, node_type: NodeType, } -impl P2P for TestPeer { - fn tcp(&self) -> &Tcp { - &self.tcp +impl Pea2Pea for TestPeer { + fn node(&self) -> &Node { + &self.node } } @@ -59,7 +58,7 @@ impl TestPeer { pub async fn new(node_type: NodeType) -> Self { let peer = Self { - tcp: Tcp::new(Config { + node: Node::new(Config { listener_ip: Some(IpAddr::V4(Ipv4Addr::LOCALHOST)), max_connections: 200, ..Default::default() @@ -85,7 +84,7 @@ impl TestPeer { #[async_trait::async_trait] impl Handshake for TestPeer { async fn perform_handshake(&self, mut conn: Connection) -> io::Result { - let local_ip = self.tcp().listening_addr().expect("listening address should be present"); + let local_ip = self.node().listening_addr().expect("listening address should be present"); let stream = self.borrow_stream(&mut conn); let mut framed = Framed::new(stream, MessageCodec::::default()); diff --git a/node/tests/new_beacon.rs b/node/tests/new_beacon.rs index 6ac564b8e6..b4c65ad4fa 100644 --- a/node/tests/new_beacon.rs +++ b/node/tests/new_beacon.rs @@ -14,13 +14,21 @@ // You should have received a copy of the GNU General Public License // along with the snarkOS library. If not, see . +#![recursion_limit = "256"] + mod common; +use common::TestPeer; use snarkos_account::Account; -use snarkos_node::{Beacon, Validator}; +use snarkos_node::{Beacon, Client, Prover, Validator}; +use snarkos_node_tcp::P2P; use snarkvm::prelude::{ConsensusMemory, Testnet3 as CurrentNetwork}; -use std::str::FromStr; +use pea2pea::Pea2Pea; + +use std::{io, net::SocketAddr, str::FromStr}; + +/* Node constructors */ async fn beacon() -> Beacon> { Beacon::new( @@ -50,51 +58,101 @@ async fn validator() -> Validator $peer_type:ident),*) => { - mod handshake_initiator_side { - use snarkos_node_tcp::P2P; +// Trait to unify Pea2Pea and P2P traits. +#[async_trait::async_trait] +trait Connect { + fn listening_addr(&self) -> SocketAddr; - $( - #[tokio::test] - async fn $peer_type () { - - let node = $crate::$node_type().await; + async fn connect(&self, target: SocketAddr) -> io::Result<()>; +} - // Spin up a test peer. - let peer = $crate::common::TestPeer::$peer_type().await; +// Implement the `Connect` trait for each node type. +macro_rules! impl_connect { + ($($node_type:ident),*) => { + $( + #[async_trait::async_trait] + impl Connect for $node_type> { + fn listening_addr(&self) -> SocketAddr { + self.tcp().listening_addr().expect("node listener should exist") + } - // Verify the handshake works when the node initiates a connection with the peer. - assert!( - node.tcp().connect(peer.tcp().listening_addr().expect("node listener should exist")).await.is_ok() - ); + async fn connect(&self, target: SocketAddr) -> io::Result<()> + where + Self: P2P, + { + self.tcp().connect(target).await } + } + )* + }; +} - )* - } +impl_connect!(Beacon, Client, Prover, Validator); - }; +// Implement the `Connect` trait for the test peer. +#[async_trait::async_trait] +impl Connect for TestPeer +where + Self: Pea2Pea, +{ + fn listening_addr(&self) -> SocketAddr { + self.node().listening_addr().expect("node listener should exist") + } - ($($node_type:tt <- $peer_type:ident),*) => { - mod handshake_responder_side { - use snarkos_node_tcp::P2P; - use snarkos_node_router::Outbound; + async fn connect(&self, target: SocketAddr) -> io::Result<()> { + self.node().connect(target).await + } +} - $( - #[tokio::test] - async fn $peer_type () { +/* Test case */ - let node = $crate::$node_type().await; +// Asserts a succesful connection was created from initiator to responder. +async fn assert_connect(initiator: T, responder: U) +where + T: Connect, + U: Connect, +{ + assert!(initiator.connect(responder.listening_addr()).await.is_ok()) +} - // Spin up a test peer. - let peer = $crate::common::TestPeer::$peer_type().await; +// Macro to simply construct handshake cases. +// Syntax: +// - (full_node -> test_peer): full node initiates a handshake to the test peer (synthetic node). +// - (full_node <- test_peer): full node receives a handshake initiated by the test peer. +// +// Test naming: full_node::handshake__side::test_peer. +macro_rules! test_handshake { + ($node_type:ident, $peer_type:ident, $is_initiator:expr) => { + #[tokio::test] + async fn $peer_type () { - // Verify the handshake works when the peer initiates a connection with the node. - assert!( - peer.tcp().connect(node.router().tcp().listening_addr().expect("node listener should exist")).await.is_ok() - ); - } + // Spin up a full node. + let node = $crate::$node_type().await; + + // Spin up a test peer (synthetic node). + let peer = $crate::common::TestPeer::$peer_type().await; + if $is_initiator { + $crate::assert_connect(node, peer).await; + } else { + $crate::assert_connect(peer, node).await; + }; + } + }; + + ($($node_type:ident -> $peer_type:ident),*) => { + mod handshake_initiator_side { + $( + test_handshake!($node_type, $peer_type, true); + )* + } + + }; + + ($($node_type:ident <- $peer_type:ident),*) => { + mod handshake_responder_side { + $( + test_handshake!($node_type, $peer_type, false); )* } @@ -102,7 +160,7 @@ macro_rules! test_handshake { } mod beacon { - // Initiator side. + // Initiator side (full node connects to full node). test_handshake! { beacon -> beacon, beacon -> client, @@ -110,7 +168,7 @@ mod beacon { beacon -> prover } - // Responder side. + // Responder side (synthetic peer connects to full node). test_handshake! { beacon <- beacon, beacon <- client, @@ -120,7 +178,7 @@ mod beacon { } mod validator { - // Initiator side. + // Initiator side (full node connects to full node). test_handshake! { validator -> beacon, validator -> client, @@ -128,7 +186,7 @@ mod validator { validator -> prover } - // Responder side. + // Responder side (synthetic peer connects to full node). test_handshake! { validator <- beacon, validator <- client,