Use MithrilFixtureBuilder with aggregator tests

Author: Alenar

mithril-aggregator/src/dependency.rs

@@ -188,7 +188,7 @@ impl DependencyManager {
     /// * Fill the [VerificationKeyStore] with the given signers keys.
     /// * Fill the [StakeStore] with the given signers stakes.
     /// * Fill the [ProtocolParametersStore] with the given parameters.
-    pub async fn simulate_genesis(
+    pub async fn prepare_for_genesis(
         &self,
         genesis_signers: Vec<SignerWithStake>,
         second_epoch_signers: Vec<SignerWithStake>,

mithril-aggregator/src/multi_signer.rs

@@ -729,6 +729,7 @@ mod tests {
         crypto_helper::tests_setup::*,
         fake_data,
         store::{adapter::MemoryAdapter, StakeStore},
+        test_utils::MithrilFixtureBuilder,
     };
     use std::{collections::HashMap, sync::Arc};
 
@@ -861,17 +862,9 @@ mod tests {
     #[tokio::test]
     async fn test_multi_signer_stake_distribution_ok() {
         let mut multi_signer = setup_multi_signer().await;
+        let fixture = MithrilFixtureBuilder::default().with_signers(5).build();
+        let mut stake_distribution_expected = fixture.protocol_stake_distribution();
 
-        let mut stake_distribution_expected: ProtocolStakeDistribution =
-            setup_signers(5, &setup_protocol_parameters())
-                .iter()
-                .map(|(signer_with_stake, _, _)| {
-                    (
-                        signer_with_stake.party_id.to_owned(),
-                        signer_with_stake.stake,
-                    )
-                })
-                .collect::<_>();
         stake_distribution_expected.sort_by_key(|k| k.0.clone());
         multi_signer
             .update_stake_distribution(&stake_distribution_expected)
@@ -918,22 +911,15 @@ mod tests {
             .await
             .expect("update protocol parameters failed");
 
-        let signers = setup_signers(5, &protocol_parameters);
+        let fixture = MithrilFixtureBuilder::default().with_signers(5).build();
+
+        let stake_distribution = fixture.protocol_stake_distribution();
 
-        let stake_distribution = &signers
-            .iter()
-            .map(|(signer_with_stake, _, _)| {
-                (
-                    signer_with_stake.party_id.to_owned(),
-                    signer_with_stake.stake,
-                )
-            })
-            .collect::<_>();
         multi_signer
-            .update_stake_distribution(stake_distribution)
+            .update_stake_distribution(&stake_distribution)
             .await
             .expect("update stake distribution failed");
-        for (signer_with_stake, _, _) in &signers {
+        for signer_with_stake in &fixture.signers_with_stake() {
             verification_key_store
                 .save_verification_key(
                     start_epoch.offset_to_recording_epoch(),
@@ -959,17 +945,20 @@ mod tests {
         let mut signatures = Vec::new();
 
         let mut expected_certificate_signers: Vec<SignerWithStake> = Vec::new();
-        for (signer_with_stake, protocol_signer, _protocol_initializer) in &signers {
-            if let Some(signature) = protocol_signer.sign(message.compute_hash().as_bytes()) {
+        for signer_fixture in fixture.signers_fixture() {
+            if let Some(signature) = signer_fixture
+                .protocol_signer
+                .sign(message.compute_hash().as_bytes())
+            {
                 let won_indexes = signature.indexes.clone();
 
                 signatures.push(entities::SingleSignatures::new(
-                    signer_with_stake.party_id.to_owned(),
+                    signer_fixture.signer_with_stake.party_id.to_owned(),
                     key_encode_hex(signature).unwrap(),
                     won_indexes,
                 ));
 
-                expected_certificate_signers.push(signer_with_stake.to_owned())
+                expected_certificate_signers.push(signer_fixture.signer_with_stake.to_owned())
             }
         }
 

mithril-aggregator/src/runtime/runner.rs

@@ -615,15 +615,13 @@ pub mod tests {
     use crate::{MithrilSignerRegisterer, ProtocolParametersStorer, SignerRegistrationRound};
     use mithril_common::chain_observer::FakeObserver;
     use mithril_common::crypto_helper::{
-        key_decode_hex,
-        tests_setup::{setup_certificate_chain, setup_signers},
-        ProtocolMultiSignature,
+        key_decode_hex, tests_setup::setup_certificate_chain, ProtocolMultiSignature,
     };
     use mithril_common::digesters::DumbImmutableFileObserver;
     use mithril_common::entities::{
-        Beacon, CertificatePending, Epoch, HexEncodedKey, ProtocolMessage, SignerWithStake,
-        StakeDistribution,
+        Beacon, CertificatePending, Epoch, HexEncodedKey, ProtocolMessage, StakeDistribution,
     };
+    use mithril_common::test_utils::MithrilFixtureBuilder;
     use mithril_common::{entities::ProtocolMessagePartKey, fake_data, store::StakeStorer};
     use mithril_common::{BeaconProviderImpl, CardanoNetwork};
     use std::collections::HashMap;
@@ -696,10 +694,18 @@ pub mod tests {
 
     #[tokio::test]
     async fn test_update_stake_distribution() {
-        let (deps, config) = initialize_dependencies().await;
+        let (mut deps, config) = initialize_dependencies().await;
+        let chain_observer = Arc::new(FakeObserver::default());
+        deps.chain_observer = chain_observer.clone();
         let deps = Arc::new(deps);
         let runner = AggregatorRunner::new(config, deps.clone());
         let beacon = runner.get_beacon_from_chain().await.unwrap();
+        let fixture = MithrilFixtureBuilder::default().with_signers(5).build();
+        let expected = fixture.stake_distribution();
+
+        chain_observer
+            .set_signers(fixture.signers_with_stake())
+            .await;
         runner
             .update_beacon(&beacon)
             .await
@@ -709,13 +715,6 @@ pub mod tests {
             .await
             .expect("updating stake distribution should not return an error");
 
-        let current_stake_distribution = deps
-            .chain_observer
-            .get_current_stake_distribution()
-            .await
-            .unwrap()
-            .expect("The stake distribution should not be None.");
-
         let saved_stake_distribution = deps
             .stake_store
             .get_stakes(beacon.epoch.offset_to_recording_epoch())
@@ -728,13 +727,7 @@ pub mod tests {
                 )
             });
 
-        assert_eq!(
-            current_stake_distribution.len(),
-            saved_stake_distribution.len()
-        );
-        for (party_id, stake) in current_stake_distribution.iter() {
-            assert_eq!(stake, saved_stake_distribution.get(party_id).unwrap());
-        }
+        assert_eq!(expected, saved_stake_distribution);
     }
 
     #[tokio::test]
@@ -815,11 +808,11 @@ pub mod tests {
         let beacon = runner.get_beacon_from_chain().await.unwrap();
         runner.update_beacon(&beacon).await.unwrap();
 
-        let signers = fake_data::signers_with_stakes(5);
-        let current_signers = signers[1..3].to_vec();
-        let next_signers = signers[2..5].to_vec();
+        let fixture = MithrilFixtureBuilder::default().with_signers(5).build();
+        let current_signers = fixture.signers_with_stake()[1..3].to_vec();
+        let next_signers = fixture.signers_with_stake()[2..5].to_vec();
         let protocol_parameters = fake_data::protocol_parameters();
-        deps.simulate_genesis(
+        deps.prepare_for_genesis(
             current_signers.clone(),
             next_signers.clone(),
             &protocol_parameters.clone(),
@@ -853,17 +846,12 @@ pub mod tests {
         let beacon = runner.get_beacon_from_chain().await.unwrap();
         runner.update_beacon(&beacon).await.unwrap();
 
-        let protocol_parameters = fake_data::protocol_parameters();
-        let signers = setup_signers(5, &protocol_parameters.clone().into())
-            .into_iter()
-            .map(|s| s.0)
-            .collect::<Vec<SignerWithStake>>();
-        let current_signers = signers[1..3].to_vec();
-        let next_signers = signers[2..5].to_vec();
-        deps.simulate_genesis(
-            current_signers.clone(),
-            next_signers.clone(),
-            &protocol_parameters.clone(),
+        let fixture = MithrilFixtureBuilder::default().with_signers(5).build();
+
+        deps.prepare_for_genesis(
+            fixture.signers_with_stake(),
+            fixture.signers_with_stake(),
+            &fixture.protocol_parameters(),
         )
         .await;
 
@@ -902,7 +890,7 @@ pub mod tests {
 
         let mut expected = WorkingCertificate::from_pending_certificate(
             &certificate_pending,
-            &current_signers,
+            &fixture.signers_with_stake(),
             &message,
             &aggregate_verification_key,
             &sut.initiated_at,
@@ -956,18 +944,15 @@ pub mod tests {
         let (deps, config) = initialize_dependencies().await;
         let deps = Arc::new(deps);
         let runner = AggregatorRunner::new(config, deps.clone());
-        let beacon = runner.get_beacon_from_chain().await.unwrap();
         let digest = "1+2+3+4=10".to_string();
-        runner.update_beacon(&beacon).await.unwrap();
-        for epoch in [
-            beacon.epoch.offset_to_signer_retrieval_epoch().unwrap(),
-            beacon.epoch.offset_to_next_signer_retrieval_epoch(),
-        ] {
-            deps.protocol_parameters_store
-                .save_protocol_parameters(epoch, fake_data::protocol_parameters())
-                .await
-                .unwrap();
-        }
+        runner.update_beacon(&fake_data::beacon()).await.unwrap();
+        let fixture = MithrilFixtureBuilder::default().build();
+        deps.prepare_for_genesis(
+            fixture.signers_with_stake(),
+            fixture.signers_with_stake(),
+            &fixture.protocol_parameters(),
+        )
+        .await;
 
         runner
             .update_message_in_multisigner(digest)
@@ -1061,15 +1046,8 @@ pub mod tests {
             .expect_get_multi_signature()
             .return_once(move || Ok(Some(multi_signature)));
         deps.multi_signer = Arc::new(RwLock::new(mock_multi_signer));
-        let certificate_store = deps.certificate_store.clone();
+        deps.init_state_from_chain(&certificate_chain, vec![]).await;
         let runner = AggregatorRunner::new(config, Arc::new(deps));
-        for certificate in certificate_chain[1..].iter().rev() {
-            certificate_store
-                .as_ref()
-                .save(certificate.to_owned())
-                .await
-                .expect("save certificate to store should not fail");
-        }
 
         let certificate = runner
             .create_and_save_certificate(&working_certificate)

mithril-aggregator/src/signer_registerer.rs

@@ -201,10 +201,9 @@ mod tests {
 
     use mithril_common::{
         chain_observer::FakeObserver,
-        crypto_helper::tests_setup::setup_signers,
         entities::{Epoch, PartyId, Signer},
-        fake_data,
         store::adapter::MemoryAdapter,
+        test_utils::MithrilFixtureBuilder,
     };
 
     use crate::{
@@ -222,13 +221,9 @@ mod tests {
         let signer_registerer =
             MithrilSignerRegisterer::new(Arc::new(chain_observer), verification_key_store.clone());
         let registration_epoch = Epoch(1);
-        let signers = setup_signers(5, &fake_data::protocol_parameters().into());
-        let signer_to_register: Signer = signers[0].0.clone().into();
-        assert!(signer_to_register.operational_certificate.is_some());
-        let stake_distribution = signers
-            .into_iter()
-            .map(|s| (s.0.party_id, s.0.stake))
-            .collect::<HashMap<_, _>>();
+        let fixture = MithrilFixtureBuilder::default().with_signers(5).build();
+        let signer_to_register: Signer = fixture.signers()[0].to_owned();
+        let stake_distribution = fixture.stake_distribution();
 
         signer_registerer
             .open_registration_round(registration_epoch, stake_distribution)
@@ -264,13 +259,12 @@ mod tests {
         let signer_registerer =
             MithrilSignerRegisterer::new(Arc::new(chain_observer), verification_key_store.clone());
         let registration_epoch = Epoch(1);
-        let signers = setup_signers(5, &fake_data::protocol_parameters().into());
-        let signer_to_register: Signer = signers[1].0.clone().into();
-        assert!(signer_to_register.operational_certificate.is_none());
-        let stake_distribution = signers
-            .into_iter()
-            .map(|s| (s.0.party_id, s.0.stake))
-            .collect::<HashMap<_, _>>();
+        let fixture = MithrilFixtureBuilder::default()
+            .with_signers(5)
+            .disable_signers_certification()
+            .build();
+        let signer_to_register: Signer = fixture.signers()[0].to_owned();
+        let stake_distribution = fixture.stake_distribution();
 
         signer_registerer
             .open_registration_round(registration_epoch, stake_distribution)
@@ -305,8 +299,8 @@ mod tests {
         ));
         let signer_registerer =
             MithrilSignerRegisterer::new(Arc::new(chain_observer), verification_key_store.clone());
-        let signers = setup_signers(5, &fake_data::protocol_parameters().into());
-        let signer_to_register: Signer = signers[1].0.clone().into();
+        let fixture = MithrilFixtureBuilder::default().with_signers(5).build();
+        let signer_to_register: Signer = fixture.signers()[0].to_owned();
 
         signer_registerer
             .register_signer(&signer_to_register)

mithril-aggregator/src/tools/genesis.rs

@@ -159,9 +159,10 @@ mod tests {
 
     use mithril_common::{
         certificate_chain::MithrilCertificateVerifier,
-        crypto_helper::{tests_setup, ProtocolClerk, ProtocolGenesisSigner},
+        crypto_helper::{ProtocolClerk, ProtocolGenesisSigner},
         fake_data,
         store::adapter::MemoryAdapter,
+        test_utils::MithrilFixtureBuilder,
     };
 
     use super::*;
@@ -179,10 +180,9 @@ mod tests {
     }
 
     fn create_fake_genesis_avk() -> ProtocolAggregateVerificationKey {
-        let protocol_parameters = tests_setup::setup_protocol_parameters();
-        let signers = tests_setup::setup_signers(5, &protocol_parameters);
-        let first_signer = &signers.first().unwrap().1;
-        let clerk = ProtocolClerk::from_signer(first_signer);
+        let fixture = MithrilFixtureBuilder::default().with_signers(5).build();
+        let first_signer = fixture.signers_fixture()[0].clone().protocol_signer;
+        let clerk = ProtocolClerk::from_signer(&first_signer);
         clerk.compute_avk()
     }
 

mithril-aggregator/tests/certificate_chain.rs

@@ -3,9 +3,12 @@ mod test_extensions;
 use std::collections::BTreeMap;
 
 use mithril_aggregator::VerificationKeyStorer;
-use mithril_common::chain_observer::ChainObserver;
-use mithril_common::crypto_helper::{tests_setup, ProtocolPartyId, ProtocolStake};
-use mithril_common::entities::{ProtocolParameters, SignerWithStake};
+use mithril_common::{
+    chain_observer::ChainObserver,
+    crypto_helper::{ProtocolPartyId, ProtocolStake},
+    entities::ProtocolParameters,
+    test_utils::MithrilFixtureBuilder,
+};
 use test_extensions::RuntimeTester;
 
 #[tokio::test]
@@ -18,18 +21,19 @@ async fn certificate_chain() {
     let mut tester = RuntimeTester::build(protocol_parameters.clone()).await;
 
     comment!("Create signers & declare stake distribution");
-    let signers = tests_setup::setup_signers(5, &protocol_parameters.clone().into());
-    let mut signers_with_stake: Vec<SignerWithStake> = signers
-        .iter()
-        .map(|(signer_with_stake, _, _)| signer_with_stake.to_owned())
-        .collect();
+    let fixture = MithrilFixtureBuilder::default()
+        .with_signers(5)
+        .with_protocol_parameters(protocol_parameters.clone())
+        .build();
+    let signers = fixture.signers_fixture();
+    let mut signers_with_stake = fixture.signers_with_stake();
     tester
         .chain_observer
         .set_signers(signers_with_stake.clone())
         .await;
     tester
         .deps
-        .simulate_genesis(
+        .prepare_for_genesis(
             signers_with_stake.clone(),
             signers_with_stake.clone(),
             &protocol_parameters,