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,