Merge pull request #2665 from kaimast/sample-batch-with-committee

Add more test-helpers for batch certificates

Author: vicsn

console/types/address/src/lib.rs

@@ -36,6 +36,7 @@ pub use snarkvm_console_types_boolean::Boolean;
 pub use snarkvm_console_types_field::Field;
 pub use snarkvm_console_types_group::Group;
 
+/// The unique ID of an Aleo account, which is derived from its view key.
 #[derive(Copy, Clone, PartialEq, Eq, Hash)]
 pub struct Address<E: Environment> {
     /// The underlying address.

ledger/committee/src/lib.rs

@@ -255,7 +255,10 @@ impl<N: Network> Committee<N> {
 #[cfg(any(test, feature = "test-helpers"))]
 pub mod test_helpers {
     use super::*;
-    use console::{account::Address, prelude::TestRng};
+    use console::{
+        account::{Address, PrivateKey},
+        prelude::TestRng,
+    };
 
     use indexmap::IndexMap;
     use rand_distr::{Distribution, Exp};
@@ -313,6 +316,28 @@ pub mod test_helpers {
         Committee::<CurrentNetwork>::new(round, members).unwrap()
     }
 
+    /// Generates a committee for the specified round, and private keys for its members.
+    ///
+    /// This is slower than `sample_committee_for_round` or `sample_committee_for_round_and_size` due to the additional key generation.
+    pub fn sample_committee_and_keys_for_round(
+        round: u64,
+        num_members: usize,
+        rng: &mut TestRng,
+    ) -> (Committee<CurrentNetwork>, Vec<PrivateKey<CurrentNetwork>>) {
+        // Sample the members.
+        let mut members = IndexMap::with_capacity(num_members);
+        let mut private_keys = Vec::with_capacity(num_members);
+        for _ in 0..num_members {
+            let private_key = PrivateKey::new(rng).unwrap();
+            let address = Address::try_from(private_key).unwrap();
+            let is_open = rng.gen();
+            private_keys.push(private_key);
+            members.insert(address, (2 * MIN_VALIDATOR_STAKE, is_open, 0));
+        }
+        // Return the committee and keys.
+        (Committee::<CurrentNetwork>::new(round, members).unwrap(), private_keys)
+    }
+
     /// Samples a random committee for a given round and members.
     pub fn sample_committee_for_round_and_members(
         round: u64,

ledger/narwhal/batch-certificate/src/lib.rs

@@ -172,25 +172,42 @@ pub mod test_helpers {
         sample_batch_certificate_for_round_with_previous_certificate_ids(round, certificate_ids, rng)
     }
 
-    /// Returns a sample batch certificate with a given round; the rest is sampled at random.
+    /// Returns a sample batch certificate with a given round and the given certificate ids as predecessors; the rest is sampled at random.
     pub fn sample_batch_certificate_for_round_with_previous_certificate_ids(
         round: u64,
         previous_certificate_ids: IndexSet<Field<CurrentNetwork>>,
         rng: &mut TestRng,
+    ) -> BatchCertificate<CurrentNetwork> {
+        let committee: Vec<_> = (0..5).map(|_| PrivateKey::new(rng).unwrap()).collect();
+        sample_batch_certificate_for_round_with_committee(
+            round,
+            previous_certificate_ids,
+            &committee[0],
+            &committee[1..],
+            rng,
+        )
+    }
+
+    /// Same as `sample_batch_certificate_for_round_with_previous_certificate_ids`, but also allows you to set the private keys that sign the certificate.
+    pub fn sample_batch_certificate_for_round_with_committee(
+        round: u64,
+        previous_certificate_ids: IndexSet<Field<CurrentNetwork>>,
+        author: &PrivateKey<CurrentNetwork>,
+        signers: &[PrivateKey<CurrentNetwork>],
+        rng: &mut TestRng,
     ) -> BatchCertificate<CurrentNetwork> {
         // Sample a batch header.
         let batch_header =
-            narwhal_batch_header::test_helpers::sample_batch_header_for_round_with_previous_certificate_ids(
+            narwhal_batch_header::test_helpers::sample_batch_header_for_round_and_key_with_previous_certificate_ids(
                 round,
+                author,
                 previous_certificate_ids,
                 rng,
             );
-        // Sample a list of signatures.
-        let mut signatures = IndexSet::with_capacity(5);
-        for _ in 0..5 {
-            let private_key = PrivateKey::new(rng).unwrap();
-            signatures.insert(private_key.sign(&[batch_header.batch_id()], rng).unwrap());
-        }
+        // Generate the endorsements.
+        let signatures: IndexSet<_> =
+            signers.iter().map(|private_key| private_key.sign(&[batch_header.batch_id()], rng).unwrap()).collect();
+
         // Return the batch certificate.
         BatchCertificate::from(batch_header, signatures).unwrap()
     }

ledger/narwhal/batch-header/src/lib.rs

@@ -276,6 +276,22 @@ pub mod test_helpers {
     ) -> BatchHeader<CurrentNetwork> {
         // Sample a private key.
         let private_key = PrivateKey::new(rng).unwrap();
+        // Generate a new certificated with the key as its author.
+        sample_batch_header_for_round_and_key_with_previous_certificate_ids(
+            round,
+            &private_key,
+            previous_certificate_ids,
+            rng,
+        )
+    }
+
+    /// Returns a sample batch header with a given round, author key, and set of previous certificate IDs; the rest is sampled at random.
+    pub fn sample_batch_header_for_round_and_key_with_previous_certificate_ids(
+        round: u64,
+        private_key: &PrivateKey<CurrentNetwork>,
+        previous_certificate_ids: IndexSet<Field<CurrentNetwork>>,
+        rng: &mut TestRng,
+    ) -> BatchHeader<CurrentNetwork> {
         // Sample the committee ID.
         let committee_id = Field::<CurrentNetwork>::rand(rng);
         // Sample transmission IDs.
@@ -284,7 +300,7 @@ pub mod test_helpers {
         // Checkpoint the timestamp for the batch.
         let timestamp = OffsetDateTime::now_utc().unix_timestamp();
         // Return the batch header.
-        BatchHeader::new(&private_key, round, timestamp, committee_id, transmission_ids, previous_certificate_ids, rng)
+        BatchHeader::new(private_key, round, timestamp, committee_id, transmission_ids, previous_certificate_ids, rng)
             .unwrap()
     }