Add MithrilFixture & its builder to common utilities

Author: Alenar

mithril-common/src/lib.rs

@@ -23,6 +23,7 @@ pub mod entities;
 pub mod fake_data;
 pub mod sqlite;
 pub mod store;
+pub mod test_utils;
 
 pub use beacon_provider::{BeaconProvider, BeaconProviderError, BeaconProviderImpl};
 pub use entities::{CardanoNetwork, MagicId};

mithril-common/src/test_utils/fixture_builder.rs

@@ -0,0 +1,196 @@
+use kes_summed_ed25519::{kes::Sum6Kes, traits::KesSk};
+use rand_chacha::ChaCha20Rng;
+use rand_core::{RngCore, SeedableRng};
+
+use crate::{
+    crypto_helper::{
+        tests_setup, tests_setup::setup_temp_directory_for_signer, ColdKeyGenerator, OpCert,
+        ProtocolStakeDistribution, SerDeShelleyFileFormat,
+    },
+    entities::{PartyId, ProtocolParameters},
+    fake_data,
+    test_utils::mithril_fixture::MithrilFixture,
+};
+
+use super::mithril_fixture::SignerFixture;
+
+/// A builder of mithril types.
+pub struct MithrilFixtureBuilder {
+    protocol_parameters: ProtocolParameters,
+    enable_signers_certification: bool,
+    number_of_signers: usize,
+    stake_distribution_generation_method: StakeDistributionGenerationMethod,
+}
+
+impl Default for MithrilFixtureBuilder {
+    fn default() -> Self {
+        Self {
+            protocol_parameters: fake_data::protocol_parameters(),
+            enable_signers_certification: true,
+            number_of_signers: 5,
+            stake_distribution_generation_method:
+                StakeDistributionGenerationMethod::RandomDistribution { seed: [0u8; 32] },
+        }
+    }
+}
+
+/// Methods that can be used to generate the stake distribution.
+pub enum StakeDistributionGenerationMethod {
+    /// Each party will have a random stake.
+    RandomDistribution {
+        /// The randomizer seed
+        seed: [u8; 32],
+    },
+}
+
+impl MithrilFixtureBuilder {
+    /// Set the protocol_parameters.
+    pub fn with_protocol_parameters(mut self, protocol_parameters: ProtocolParameters) -> Self {
+        self.protocol_parameters = protocol_parameters;
+        self
+    }
+
+    /// Set the number of signers that will be generated.
+    pub fn with_signers(mut self, number_of_signers: usize) -> Self {
+        self.number_of_signers = number_of_signers;
+        self
+    }
+
+    /// If set the generated signers won't be certified (meaning that they won't
+    /// have a operational certificate).
+    pub fn disable_signers_certification(mut self) -> Self {
+        self.enable_signers_certification = false;
+        self
+    }
+
+    /// Set the generation method used to compute the stake distribution.
+    pub fn with_stake_distribution(
+        mut self,
+        stake_distribution_generation_method: StakeDistributionGenerationMethod,
+    ) -> Self {
+        self.stake_distribution_generation_method = stake_distribution_generation_method;
+        self
+    }
+
+    /// Transform the specified parameters to a [MithrilFixture].
+    pub fn build(self) -> MithrilFixture {
+        let protocol_stake_distribution = self.generate_stake_distribution();
+        let signers = tests_setup::setup_signers_from_stake_distribution(
+            &protocol_stake_distribution,
+            &self.protocol_parameters.clone().into(),
+        )
+        .into_iter()
+        .map(
+            |(signer_with_stake, protocol_signer, protocol_initializer)| SignerFixture {
+                signer_with_stake,
+                protocol_signer,
+                protocol_initializer,
+            },
+        )
+        .collect();
+
+        MithrilFixture::new(
+            self.protocol_parameters,
+            signers,
+            protocol_stake_distribution,
+        )
+    }
+
+    fn generate_stake_distribution(&self) -> ProtocolStakeDistribution {
+        let mut kes_keys_seed = [0u8; 32];
+        let signers_party_ids = (0..self.number_of_signers).into_iter().map(|party_index| {
+            if self.enable_signers_certification {
+                build_party_with_operational_certificate(party_index, &mut kes_keys_seed)
+            } else {
+                format!("{:<032}", party_index)
+            }
+        });
+
+        match self.stake_distribution_generation_method {
+            StakeDistributionGenerationMethod::RandomDistribution { seed } => {
+                let mut stake_rng = ChaCha20Rng::from_seed(seed);
+
+                signers_party_ids
+                    .map(|party_id| {
+                        let stake = 1 + stake_rng.next_u64() % 999;
+                        (party_id, stake)
+                    })
+                    .collect::<Vec<_>>()
+            }
+        }
+    }
+}
+
+fn build_party_with_operational_certificate(
+    party_index: usize,
+    kes_key_seed: &mut [u8],
+) -> PartyId {
+    let keypair = ColdKeyGenerator::create_deterministic_keypair([party_index as u8; 32]);
+    let (kes_secret_key, kes_verification_key) = Sum6Kes::keygen(kes_key_seed);
+    let operational_certificate = OpCert::new(kes_verification_key, 0, 0, keypair);
+    let party_id = operational_certificate
+        .compute_protocol_party_id()
+        .expect("compute protocol party id should not fail");
+    let temp_dir = setup_temp_directory_for_signer(&party_id, true)
+        .expect("setup temp directory should return a value");
+    if !temp_dir.join("kes.sk").exists() {
+        kes_secret_key
+            .to_file(temp_dir.join("kes.sk"))
+            .expect("KES secret key file export should not fail");
+    }
+    if !temp_dir.join("opcert.cert").exists() {
+        operational_certificate
+            .to_file(temp_dir.join("opcert.cert"))
+            .expect("operational certificate file export should not fail");
+    }
+    party_id
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::collections::BTreeSet;
+
+    #[test]
+    fn with_protocol_params() {
+        let protocol_parameters = ProtocolParameters::new(1, 10, 0.56);
+        let result = MithrilFixtureBuilder::default()
+            .with_protocol_parameters(protocol_parameters.clone())
+            .build();
+
+        assert_eq!(protocol_parameters, result.protocol_parameters());
+    }
+
+    #[test]
+    fn with_signers() {
+        let result = MithrilFixtureBuilder::default().with_signers(4).build();
+
+        assert_eq!(4, result.signers_with_stake().len());
+    }
+
+    #[test]
+    fn random_stake_distribution_generates_as_many_signers_as_parties() {
+        let result = MithrilFixtureBuilder::default()
+            .with_stake_distribution(StakeDistributionGenerationMethod::RandomDistribution {
+                seed: [0u8; 32],
+            })
+            .with_signers(4)
+            .build();
+
+        assert_eq!(4, result.stake_distribution().len());
+    }
+
+    #[test]
+    fn each_parties_generated_with_random_stake_distribution_have_different_stakes() {
+        let result = MithrilFixtureBuilder::default()
+            .with_stake_distribution(StakeDistributionGenerationMethod::RandomDistribution {
+                seed: [0u8; 32],
+            })
+            .with_signers(5)
+            .build();
+        let stakes = result.stake_distribution();
+
+        // BtreeSet dedup values
+        assert_eq!(stakes.len(), BTreeSet::from_iter(stakes.values()).len());
+    }
+}

mithril-common/src/test_utils/mithril_fixture.rs

@@ -0,0 +1,71 @@
+use crate::{
+    crypto_helper::{ProtocolInitializer, ProtocolSigner, ProtocolStakeDistribution},
+    entities::{ProtocolParameters, SignerWithStake, StakeDistribution},
+};
+use std::collections::HashMap;
+
+/// A fixture of Mithril data types.
+#[derive(Debug, Clone)]
+pub struct MithrilFixture {
+    protocol_parameters: ProtocolParameters,
+    signers: Vec<SignerFixture>,
+    stake_distribution: ProtocolStakeDistribution,
+}
+
+/// A signer fixture, containing a [signer entity][SignerWithStake] with its
+/// corresponding protocol [signer][ProtocolSigner] and
+/// [initializer][ProtocolInitializer]
+#[derive(Debug, Clone)]
+pub struct SignerFixture {
+    /// A [SignerWithStake].
+    pub signer_with_stake: SignerWithStake,
+    /// A [ProtocolSigner].
+    pub protocol_signer: ProtocolSigner,
+    /// A [ProtocolSigner].
+    pub protocol_initializer: ProtocolInitializer,
+}
+
+impl MithrilFixture {
+    /// [MithrilFixture] factory.
+    #[cfg(any(test, feature = "test_only"))]
+    pub fn new(
+        protocol_parameters: ProtocolParameters,
+        signers: Vec<SignerFixture>,
+        stake_distribution: ProtocolStakeDistribution,
+    ) -> Self {
+        Self {
+            protocol_parameters,
+            signers,
+            stake_distribution,
+        }
+    }
+
+    /// Get the fixture protocol parameters.
+    pub fn protocol_parameters(&self) -> ProtocolParameters {
+        self.protocol_parameters.clone()
+    }
+
+    /// Get the fixture signers.
+    pub fn signers_fixture(&self) -> Vec<SignerFixture> {
+        self.signers.clone()
+    }
+
+    /// Get the fixture signers with stake.
+    pub fn signers_with_stake(&self) -> Vec<SignerWithStake> {
+        self.signers
+            .iter()
+            .map(|s| &s.signer_with_stake)
+            .cloned()
+            .collect()
+    }
+
+    /// Get the fixture stake distribution.
+    pub fn stake_distribution(&self) -> StakeDistribution {
+        HashMap::from_iter(self.stake_distribution.clone())
+    }
+
+    /// Get the fixture protocol stake distribution.
+    pub fn protocol_stake_distribution(&self) -> ProtocolStakeDistribution {
+        self.stake_distribution.clone()
+    }
+}

mithril-common/src/test_utils/mod.rs

@@ -0,0 +1,7 @@
+#[cfg(any(test, feature = "test_only"))]
+mod fixture_builder;
+mod mithril_fixture;
+
+#[cfg(any(test, feature = "test_only"))]
+pub use fixture_builder::{MithrilFixtureBuilder, StakeDistributionGenerationMethod};
+pub use mithril_fixture::{MithrilFixture, SignerFixture};