1  // Copyright (c) 2025 ADnet Contributors
  2  // This file is part of the AlphaOS library.
  3  
  4  // Licensed under the Apache License, Version 2.0 (the "License");
  5  // you may not use this file except in compliance with the License.
  6  // You may obtain a copy of the License at:
  7  
  8  // http://www.apache.org/licenses/LICENSE-2.0
  9  
 10  // Unless required by applicable law or agreed to in writing, software
 11  // distributed under the License is distributed on an "AS IS" BASIS,
 12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 13  // See the License for the specific language governing permissions and
 14  // limitations under the License.
 15  
 16  #[allow(dead_code)]
 17  pub mod router;
 18  pub use router::*;
 19  
 20  use std::{
 21      net::{IpAddr, Ipv4Addr, SocketAddr},
 22      sync::Arc,
 23  };
 24  
 25  use alpha_std::StorageMode;
 26  use alphaos_account::Account;
 27  use alphaos_node_bft_ledger_service::MockLedgerService;
 28  use alphaos_node_network::NodeType;
 29  use alphaos_node_router::Router;
 30  use alphavm::{
 31      prelude::{FromBytes, MainnetV0 as CurrentNetwork, Network, PrivateKey, block::Block},
 32      utilities::TestRng,
 33  };
 34  
 35  /// A helper macro to print the TCP listening address, along with the connected and connecting peers.
 36  #[macro_export]
 37  macro_rules! print_tcp {
 38      ($node:expr) => {
 39          println!(
 40              "{}: Active - {:?}, Pending - {:?}",
 41              $node.local_ip(),
 42              $node.tcp().connected_addrs(),
 43              $node.tcp().connecting_addrs()
 44          );
 45      };
 46  }
 47  
 48  /// Returns a fixed account.
 49  pub fn sample_account(rng: &mut TestRng) -> Account<CurrentNetwork> {
 50      let private_key = PrivateKey::<CurrentNetwork>::new(rng).unwrap();
 51      Account::<CurrentNetwork>::try_from(&private_key).unwrap()
 52  }
 53  
 54  /// Loads the current network's genesis block.
 55  pub fn sample_genesis_block<N: Network>() -> Block<N> {
 56      Block::<N>::from_bytes_le(N::genesis_bytes()).unwrap()
 57  }
 58  
 59  /// Initializes a client router. Setting the `listening_port = 0` will result in a random port being assigned.
 60  pub async fn client(listening_port: u16, max_peers: u16, rng: &mut TestRng) -> TestRouter<CurrentNetwork> {
 61      let committee = alphavm::ledger::committee::test_helpers::sample_committee(rng);
 62      let ledger_service = Arc::new(MockLedgerService::new(committee));
 63      Router::new(
 64          SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), listening_port),
 65          NodeType::Client,
 66          sample_account(rng),
 67          ledger_service,
 68          &[],
 69          max_peers,
 70          false,
 71          StorageMode::new_test(None),
 72          true,
 73      )
 74      .await
 75      .expect("couldn't create client router")
 76      .into()
 77  }
 78  
 79  /// Initializes a prover router. Setting the `listening_port = 0` will result in a random port being assigned.
 80  #[allow(dead_code)]
 81  pub async fn prover(listening_port: u16, max_peers: u16, rng: &mut TestRng) -> TestRouter<CurrentNetwork> {
 82      let committee = alphavm::ledger::committee::test_helpers::sample_committee(rng);
 83      let ledger_service = Arc::new(MockLedgerService::new(committee));
 84      Router::new(
 85          SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), listening_port),
 86          NodeType::Prover,
 87          sample_account(rng),
 88          ledger_service,
 89          &[],
 90          max_peers,
 91          false,
 92          StorageMode::new_test(None),
 93          true,
 94      )
 95      .await
 96      .expect("couldn't create prover router")
 97      .into()
 98  }
 99  
100  /// Initializes a validator router. Setting the `listening_port = 0` will result in a random port being assigned.
101  #[allow(dead_code)]
102  pub async fn validator(
103      listening_port: u16,
104      max_peers: u16,
105      trusted_peers: &[SocketAddr],
106      trusted_peers_only: bool,
107      rng: &mut TestRng,
108  ) -> TestRouter<CurrentNetwork> {
109      let committee = alphavm::ledger::committee::test_helpers::sample_committee(rng);
110      let ledger_service = Arc::new(MockLedgerService::new(committee));
111      Router::new(
112          SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), listening_port),
113          NodeType::Validator,
114          sample_account(rng),
115          ledger_service,
116          trusted_peers,
117          max_peers,
118          trusted_peers_only,
119          StorageMode::new_test(None),
120          true,
121      )
122      .await
123      .expect("couldn't create validator router")
124      .into()
125  }