feat: generate real workers to proptest, test module chores

Author: vvp

node/narwhal/src/gateway.rs

@@ -789,11 +789,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;
@@ -808,12 +816,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 {
@@ -826,6 +837,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]
@@ -857,18 +892,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}");

node/narwhal/src/helpers/committee.rs

@@ -103,7 +103,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;
@@ -115,7 +115,7 @@ pub mod tests {
 
     type N = Testnet3;
 
-    #[derive(Arbitrary, Debug)]
+    #[derive(Arbitrary, Debug, Clone)]
     pub struct CommitteeInput {
         #[strategy(0u64..)]
         pub round: u64,

node/narwhal/src/helpers/storage.rs

@@ -329,11 +329,12 @@ impl<N: Network> Storage<N> {
 }
 
 #[cfg(test)]
-mod tests {
+pub mod storage_tests {
     use super::*;
     use snarkvm::prelude::TestRng;
 
     use indexmap::indexset;
+    use test_strategy::Arbitrary;
 
     type CurrentNetwork = snarkvm::prelude::Testnet3;
 
@@ -418,4 +419,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)
+        }
+    }
 }

node/narwhal/src/worker.rs

@@ -289,3 +289,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,
+    }
+}