lighthouse/beacon_node/beacon_chain/tests/utils/test_rig.rs

91 lines
2.8 KiB
Rust
Raw Normal View History

2019-01-25 02:52:21 +00:00
use super::TestValidator;
2019-01-25 20:20:58 +00:00
pub use beacon_chain::dump::{Error as DumpError, SlotDump};
2019-01-24 06:05:48 +00:00
use beacon_chain::BeaconChain;
#[cfg(test)]
use db::{
stores::{BeaconBlockStore, BeaconStateStore},
MemoryDB,
};
use slot_clock::TestingSlotClock;
2019-01-25 02:52:21 +00:00
use std::sync::Arc;
use types::{ChainSpec, Keypair, Validator};
2019-01-24 06:05:48 +00:00
2019-01-25 00:30:06 +00:00
pub struct TestRig {
2019-01-24 06:05:48 +00:00
db: Arc<MemoryDB>,
2019-01-25 00:30:06 +00:00
beacon_chain: Arc<BeaconChain<MemoryDB, TestingSlotClock>>,
2019-01-24 06:05:48 +00:00
block_store: Arc<BeaconBlockStore<MemoryDB>>,
state_store: Arc<BeaconStateStore<MemoryDB>>,
2019-01-25 00:30:06 +00:00
validators: Vec<TestValidator>,
2019-01-24 06:05:48 +00:00
}
2019-01-25 00:30:06 +00:00
impl TestRig {
pub fn new(mut spec: ChainSpec, validator_count: usize) -> Self {
2019-01-24 06:05:48 +00:00
let db = Arc::new(MemoryDB::open());
let block_store = Arc::new(BeaconBlockStore::new(db.clone()));
let state_store = Arc::new(BeaconStateStore::new(db.clone()));
let slot_clock = TestingSlotClock::new(0);
2019-01-25 00:30:06 +00:00
// Remove the validators present in the spec (if any).
spec.initial_validators = Vec::with_capacity(validator_count);
spec.initial_balances = Vec::with_capacity(validator_count);
// Insert `validator_count` new `Validator` records into the spec, retaining the keypairs
// for later user.
let mut keypairs = Vec::with_capacity(validator_count);
for _ in 0..validator_count {
let keypair = Keypair::random();
spec.initial_validators.push(Validator {
pubkey: keypair.pk.clone(),
..std::default::Default::default()
});
spec.initial_balances.push(32_000_000_000); // 32 ETH
keypairs.push(keypair);
}
// Create the Beacon Chain
let beacon_chain = Arc::new(
2019-01-24 06:05:48 +00:00
BeaconChain::genesis(state_store.clone(), block_store.clone(), slot_clock, spec)
2019-01-25 00:30:06 +00:00
.unwrap(),
);
2019-01-24 06:05:48 +00:00
2019-01-25 00:30:06 +00:00
// Spawn the test validator instances.
let mut validators = Vec::with_capacity(validator_count);
for keypair in keypairs {
validators.push(TestValidator::new(keypair.clone(), beacon_chain.clone()));
}
2019-01-24 06:05:48 +00:00
Self {
db,
beacon_chain,
block_store,
state_store,
2019-01-25 00:30:06 +00:00
validators,
2019-01-24 06:05:48 +00:00
}
}
2019-01-25 00:30:06 +00:00
pub fn produce_next_slot(&mut self) {
2019-01-24 06:05:48 +00:00
let slot = self
.beacon_chain
.present_slot()
.expect("Unable to determine slot.")
+ 1;
2019-01-25 00:30:06 +00:00
2019-01-24 06:05:48 +00:00
self.beacon_chain.slot_clock.set_slot(slot);
2019-01-25 00:30:06 +00:00
let proposer = self
2019-01-24 06:05:48 +00:00
.beacon_chain
.block_proposer(slot)
.expect("Unable to determine proposer.");
2019-01-25 00:30:06 +00:00
self.validators[proposer].set_slot(slot);
self.validators[proposer].produce_block().unwrap();
2019-01-24 06:05:48 +00:00
}
2019-01-25 20:20:58 +00:00
pub fn chain_dump(&self) -> Result<Vec<SlotDump>, DumpError> {
self.beacon_chain.chain_dump()
}
2019-01-24 06:05:48 +00:00
}