Skip to content

Commit

Permalink
Simplify fake data building in test
Browse files Browse the repository at this point in the history
  • Loading branch information
@sfauvel
sfauvel committed Nov 27, 2024
1 parent b974e02 commit 28035c4
Showing 1 changed file with 55 additions and 59 deletions.
114 changes: 55 additions & 59 deletions mithril-aggregator/src/database/repository/signer_registration_store.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -108,43 +108,45 @@ mod tests {

use mithril_common::entities::{Epoch, PartyId, Signer, SignerWithStake};
use mithril_common::test_utils::fake_keys;
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::collections::HashMap;
use std::sync::Arc;

use crate::VerificationKeyStorer;

use super::*;

/// Build simple fake signers with stakes.
/// It could be done by `fake_data::signers_with_stakes` which produce verification keys dynamically
/// but take longer.
fn build_fake_signers_with_stakes(nb: u64) -> Vec<SignerWithStake> {
let verification_keys = fake_keys::signer_verification_key();
let nb_keys = verification_keys.len() as u64;
(1..=nb)
.map(|party_idx| SignerWithStake {
party_id: format!("party_id:{party_idx}"),
verification_key: verification_keys[(party_idx % nb_keys) as usize]
.try_into()
.unwrap(),
verification_key_signature: None,
operational_certificate: None,
kes_period: None,
stake: 10,
})
.collect()
}

fn build_signers(
nb_epoch: u64,
signers_per_epoch: usize,
) -> Vec<(Epoch, HashMap<PartyId, SignerWithStake>)> {
let mut values = vec![];

for epoch in 1..=nb_epoch {
let mut signers: HashMap<PartyId, SignerWithStake> =
HashMap::with_capacity(signers_per_epoch);

for party_idx in 1..=signers_per_epoch {
let party_id = format!("party_id:e{epoch}:{party_idx}");
signers.insert(
party_id.clone(),
SignerWithStake {
party_id: party_id.clone(),
verification_key: fake_keys::signer_verification_key()[0]
.try_into()
.unwrap(),
verification_key_signature: None,
operational_certificate: None,
kes_period: None,
stake: 10,
},
);
}
values.push((Epoch(epoch), signers));
}

values
) -> HashMap<Epoch, Vec<SignerWithStake>> {
(1..=nb_epoch)
.map(|epoch| {
(
Epoch(epoch),
build_fake_signers_with_stakes(signers_per_epoch as u64),
)
})
.collect()
}

fn insert_golden_signer_registration(connection: &SqliteConnection) {
Expand Down Expand Up @@ -181,13 +183,11 @@ mod tests {
}

pub fn init_signer_registration_store(
initial_data: Vec<(Epoch, HashMap<PartyId, SignerWithStake>)>,
initial_data: HashMap<Epoch, Vec<SignerWithStake>>,
) -> Arc<SignerRegistrationStore> {
let connection = main_db_connection().unwrap();
let initial_data: Vec<(Epoch, Vec<SignerWithStake>)> = initial_data
.into_iter()
.map(|(e, signers)| (e, signers.into_values().collect::<Vec<_>>()))
.collect();

let initial_data = initial_data.into_iter().collect();
insert_signer_registrations(&connection, initial_data).unwrap();

Arc::new(SignerRegistrationStore::new(Arc::new(connection)))
Expand Down Expand Up @@ -218,12 +218,15 @@ mod tests {
#[tokio::test]
pub async fn update_signer_in_store() {
let signers = build_signers(1, 1);
let signers_on_epoch = signers.get(&Epoch(1)).unwrap().clone();
let first_signer = signers_on_epoch.first().unwrap();

let store = init_signer_registration_store(signers);
let res = store
.save_verification_key(
Epoch(1),
SignerWithStake {
party_id: "party_id:e1:1".to_string(),
party_id: first_signer.party_id.clone(),
verification_key: fake_keys::signer_verification_key()[2].try_into().unwrap(),
verification_key_signature: None,
operational_certificate: None,
Expand All @@ -236,7 +239,7 @@ mod tests {

assert_eq!(
Some(SignerWithStake {
party_id: "party_id:e1:1".to_string(),
party_id: first_signer.party_id.clone(),
verification_key: fake_keys::signer_verification_key()[2].try_into().unwrap(),
verification_key_signature: None,
operational_certificate: None,
Expand Down Expand Up @@ -270,18 +273,15 @@ mod tests {
let signers = build_signers(2, 2);
let store = init_signer_registration_store(signers.clone());

let expected_signers: Option<BTreeMap<PartyId, Signer>> = signers
.into_iter()
.filter(|(e, _)| e == 1)
.map(|(_, signers)| {
BTreeMap::from_iter(signers.into_iter().map(|(p, s)| (p, s.into())))
})
.next();
let res = store
.get_verification_keys(Epoch(1))
.await
let epoch = Epoch(1);
let expected_signers = signers
.get(&epoch)
.unwrap()
.map(|x| BTreeMap::from_iter(x.into_iter()));
.iter()
.map(|s| (s.party_id.clone(), Signer::from(s.clone())))
.collect::<HashMap<PartyId, Signer>>();

let res = store.get_verification_keys(epoch).await.unwrap().unwrap();

assert_eq!(expected_signers, res);
}
Expand All @@ -291,16 +291,12 @@ mod tests {
let signers = build_signers(2, 2);
let store = init_signer_registration_store(signers.clone());

let expected_signers: Option<BTreeSet<SignerWithStake>> = signers
.into_iter()
.filter(|(e, _)| e == 1)
.map(|(_, signers)| BTreeSet::from_iter(signers.into_values()))
.next();
let res = store
.get_signers(Epoch(1))
.await
.unwrap()
.map(|x| BTreeSet::from_iter(x.into_iter()));
let epoch = Epoch(1);
let mut expected_signers = signers.get(&epoch).unwrap().clone();
expected_signers.sort_by(|a, b| a.party_id.cmp(&b.party_id));

let mut res = store.get_signers(epoch).await.unwrap().unwrap();
res.sort_by(|a, b| a.party_id.cmp(&b.party_id));

assert_eq!(expected_signers, res);
}
Expand Down Expand Up @@ -334,8 +330,8 @@ mod tests {
until_epoch: Epoch,
) -> Vec<Epoch> {
let mut epochs_in_database = vec![];
let mut current_epoch = Epoch(1);
while current_epoch <= until_epoch {
for epoch_number in 1..=(*until_epoch) {
let current_epoch = Epoch(epoch_number);
if store
.get_verification_keys(current_epoch)
.await
Expand All @@ -344,8 +340,8 @@ mod tests {
{
epochs_in_database.push(current_epoch);
}
current_epoch += 1;
}

epochs_in_database
}

Expand Down

0 comments on commit 28035c4

Please sign in to comment.