client/beefy: more flexible test runtime api

client/beefy/src/tests.rs

@@ -77,8 +77,8 @@ const BAD_MMR_ROOT: MmrRootHash = MmrRootHash::repeat_byte(0x42);
 type BeefyBlockImport = crate::BeefyBlockImport<
 	Block,
 	substrate_test_runtime_client::Backend,
-	two_validators::TestApi,
-	BlockImportAdapter<PeersClient, sp_api::TransactionFor<two_validators::TestApi, Block>>,
+	TestApi,
+	BlockImportAdapter<PeersClient, sp_api::TransactionFor<TestApi, Block>>,
 >;
 
 pub(crate) type BeefyValidatorSet = ValidatorSet<AuthorityId>;
@@ -198,12 +198,12 @@ impl TestNetFactory for BeefyTestNet {
 		Option<BoxJustificationImport<Block>>,
 		Self::PeerData,
 	) {
+		let keys = &[BeefyKeyring::Alice, BeefyKeyring::Bob];
+		let validator_set = ValidatorSet::new(make_beefy_ids(keys), 0).unwrap();
+		let api = Arc::new(TestApi::with_validator_set(&validator_set));
 		let inner = BlockImportAdapter::new(client.clone());
-		let (block_import, voter_links, rpc_links) = beefy_block_import_and_links(
-			inner,
-			client.as_backend(),
-			Arc::new(two_validators::TestApi {}),
-		);
+		let (block_import, voter_links, rpc_links) =
+			beefy_block_import_and_links(inner, client.as_backend(), api);
 		let peer_data = PeerData {
 			beefy_rpc_links: Mutex::new(Some(rpc_links)),
 			beefy_voter_links: Mutex::new(Some(voter_links)),
@@ -230,83 +230,79 @@ impl TestNetFactory for BeefyTestNet {
 	}
 }
 
-macro_rules! create_test_api {
-    ( $api_name:ident, mmr_root: $mmr_root:expr, $($inits:expr),+ ) => {
-		pub(crate) mod $api_name {
-			use super::*;
-
-			#[derive(Clone, Default)]
-			pub(crate) struct TestApi {}
+#[derive(Clone)]
+pub(crate) struct TestApi {
+	pub beefy_genesis: u64,
+	pub validator_set: BeefyValidatorSet,
+	pub mmr_root_hash: MmrRootHash,
+}
 
-			// compiler gets confused and warns us about unused inner
-			#[allow(dead_code)]
-			pub(crate) struct RuntimeApi {
-				inner: TestApi,
-			}
+impl TestApi {
+	pub fn new(
+		beefy_genesis: u64,
+		validator_set: &BeefyValidatorSet,
+		mmr_root_hash: MmrRootHash,
+	) -> Self {
+		TestApi { beefy_genesis, validator_set: validator_set.clone(), mmr_root_hash }
+	}
 
-			impl ProvideRuntimeApi<Block> for TestApi {
-				type Api = RuntimeApi;
-				fn runtime_api<'a>(&'a self) -> ApiRef<'a, Self::Api> {
-					RuntimeApi { inner: self.clone() }.into()
-				}
-			}
-			sp_api::mock_impl_runtime_apis! {
-				impl BeefyApi<Block> for RuntimeApi {
-					fn beefy_genesis() -> Option<NumberFor<Block>> {
-						Some(sp_runtime::traits::One::one())
-					}
-
-					fn validator_set() -> Option<BeefyValidatorSet> {
-						BeefyValidatorSet::new(make_beefy_ids(&[$($inits),+]), 0)
-					}
-				}
-
-				impl MmrApi<Block, MmrRootHash, NumberFor<Block>> for RuntimeApi {
-					fn mmr_root() -> Result<MmrRootHash, MmrError> {
-						Ok($mmr_root)
-					}
-
-					fn generate_proof(
-						_block_numbers: Vec<u64>,
-						_best_known_block_number: Option<u64>
-					) -> Result<(Vec<EncodableOpaqueLeaf>, Proof<MmrRootHash>), MmrError> {
-						unimplemented!()
-					}
-
-					fn verify_proof(_leaves: Vec<EncodableOpaqueLeaf>, _proof: Proof<MmrRootHash>) -> Result<(), MmrError> {
-						unimplemented!()
-					}
-
-					fn verify_proof_stateless(
-						_root: MmrRootHash,
-						_leaves: Vec<EncodableOpaqueLeaf>,
-						_proof: Proof<MmrRootHash>
-					) -> Result<(), MmrError> {
-						unimplemented!()
-					}
-				}
-			}
+	pub fn with_validator_set(validator_set: &BeefyValidatorSet) -> Self {
+		TestApi {
+			beefy_genesis: 1,
+			validator_set: validator_set.clone(),
+			mmr_root_hash: GOOD_MMR_ROOT,
 		}
-	};
+	}
+}
+
+// compiler gets confused and warns us about unused inner
+#[allow(dead_code)]
+pub(crate) struct RuntimeApi {
+	inner: TestApi,
 }
 
-create_test_api!(two_validators, mmr_root: GOOD_MMR_ROOT, BeefyKeyring::Alice, BeefyKeyring::Bob);
-create_test_api!(
-	four_validators,
-	mmr_root: GOOD_MMR_ROOT,
-	BeefyKeyring::Alice,
-	BeefyKeyring::Bob,
-	BeefyKeyring::Charlie,
-	BeefyKeyring::Dave
-);
-create_test_api!(
-	bad_four_validators,
-	mmr_root: BAD_MMR_ROOT,
-	BeefyKeyring::Alice,
-	BeefyKeyring::Bob,
-	BeefyKeyring::Charlie,
-	BeefyKeyring::Dave
-);
+impl ProvideRuntimeApi<Block> for TestApi {
+	type Api = RuntimeApi;
+	fn runtime_api<'a>(&'a self) -> ApiRef<'a, Self::Api> {
+		RuntimeApi { inner: self.clone() }.into()
+	}
+}
+sp_api::mock_impl_runtime_apis! {
+	impl BeefyApi<Block> for RuntimeApi {
+		fn beefy_genesis() -> Option<NumberFor<Block>> {
+			Some(self.inner.beefy_genesis)
+		}
+
+		fn validator_set() -> Option<BeefyValidatorSet> {
+			Some(self.inner.validator_set.clone())
+		}
+	}
+
+	impl MmrApi<Block, MmrRootHash, NumberFor<Block>> for RuntimeApi {
+		fn mmr_root() -> Result<MmrRootHash, MmrError> {
+			Ok(self.inner.mmr_root_hash)
+		}
+
+		fn generate_proof(
+			_block_numbers: Vec<u64>,
+			_best_known_block_number: Option<u64>
+		) -> Result<(Vec<EncodableOpaqueLeaf>, Proof<MmrRootHash>), MmrError> {
+			unimplemented!()
+		}
+
+		fn verify_proof(_leaves: Vec<EncodableOpaqueLeaf>, _proof: Proof<MmrRootHash>) -> Result<(), MmrError> {
+			unimplemented!()
+		}
+
+		fn verify_proof_stateless(
+			_root: MmrRootHash,
+			_leaves: Vec<EncodableOpaqueLeaf>,
+			_proof: Proof<MmrRootHash>
+		) -> Result<(), MmrError> {
+			unimplemented!()
+		}
+	}
+}
 
 fn add_mmr_digest(header: &mut Header, mmr_hash: MmrRootHash) {
 	header.digest_mut().push(DigestItem::Consensus(
@@ -336,9 +332,9 @@ pub(crate) fn create_beefy_keystore(authority: BeefyKeyring) -> SyncCryptoStoreP
 fn voter_init_setup(
 	net: &mut BeefyTestNet,
 	finality: &mut futures::stream::Fuse<FinalityNotifications<Block>>,
+	api: &TestApi,
 ) -> sp_blockchain::Result<PersistedState<Block>> {
 	let backend = net.peer(0).client().as_backend();
-	let api = Arc::new(crate::tests::two_validators::TestApi {});
 	let known_peers = Arc::new(Mutex::new(KnownPeers::new()));
 	let gossip_validator =
 		Arc::new(crate::communication::gossip::GossipValidator::new(known_peers));
@@ -349,9 +345,9 @@ fn voter_init_setup(
 		None,
 	);
 	let best_grandpa =
-		futures::executor::block_on(wait_for_runtime_pallet(&*api, &mut gossip_engine, finality))
+		futures::executor::block_on(wait_for_runtime_pallet(api, &mut gossip_engine, finality))
 			.unwrap();
-	load_or_init_voter_state(&*backend, &*api, best_grandpa, 1)
+	load_or_init_voter_state(&*backend, api, best_grandpa, 1)
 }
 
 // Spawns beefy voters. Returns a future to spawn on the runtime.
@@ -361,7 +357,7 @@ fn initialize_beefy<API>(
 	min_block_delta: u32,
 ) -> impl Future<Output = ()>
 where
-	API: ProvideRuntimeApi<Block> + Default + Sync + Send,
+	API: ProvideRuntimeApi<Block> + Sync + Send,
 	API::Api: BeefyApi<Block> + MmrApi<Block, MmrRootHash, NumberFor<Block>>,
 {
 	let tasks = FuturesUnordered::new();
@@ -548,7 +544,7 @@ async fn beefy_finalizing_blocks() {
 
 	let mut net = BeefyTestNet::new(2);
 
-	let api = Arc::new(two_validators::TestApi {});
+	let api = Arc::new(TestApi::with_validator_set(&validator_set));
 	let beefy_peers = peers.iter().enumerate().map(|(id, key)| (id, key, api.clone())).collect();
 	tokio::spawn(initialize_beefy(&mut net, beefy_peers, min_block_delta));
 
@@ -586,7 +582,7 @@ async fn lagging_validators() {
 	let min_block_delta = 1;
 
 	let mut net = BeefyTestNet::new(2);
-	let api = Arc::new(two_validators::TestApi {});
+	let api = Arc::new(TestApi::with_validator_set(&validator_set));
 	let beefy_peers = peers.iter().enumerate().map(|(id, key)| (id, key, api.clone())).collect();
 	tokio::spawn(initialize_beefy(&mut net, beefy_peers, min_block_delta));
 
@@ -664,7 +660,7 @@ async fn correct_beefy_payload() {
 	let mut net = BeefyTestNet::new(4);
 
 	// Alice, Bob, Charlie will vote on good payloads
-	let good_api = Arc::new(four_validators::TestApi {});
+	let good_api = Arc::new(TestApi::new(1, &validator_set, GOOD_MMR_ROOT));
 	let good_peers = [BeefyKeyring::Alice, BeefyKeyring::Bob, BeefyKeyring::Charlie]
 		.iter()
 		.enumerate()
@@ -673,7 +669,7 @@ async fn correct_beefy_payload() {
 	tokio::spawn(initialize_beefy(&mut net, good_peers, min_block_delta));
 
 	// Dave will vote on bad mmr roots
-	let bad_api = Arc::new(bad_four_validators::TestApi {});
+	let bad_api = Arc::new(TestApi::new(1, &validator_set, BAD_MMR_ROOT));
 	let bad_peers = vec![(3, &BeefyKeyring::Dave, bad_api)];
 	tokio::spawn(initialize_beefy(&mut net, bad_peers, min_block_delta));
 
@@ -718,6 +714,8 @@ async fn beefy_importing_blocks() {
 	sp_tracing::try_init_simple();
 
 	let mut net = BeefyTestNet::new(2);
+	let keys = &[BeefyKeyring::Alice, BeefyKeyring::Bob];
+	let good_set = ValidatorSet::new(make_beefy_ids(keys), 0).unwrap();
 
 	let client = net.peer(0).client().clone();
 	let (mut block_import, _, peer_data) = net.make_block_import(client.clone());
@@ -773,9 +771,7 @@ async fn beefy_importing_blocks() {
 	// Import with valid justification.
 	let parent_id = BlockId::Number(1);
 	let block_num = 2;
-	let keys = &[BeefyKeyring::Alice, BeefyKeyring::Bob];
-	let validator_set = ValidatorSet::new(make_beefy_ids(keys), 0).unwrap();
-	let proof = crate::justification::tests::new_finality_proof(block_num, &validator_set, keys);
+	let proof = crate::justification::tests::new_finality_proof(block_num, &good_set, keys);
 	let versioned_proof: VersionedFinalityProof<NumberFor<Block>, Signature> = proof.into();
 	let encoded = versioned_proof.encode();
 	let justif = Some(Justifications::from((BEEFY_ENGINE_ID, encoded)));
@@ -818,8 +814,8 @@ async fn beefy_importing_blocks() {
 	let parent_id = BlockId::Number(2);
 	let block_num = 3;
 	let keys = &[BeefyKeyring::Alice];
-	let validator_set = ValidatorSet::new(make_beefy_ids(keys), 1).unwrap();
-	let proof = crate::justification::tests::new_finality_proof(block_num, &validator_set, keys);
+	let bad_set = ValidatorSet::new(make_beefy_ids(keys), 1).unwrap();
+	let proof = crate::justification::tests::new_finality_proof(block_num, &bad_set, keys);
 	let versioned_proof: VersionedFinalityProof<NumberFor<Block>, Signature> = proof.into();
 	let encoded = versioned_proof.encode();
 	let justif = Some(Justifications::from((BEEFY_ENGINE_ID, encoded)));
@@ -869,7 +865,7 @@ async fn voter_initialization() {
 	let min_block_delta = 10;
 
 	let mut net = BeefyTestNet::new(2);
-	let api = Arc::new(two_validators::TestApi {});
+	let api = Arc::new(TestApi::with_validator_set(&validator_set));
 	let beefy_peers = peers.iter().enumerate().map(|(id, key)| (id, key, api.clone())).collect();
 	tokio::spawn(initialize_beefy(&mut net, beefy_peers, min_block_delta));
 
@@ -902,7 +898,7 @@ async fn on_demand_beefy_justification_sync() {
 	let mut net = BeefyTestNet::new(4);
 
 	// Alice, Bob, Charlie start first and make progress through voting.
-	let api = Arc::new(four_validators::TestApi {});
+	let api = Arc::new(TestApi::with_validator_set(&validator_set));
 	let fast_peers = [BeefyKeyring::Alice, BeefyKeyring::Bob, BeefyKeyring::Charlie];
 	let voting_peers =
 		fast_peers.iter().enumerate().map(|(id, key)| (id, key, api.clone())).collect();
@@ -980,8 +976,9 @@ async fn should_initialize_voter_at_genesis() {
 	// finalize 13 without justifications
 	net.peer(0).client().as_client().finalize_block(hashes[13], None).unwrap();
 
+	let api = TestApi::with_validator_set(&validator_set);
 	// load persistent state - nothing in DB, should init at session boundary
-	let persisted_state = voter_init_setup(&mut net, &mut finality).unwrap();
+	let persisted_state = voter_init_setup(&mut net, &mut finality, &api).unwrap();
 
 	// Test initialization at session boundary.
 	// verify voter initialized with two sessions starting at blocks 1 and 10
@@ -1041,8 +1038,9 @@ async fn should_initialize_voter_when_last_final_is_session_boundary() {
 	// Test corner-case where session boundary == last beefy finalized,
 	// expect rounds initialized at last beefy finalized 10.
 
+	let api = TestApi::with_validator_set(&validator_set);
 	// load persistent state - nothing in DB, should init at session boundary
-	let persisted_state = voter_init_setup(&mut net, &mut finality).unwrap();
+	let persisted_state = voter_init_setup(&mut net, &mut finality, &api).unwrap();
 
 	// verify voter initialized with single session starting at block 10
 	assert_eq!(persisted_state.voting_oracle().sessions().len(), 1);
@@ -1098,8 +1096,9 @@ async fn should_initialize_voter_at_latest_finalized() {
 
 	// Test initialization at last BEEFY finalized.
 
+	let api = TestApi::with_validator_set(&validator_set);
 	// load persistent state - nothing in DB, should init at last BEEFY finalized
-	let persisted_state = voter_init_setup(&mut net, &mut finality).unwrap();
+	let persisted_state = voter_init_setup(&mut net, &mut finality, &api).unwrap();
 
 	// verify voter initialized with single session starting at block 12
 	assert_eq!(persisted_state.voting_oracle().sessions().len(), 1);

client/beefy/src/worker.rs

@@ -974,8 +974,8 @@ pub(crate) mod tests {
 		communication::notification::{BeefyBestBlockStream, BeefyVersionedFinalityProofStream},
 		keystore::tests::Keyring,
 		tests::{
-			create_beefy_keystore, get_beefy_streams, make_beefy_ids, two_validators::TestApi,
-			BeefyPeer, BeefyTestNet,
+			create_beefy_keystore, get_beefy_streams, make_beefy_ids, BeefyPeer, BeefyTestNet,
+			TestApi,
 		},
 		BeefyRPCLinks, KnownPeers,
 	};
@@ -1048,7 +1048,7 @@ pub(crate) mod tests {
 		};
 
 		let backend = peer.client().as_backend();
-		let api = Arc::new(TestApi {});
+		let api = Arc::new(TestApi::with_validator_set(&genesis_validator_set));
 		let network = peer.network_service().clone();
 		let known_peers = Arc::new(Mutex::new(KnownPeers::new()));
 		let gossip_validator = Arc::new(GossipValidator::new(known_peers.clone()));