Skip to content

Commit

Permalink
feat: generate real workers to proptest, test module chores
Browse files Browse the repository at this point in the history
  • Loading branch information
vvp committed Jul 6, 2023
1 parent 7a896a6 commit 326d412
Show file tree
Hide file tree
Showing 4 changed files with 76 additions and 12 deletions.
55 changes: 46 additions & 9 deletions node/narwhal/src/gateway.rs
Original file line number Diff line number Diff line change
Expand Up @@ -795,11 +795,19 @@ impl<N: Network> Gateway<N> {
}
}

mod tests {
#[cfg(test)]
pub mod gateway_tests {
use crate::{
helpers::tests::{CommitteeInput, Validator},
helpers::{
committee_tests::{CommitteeInput, Validator},
init_worker_channels,
storage_tests::StorageInput,
WorkerSender,
},
Gateway,
Worker,
MAX_COMMITTEE_SIZE,
MAX_WORKERS,
MEMORY_POOL_PORT,
};
use indexmap::IndexMap;
Expand All @@ -814,12 +822,15 @@ mod tests {

type N = Testnet3;

#[derive(Arbitrary, Debug)]
struct GatewayInput {
#[derive(Arbitrary, Debug, Clone)]
pub struct GatewayInput {
#[filter(CommitteeInput::is_valid)]
committee_input: CommitteeInput,
node_validator: Validator,
dev: Option<u8>,
pub committee_input: CommitteeInput,
pub node_validator: Validator,
pub dev: Option<u8>,
#[strategy(0..MAX_WORKERS)]
pub workers_count: u8,
pub worker_storage: StorageInput,
}

impl GatewayInput {
Expand All @@ -832,6 +843,30 @@ mod tests {
};
Gateway::new(Arc::new(RwLock::new(committee)), account, dev).unwrap()
}

pub async fn generate_workers(
&self,
gateway: &Gateway<N>,
) -> (IndexMap<u8, Worker<N>>, IndexMap<u8, WorkerSender<N>>) {
// Construct a map of the worker senders.
let mut tx_workers = IndexMap::new();
let mut workers = IndexMap::new();

// Initialize the workers.
for id in 0..self.workers_count {
// Construct the worker channels.
let (tx_worker, rx_worker) = init_worker_channels();
// Construct the worker instance.
let mut worker = Worker::new(id, gateway.clone(), self.worker_storage.to_storage()).unwrap();
// Run the worker instance.
worker.run(rx_worker).await.unwrap();

// Add the worker and the worker sender to maps
workers.insert(id, worker);
tx_workers.insert(id, tx_worker);
}
(workers, tx_workers)
}
}

#[proptest]
Expand Down Expand Up @@ -863,18 +898,20 @@ mod tests {
#[proptest(async = "tokio")]
async fn gateway_start(#[filter(|x| x.dev.is_some())] input: GatewayInput) {
let Some(dev) = input.dev else { unreachable!() };
println!("{input:?}");
let mut gateway = input.to_gateway();
let tcp_config = gateway.tcp().config();
assert_eq!(tcp_config.listener_ip, Some(IpAddr::V4(Ipv4Addr::LOCALHOST)));
assert_eq!(tcp_config.desired_listening_port, Some(MEMORY_POOL_PORT + (dev as u16)));

match gateway.run(IndexMap::new()).await {
let (workers, worker_senders) = input.generate_workers(&gateway).await;
match gateway.run(worker_senders).await {
Ok(_) => {
assert_eq!(
gateway.local_ip(),
SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), MEMORY_POOL_PORT + (dev as u16))
);
assert_eq!(gateway.num_workers(), 0);
assert_eq!(gateway.num_workers(), workers.len() as u8);
}
Err(err) => {
panic!("Shouldn't fail because of {err}");
Expand Down
4 changes: 2 additions & 2 deletions node/narwhal/src/helpers/committee.rs
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ impl<N: Network> Committee<N> {
}

#[cfg(test)]
pub mod tests {
pub mod committee_tests {
use crate::helpers::Committee;
use anyhow::Result;
use indexmap::IndexMap;
Expand All @@ -131,7 +131,7 @@ pub mod tests {

type N = Testnet3;

#[derive(Arbitrary, Debug)]
#[derive(Arbitrary, Debug, Clone)]
pub struct CommitteeInput {
#[strategy(0u64..)]
pub round: u64,
Expand Down
14 changes: 13 additions & 1 deletion node/narwhal/src/helpers/storage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -406,7 +406,7 @@ impl<N: Network> Storage<N> {
}

#[cfg(test)]
mod tests {
pub mod storage_tests {
use super::*;
use snarkvm::{
ledger::narwhal::Data,
Expand All @@ -416,6 +416,7 @@ mod tests {
use ::bytes::Bytes;

use indexmap::indexset;
use test_strategy::Arbitrary;

type CurrentNetwork = snarkvm::prelude::Testnet3;

Expand Down Expand Up @@ -560,4 +561,15 @@ mod tests {
// Ensure the certificate is no longer stored in the round.
assert!(storage.get_certificates_for_round(round).is_empty());
}

#[derive(Arbitrary, Debug, Clone)]
pub struct StorageInput {
pub gc_rounds: u64,
}

impl StorageInput {
pub fn to_storage(&self) -> Storage<CurrentNetwork> {
Storage::new(self.gc_rounds)
}
}
}
15 changes: 15 additions & 0 deletions node/narwhal/src/worker.rs
Original file line number Diff line number Diff line change
Expand Up @@ -266,3 +266,18 @@ impl<N: Network> Worker<N> {
self.handles.lock().iter().for_each(|handle| handle.abort());
}
}

#[cfg(test)]
mod worker_tests {
use crate::{gateway_tests::GatewayInput, helpers::storage_tests::StorageInput, Gateway, Worker, MAX_WORKERS};
use snarkvm::prelude::Testnet3;
use test_strategy::Arbitrary;

type N = Testnet3;

#[derive(Arbitrary, Debug, Clone)]
pub struct WorkerInput {
pub id: u8,
pub storage: StorageInput,
}
}

0 comments on commit 326d412

Please sign in to comment.