Expose MMR root through runtime API - use it in BEEFY client (#11183)

* beefy-gadget: allow custom runtime api provider

* beefy-gadget: use mock runtime api in tests

* pallet-mmr: expose mmr root from state through runtime API

* beefy-gadget: get mmr root from runtime state

* pallet-beefy-mmr: remove MmrRoot from header digests

* frame/mmr: move mmr primitives out of frame

* frame/mmr: completely move primitives out of frame

* address review comments

* beefy-mmr: bring back mmr root from header digest

* clippy fixes for rustc 1.60

* address review comments

substrate/client/beefy/src/lib.rs

@@ -27,9 +27,10 @@ use sp_api::ProvideRuntimeApi;
 use sp_blockchain::HeaderBackend;
 use sp_consensus::SyncOracle;
 use sp_keystore::SyncCryptoStorePtr;
+use sp_mmr_primitives::MmrApi;
 use sp_runtime::traits::Block;
 
-use beefy_primitives::BeefyApi;
+use beefy_primitives::{BeefyApi, MmrRootHash};
 
 use crate::notification::{BeefyBestBlockSender, BeefySignedCommitmentSender};
 
@@ -87,7 +88,7 @@ pub fn beefy_peers_set_config(
 /// of today, Rust does not allow a type alias to be used as a trait bound. Tracking
 /// issue is <https://github.com/rust-lang/rust/issues/41517>.
 pub trait Client<B, BE>:
-	BlockchainEvents<B> + HeaderBackend<B> + Finalizer<B, BE> + ProvideRuntimeApi<B> + Send + Sync
+	BlockchainEvents<B> + HeaderBackend<B> + Finalizer<B, BE> + Send + Sync
 where
 	B: Block,
 	BE: Backend<B>,
@@ -110,18 +111,21 @@ where
 }
 
 /// BEEFY gadget initialization parameters.
-pub struct BeefyParams<B, BE, C, N>
+pub struct BeefyParams<B, BE, C, N, R>
 where
 	B: Block,
 	BE: Backend<B>,
 	C: Client<B, BE>,
-	C::Api: BeefyApi<B>,
+	R: ProvideRuntimeApi<B>,
+	R::Api: BeefyApi<B> + MmrApi<B, MmrRootHash>,
 	N: GossipNetwork<B> + Clone + SyncOracle + Send + Sync + 'static,
 {
 	/// BEEFY client
 	pub client: Arc<C>,
 	/// Client Backend
 	pub backend: Arc<BE>,
+	/// Runtime Api Provider
+	pub runtime: Arc<R>,
 	/// Local key store
 	pub key_store: Option<SyncCryptoStorePtr>,
 	/// Gossip network
@@ -138,21 +142,22 @@ where
 	pub protocol_name: std::borrow::Cow<'static, str>,
 }
 
-#[cfg(not(test))]
 /// Start the BEEFY gadget.
 ///
 /// This is a thin shim around running and awaiting a BEEFY worker.
-pub async fn start_beefy_gadget<B, BE, C, N>(beefy_params: BeefyParams<B, BE, C, N>)
+pub async fn start_beefy_gadget<B, BE, C, N, R>(beefy_params: BeefyParams<B, BE, C, N, R>)
 where
 	B: Block,
 	BE: Backend<B>,
 	C: Client<B, BE>,
-	C::Api: BeefyApi<B>,
+	R: ProvideRuntimeApi<B>,
+	R::Api: BeefyApi<B> + MmrApi<B, MmrRootHash>,
 	N: GossipNetwork<B> + Clone + SyncOracle + Send + Sync + 'static,
 {
 	let BeefyParams {
 		client,
 		backend,
+		runtime,
 		key_store,
 		network,
 		signed_commitment_sender,
@@ -188,6 +193,7 @@ where
 	let worker_params = worker::WorkerParams {
 		client,
 		backend,
+		runtime,
 		key_store: key_store.into(),
 		signed_commitment_sender,
 		beefy_best_block_sender,
@@ -198,7 +204,7 @@ where
 		sync_oracle,
 	};
 
-	let worker = worker::BeefyWorker::<_, _, _, _>::new(worker_params);
+	let worker = worker::BeefyWorker::<_, _, _, _, _>::new(worker_params);
 
 	worker.run().await
 }

substrate/client/beefy/src/metrics.rs

@@ -18,9 +18,7 @@
 
 //! BEEFY Prometheus metrics definition
 
-#[cfg(not(test))]
-use prometheus::{register, PrometheusError, Registry};
-use prometheus::{Counter, Gauge, U64};
+use prometheus::{register, Counter, Gauge, PrometheusError, Registry, U64};
 
 /// BEEFY metrics exposed through Prometheus
 pub(crate) struct Metrics {
@@ -39,7 +37,6 @@ pub(crate) struct Metrics {
 }
 
 impl Metrics {
-	#[cfg(not(test))]
 	pub(crate) fn register(registry: &Registry) -> Result<Self, PrometheusError> {
 		Ok(Self {
 			beefy_validator_set_id: register(

substrate/client/beefy/src/tests.rs

@@ -28,18 +28,20 @@ use sc_chain_spec::{ChainSpec, GenericChainSpec};
 use sc_client_api::HeaderBackend;
 use sc_consensus::BoxJustificationImport;
 use sc_keystore::LocalKeystore;
-use sc_network::{config::ProtocolConfig, NetworkService};
-use sc_network_gossip::GossipEngine;
+use sc_network::config::ProtocolConfig;
 use sc_network_test::{
 	Block, BlockImportAdapter, FullPeerConfig, PassThroughVerifier, Peer, PeersClient,
-	PeersFullClient, TestNetFactory,
+	TestNetFactory,
 };
 use sc_utils::notification::NotificationReceiver;
 
 use beefy_primitives::{
-	crypto::AuthorityId, ConsensusLog, MmrRootHash, ValidatorSet, BEEFY_ENGINE_ID,
+	crypto::AuthorityId, BeefyApi, ConsensusLog, MmrRootHash, ValidatorSet, BEEFY_ENGINE_ID,
 	KEY_TYPE as BeefyKeyType,
 };
+use sp_mmr_primitives::{EncodableOpaqueLeaf, Error as MmrError, LeafIndex, MmrApi, Proof};
+
+use sp_api::{ApiRef, ProvideRuntimeApi};
 use sp_consensus::BlockOrigin;
 use sp_core::H256;
 use sp_keystore::{SyncCryptoStore, SyncCryptoStorePtr};
@@ -47,19 +49,16 @@ use sp_runtime::{
 	codec::Encode, generic::BlockId, traits::Header as HeaderT, BuildStorage, DigestItem, Storage,
 };
 
-use substrate_test_runtime_client::{runtime::Header, Backend, ClientExt};
+use substrate_test_runtime_client::{runtime::Header, ClientExt};
 
-use crate::{
-	beefy_protocol_name,
-	keystore::tests::Keyring as BeefyKeyring,
-	notification::*,
-	worker::{tests::TestModifiers, BeefyWorker},
-};
+use crate::{beefy_protocol_name, keystore::tests::Keyring as BeefyKeyring, notification::*};
 
-const BEEFY_PROTOCOL_NAME: &'static str = "/beefy/1";
+pub(crate) const BEEFY_PROTOCOL_NAME: &'static str = "/beefy/1";
+const GOOD_MMR_ROOT: MmrRootHash = MmrRootHash::repeat_byte(0xbf);
+const BAD_MMR_ROOT: MmrRootHash = MmrRootHash::repeat_byte(0x42);
 
-type BeefyValidatorSet = ValidatorSet<AuthorityId>;
-type BeefyPeer = Peer<PeerData, PeersClient>;
+pub(crate) type BeefyValidatorSet = ValidatorSet<AuthorityId>;
+pub(crate) type BeefyPeer = Peer<PeerData, PeersClient>;
 
 #[derive(Debug, Serialize, Deserialize)]
 struct Genesis(std::collections::BTreeMap<String, String>);
@@ -101,27 +100,13 @@ fn beefy_protocol_name() {
 #[allow(dead_code)]
 #[derive(Clone)]
 pub(crate) struct BeefyLinkHalf {
-	signed_commitment_stream: BeefySignedCommitmentStream<Block>,
-	beefy_best_block_stream: BeefyBestBlockStream<Block>,
+	pub signed_commitment_stream: BeefySignedCommitmentStream<Block>,
+	pub beefy_best_block_stream: BeefyBestBlockStream<Block>,
 }
 
 #[derive(Default)]
 pub(crate) struct PeerData {
 	pub(crate) beefy_link_half: Mutex<Option<BeefyLinkHalf>>,
-	pub(crate) test_modifiers: Option<TestModifiers>,
-}
-
-impl PeerData {
-	pub(crate) fn use_validator_set(&mut self, validator_set: &ValidatorSet<AuthorityId>) {
-		if let Some(tm) = self.test_modifiers.as_mut() {
-			tm.active_validators = validator_set.clone();
-		} else {
-			self.test_modifiers = Some(TestModifiers {
-				active_validators: validator_set.clone(),
-				corrupt_mmr_roots: false,
-			});
-		}
-	}
 }
 
 pub(crate) struct BeefyTestNet {
@@ -153,17 +138,19 @@ impl BeefyTestNet {
 		count: usize,
 		session_length: u64,
 		validator_set: &BeefyValidatorSet,
+		include_mmr_digest: bool,
 	) {
 		self.peer(0).generate_blocks(count, BlockOrigin::File, |builder| {
 			let mut block = builder.build().unwrap().block;
 
-			let block_num = *block.header.number();
-			let num_byte = block_num.to_le_bytes().into_iter().next().unwrap();
-			let mmr_root = MmrRootHash::repeat_byte(num_byte);
+			if include_mmr_digest {
+				let block_num = *block.header.number();
+				let num_byte = block_num.to_le_bytes().into_iter().next().unwrap();
+				let mmr_root = MmrRootHash::repeat_byte(num_byte);
+				add_mmr_digest(&mut block.header, mmr_root);
+			}
 
-			add_mmr_digest(&mut block.header, mmr_root);
-
-			if block_num % session_length == 0 {
+			if *block.header.number() % session_length == 0 {
 				add_auth_change_digest(&mut block.header, validator_set.clone());
 			}
 
@@ -223,6 +210,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 {}
+
+			// compiler gets confused and warns us about unused inner
+			#[allow(dead_code)]
+			pub(crate) struct RuntimeApi {
+				inner: TestApi,
+			}
+
+			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 validator_set() -> Option<BeefyValidatorSet> {
+						BeefyValidatorSet::new(make_beefy_ids(&[$($inits),+]), 0)
+					}
+				}
+
+				impl MmrApi<Block, MmrRootHash> for RuntimeApi {
+					fn generate_proof(_leaf_index: LeafIndex)
+						-> Result<(EncodableOpaqueLeaf, Proof<MmrRootHash>), MmrError> {
+						unimplemented!()
+					}
+
+					fn verify_proof(_leaf: EncodableOpaqueLeaf, _proof: Proof<MmrRootHash>)
+						-> Result<(), MmrError> {
+						unimplemented!()
+					}
+
+					fn verify_proof_stateless(
+						_root: MmrRootHash,
+						_leaf: EncodableOpaqueLeaf,
+						_proof: Proof<MmrRootHash>
+					) -> Result<(), MmrError> {
+						unimplemented!()
+					}
+
+					fn mmr_root() -> Result<MmrRootHash, MmrError> {
+						Ok($mmr_root)
+					}
+				}
+			}
+		}
+	};
+}
+
+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
+);
+
 fn add_mmr_digest(header: &mut Header, mmr_hash: MmrRootHash) {
 	header.digest_mut().push(DigestItem::Consensus(
 		BEEFY_ENGINE_ID,
@@ -248,54 +308,43 @@ pub(crate) fn create_beefy_keystore(authority: BeefyKeyring) -> SyncCryptoStoreP
 	keystore
 }
 
-pub(crate) fn create_beefy_worker(
-	peer: &BeefyPeer,
-	key: &BeefyKeyring,
-	min_block_delta: u32,
-) -> BeefyWorker<Block, PeersFullClient, Backend, Arc<NetworkService<Block, H256>>> {
-	let keystore = create_beefy_keystore(*key);
-
-	let (signed_commitment_sender, signed_commitment_stream) =
-		BeefySignedCommitmentStream::<Block>::channel();
-	let (beefy_best_block_sender, beefy_best_block_stream) =
-		BeefyBestBlockStream::<Block>::channel();
-
-	let beefy_link_half = BeefyLinkHalf { signed_commitment_stream, beefy_best_block_stream };
-	*peer.data.beefy_link_half.lock() = Some(beefy_link_half);
-	let test_modifiers = peer.data.test_modifiers.clone().unwrap();
-
-	let network = peer.network_service().clone();
-	let sync_oracle = network.clone();
-	let gossip_validator = Arc::new(crate::gossip::GossipValidator::new());
-	let gossip_engine =
-		GossipEngine::new(network, BEEFY_PROTOCOL_NAME, gossip_validator.clone(), None);
-	let worker_params = crate::worker::WorkerParams {
-		client: peer.client().as_client(),
-		backend: peer.client().as_backend(),
-		key_store: Some(keystore).into(),
-		signed_commitment_sender,
-		beefy_best_block_sender,
-		gossip_engine,
-		gossip_validator,
-		min_block_delta,
-		metrics: None,
-		sync_oracle,
-	};
-
-	BeefyWorker::<_, _, _, _>::new(worker_params, test_modifiers)
-}
-
 // Spawns beefy voters. Returns a future to spawn on the runtime.
-fn initialize_beefy(
+fn initialize_beefy<API>(
 	net: &mut BeefyTestNet,
-	peers: &[BeefyKeyring],
+	peers: Vec<(usize, &BeefyKeyring, Arc<API>)>,
 	min_block_delta: u32,
-) -> impl Future<Output = ()> {
+) -> impl Future<Output = ()>
+where
+	API: ProvideRuntimeApi<Block> + Default + Sync + Send,
+	API::Api: BeefyApi<Block> + MmrApi<Block, MmrRootHash>,
+{
 	let voters = FuturesUnordered::new();
 
-	for (peer_id, key) in peers.iter().enumerate() {
-		let worker = create_beefy_worker(&net.peers[peer_id], key, min_block_delta);
-		let gadget = worker.run();
+	for (peer_id, key, api) in peers.into_iter() {
+		let peer = &net.peers[peer_id];
+
+		let keystore = create_beefy_keystore(*key);
+
+		let (signed_commitment_sender, signed_commitment_stream) =
+			BeefySignedCommitmentStream::<Block>::channel();
+		let (beefy_best_block_sender, beefy_best_block_stream) =
+			BeefyBestBlockStream::<Block>::channel();
+		let beefy_link_half = BeefyLinkHalf { signed_commitment_stream, beefy_best_block_stream };
+		*peer.data.beefy_link_half.lock() = Some(beefy_link_half);
+
+		let beefy_params = crate::BeefyParams {
+			client: peer.client().as_client(),
+			backend: peer.client().as_backend(),
+			runtime: api.clone(),
+			key_store: Some(keystore),
+			network: peer.network_service().clone(),
+			signed_commitment_sender,
+			beefy_best_block_sender,
+			min_block_delta,
+			prometheus_registry: None,
+			protocol_name: BEEFY_PROTOCOL_NAME.into(),
+		};
+		let gadget = crate::start_beefy_gadget::<_, _, _, _, _>(beefy_params);
 
 		fn assert_send<T: Send>(_: &T) {}
 		assert_send(&gadget);
@@ -443,13 +492,12 @@ fn beefy_finalizing_blocks() {
 
 	let mut net = BeefyTestNet::new(2, 0);
 
-	for i in 0..peers.len() {
-		net.peer(i).data.use_validator_set(&validator_set);
-	}
-	runtime.spawn(initialize_beefy(&mut net, peers, min_block_delta));
+	let api = Arc::new(two_validators::TestApi {});
+	let beefy_peers = peers.iter().enumerate().map(|(id, key)| (id, key, api.clone())).collect();
+	runtime.spawn(initialize_beefy(&mut net, beefy_peers, min_block_delta));
 
 	// push 42 blocks including `AuthorityChange` digests every 10 blocks.
-	net.generate_blocks(42, session_len, &validator_set);
+	net.generate_blocks(42, session_len, &validator_set, true);
 	net.block_until_sync();
 
 	let net = Arc::new(Mutex::new(net));
@@ -477,19 +525,18 @@ fn lagging_validators() {
 	sp_tracing::try_init_simple();
 
 	let mut runtime = Runtime::new().unwrap();
-	let peers = &[BeefyKeyring::Charlie, BeefyKeyring::Dave];
+	let peers = &[BeefyKeyring::Alice, BeefyKeyring::Bob];
 	let validator_set = ValidatorSet::new(make_beefy_ids(peers), 0).unwrap();
 	let session_len = 30;
 	let min_block_delta = 1;
 
 	let mut net = BeefyTestNet::new(2, 0);
-	for i in 0..peers.len() {
-		net.peer(i).data.use_validator_set(&validator_set);
-	}
-	runtime.spawn(initialize_beefy(&mut net, peers, min_block_delta));
+	let api = Arc::new(two_validators::TestApi {});
+	let beefy_peers = peers.iter().enumerate().map(|(id, key)| (id, key, api.clone())).collect();
+	runtime.spawn(initialize_beefy(&mut net, beefy_peers, min_block_delta));
 
 	// push 42 blocks including `AuthorityChange` digests every 30 blocks.
-	net.generate_blocks(42, session_len, &validator_set);
+	net.generate_blocks(42, session_len, &validator_set, true);
 	net.block_until_sync();
 
 	let net = Arc::new(Mutex::new(net));
@@ -498,7 +545,7 @@ fn lagging_validators() {
 	// diff-power-of-two rule.
 	finalize_block_and_wait_for_beefy(&net, peers, &mut runtime, &[15], &[1, 9, 13, 14, 15]);
 
-	// Charlie finalizes #25, Dave lags behind
+	// Alice finalizes #25, Bob lags behind
 	let finalize = BlockId::number(25);
 	let (best_blocks, signed_commitments) = get_beefy_streams(&mut *net.lock(), peers);
 	net.lock().peer(0).client().as_client().finalize_block(finalize, None).unwrap();
@@ -507,7 +554,7 @@ fn lagging_validators() {
 	streams_empty_after_timeout(best_blocks, &net, &mut runtime, timeout);
 	streams_empty_after_timeout(signed_commitments, &net, &mut runtime, None);
 
-	// Dave catches up and also finalizes #25
+	// Bob catches up and also finalizes #25
 	let (best_blocks, signed_commitments) = get_beefy_streams(&mut *net.lock(), peers);
 	net.lock().peer(1).client().as_client().finalize_block(finalize, None).unwrap();
 	// expected beefy finalizes block #17 from diff-power-of-two
@@ -530,16 +577,23 @@ fn correct_beefy_payload() {
 	let min_block_delta = 2;
 
 	let mut net = BeefyTestNet::new(4, 0);
-	for i in 0..peers.len() {
-		net.peer(i).data.use_validator_set(&validator_set);
-	}
+
+	// Alice, Bob, Charlie will vote on good payloads
+	let good_api = Arc::new(four_validators::TestApi {});
+	let good_peers = [BeefyKeyring::Alice, BeefyKeyring::Bob, BeefyKeyring::Charlie]
+		.iter()
+		.enumerate()
+		.map(|(id, key)| (id, key, good_api.clone()))
+		.collect();
+	runtime.spawn(initialize_beefy(&mut net, good_peers, min_block_delta));
 
 	// Dave will vote on bad mmr roots
-	net.peer(3).data.test_modifiers.as_mut().map(|tm| tm.corrupt_mmr_roots = true);
-	runtime.spawn(initialize_beefy(&mut net, peers, min_block_delta));
+	let bad_api = Arc::new(bad_four_validators::TestApi {});
+	let bad_peers = vec![(3, &BeefyKeyring::Dave, bad_api)];
+	runtime.spawn(initialize_beefy(&mut net, bad_peers, min_block_delta));
 
 	// push 10 blocks
-	net.generate_blocks(12, session_len, &validator_set);
+	net.generate_blocks(12, session_len, &validator_set, false);
 	net.block_until_sync();
 
 	let net = Arc::new(Mutex::new(net));

substrate/client/beefy/src/worker.rs

@@ -26,9 +26,10 @@ use parking_lot::Mutex;
 use sc_client_api::{Backend, FinalityNotification, FinalityNotifications};
 use sc_network_gossip::GossipEngine;
 
-use sp_api::BlockId;
+use sp_api::{BlockId, ProvideRuntimeApi};
 use sp_arithmetic::traits::AtLeast32Bit;
 use sp_consensus::SyncOracle;
+use sp_mmr_primitives::MmrApi;
 use sp_runtime::{
 	generic::OpaqueDigestItemId,
 	traits::{Block, Header, NumberFor},
@@ -52,12 +53,10 @@ use crate::{
 	Client,
 };
 
-pub(crate) struct WorkerParams<B, BE, C, SO>
-where
-	B: Block,
-{
+pub(crate) struct WorkerParams<B: Block, BE, C, R, SO> {
 	pub client: Arc<C>,
 	pub backend: Arc<BE>,
+	pub runtime: Arc<R>,
 	pub key_store: BeefyKeystore,
 	pub signed_commitment_sender: BeefySignedCommitmentSender<B>,
 	pub beefy_best_block_sender: BeefyBestBlockSender<B>,
@@ -69,15 +68,10 @@ where
 }
 
 /// A BEEFY worker plays the BEEFY protocol
-pub(crate) struct BeefyWorker<B, C, BE, SO>
-where
-	B: Block,
-	BE: Backend<B>,
-	C: Client<B, BE>,
-	SO: SyncOracle + Send + Sync + Clone + 'static,
-{
+pub(crate) struct BeefyWorker<B: Block, BE, C, R, SO> {
 	client: Arc<C>,
 	backend: Arc<BE>,
+	runtime: Arc<R>,
 	key_store: BeefyKeystore,
 	signed_commitment_sender: BeefySignedCommitmentSender<B>,
 	gossip_engine: Arc<Mutex<GossipEngine<B>>>,
@@ -99,17 +93,15 @@ where
 	sync_oracle: SO,
 	// keep rustc happy
 	_backend: PhantomData<BE>,
-	#[cfg(test)]
-	// behavior modifiers used in tests
-	test_res: tests::TestModifiers,
 }
 
-impl<B, C, BE, SO> BeefyWorker<B, C, BE, SO>
+impl<B, BE, C, R, SO> BeefyWorker<B, BE, C, R, SO>
 where
 	B: Block + Codec,
 	BE: Backend<B>,
 	C: Client<B, BE>,
-	C::Api: BeefyApi<B>,
+	R: ProvideRuntimeApi<B>,
+	R::Api: BeefyApi<B> + MmrApi<B, MmrRootHash>,
 	SO: SyncOracle + Send + Sync + Clone + 'static,
 {
 	/// Return a new BEEFY worker instance.
@@ -118,15 +110,11 @@ where
 	/// BEEFY pallet has been deployed on-chain.
 	///
 	/// The BEEFY pallet is needed in order to keep track of the BEEFY authority set.
-	pub(crate) fn new(
-		worker_params: WorkerParams<B, BE, C, SO>,
-		#[cfg(test)]
-		// behavior modifiers used in tests
-		test_res: tests::TestModifiers,
-	) -> Self {
+	pub(crate) fn new(worker_params: WorkerParams<B, BE, C, R, SO>) -> Self {
 		let WorkerParams {
 			client,
 			backend,
+			runtime,
 			key_store,
 			signed_commitment_sender,
 			beefy_best_block_sender,
@@ -144,6 +132,7 @@ where
 		BeefyWorker {
 			client: client.clone(),
 			backend,
+			runtime,
 			key_store,
 			signed_commitment_sender,
 			gossip_engine: Arc::new(Mutex::new(gossip_engine)),
@@ -159,20 +148,9 @@ where
 			beefy_best_block_sender,
 			sync_oracle,
 			_backend: PhantomData,
-			#[cfg(test)]
-			test_res,
 		}
 	}
-}
 
-impl<B, C, BE, SO> BeefyWorker<B, C, BE, SO>
-where
-	B: Block,
-	BE: Backend<B>,
-	C: Client<B, BE>,
-	C::Api: BeefyApi<B>,
-	SO: SyncOracle + Send + Sync + Clone + 'static,
-{
 	/// Return `Some(number)` if we should be voting on block `number` now,
 	/// return `None` if there is no block we should vote on now.
 	fn current_vote_target(&self) -> Option<NumberFor<B>> {
@@ -396,7 +374,7 @@ where
 		};
 		let target_hash = target_header.hash();
 
-		let mmr_root = if let Some(hash) = self.extract_mmr_root_digest(&target_header) {
+		let mmr_root = if let Some(hash) = self.get_mmr_root_digest(&target_header) {
 			hash
 		} else {
 			warn!(target: "beefy", "🥩 No MMR root digest found for: {:?}", target_hash);
@@ -458,13 +436,12 @@ where
 	}
 
 	/// Wait for BEEFY runtime pallet to be available.
-	#[cfg(not(test))]
 	async fn wait_for_runtime_pallet(&mut self) {
 		self.client
 			.finality_notification_stream()
 			.take_while(|notif| {
 				let at = BlockId::hash(notif.header.hash());
-				if let Some(active) = self.client.runtime_api().validator_set(&at).ok().flatten() {
+				if let Some(active) = self.runtime.runtime_api().validator_set(&at).ok().flatten() {
 					if active.id() == GENESIS_AUTHORITY_SET_ID {
 						// When starting from genesis, there is no session boundary digest.
 						// Just initialize `rounds` to Block #1 as BEEFY mandatory block.
@@ -488,13 +465,6 @@ where
 		self.finality_notifications = self.client.finality_notification_stream();
 	}
 
-	/// For tests don't use runtime pallet. Start rounds from block #1.
-	#[cfg(test)]
-	async fn wait_for_runtime_pallet(&mut self) {
-		let active = self.test_res.active_validators.clone();
-		self.init_session_at(active, 1u32.into());
-	}
-
 	/// Main loop for BEEFY worker.
 	///
 	/// Wait for BEEFY runtime pallet to be available, then start the main async loop
@@ -550,20 +520,15 @@ where
 		}
 	}
 
-	/// Simple wrapper over mmr root extraction.
-	#[cfg(not(test))]
-	fn extract_mmr_root_digest(&self, header: &B::Header) -> Option<MmrRootHash> {
-		find_mmr_root_digest::<B>(header)
-	}
-
-	/// For tests, have the option to modify mmr root.
-	#[cfg(test)]
-	fn extract_mmr_root_digest(&self, header: &B::Header) -> Option<MmrRootHash> {
-		let mut mmr_root = find_mmr_root_digest::<B>(header);
-		if self.test_res.corrupt_mmr_roots {
-			mmr_root.as_mut().map(|hash| *hash ^= MmrRootHash::random());
-		}
-		mmr_root
+	/// Simple wrapper that gets MMR root from header digests or from client state.
+	fn get_mmr_root_digest(&self, header: &B::Header) -> Option<MmrRootHash> {
+		find_mmr_root_digest::<B>(header).or_else(|| {
+			self.runtime
+				.runtime_api()
+				.mmr_root(&BlockId::hash(header.hash()))
+				.ok()
+				.and_then(|r| r.ok())
+		})
 	}
 }
 
@@ -658,11 +623,16 @@ pub(crate) mod tests {
 	use super::*;
 	use crate::{
 		keystore::tests::Keyring,
-		tests::{create_beefy_worker, get_beefy_streams, make_beefy_ids, BeefyTestNet},
+		notification::{BeefyBestBlockStream, BeefySignedCommitmentStream},
+		tests::{
+			create_beefy_keystore, get_beefy_streams, make_beefy_ids, two_validators::TestApi,
+			BeefyPeer, BeefyTestNet, BEEFY_PROTOCOL_NAME,
+		},
 	};
 
 	use futures::{executor::block_on, future::poll_fn, task::Poll};
 
+	use crate::tests::BeefyLinkHalf;
 	use sc_client_api::HeaderBackend;
 	use sc_network::NetworkService;
 	use sc_network_test::{PeersFullClient, TestNetFactory};
@@ -672,10 +642,40 @@ pub(crate) mod tests {
 		Backend,
 	};
 
-	#[derive(Clone)]
-	pub struct TestModifiers {
-		pub active_validators: ValidatorSet<AuthorityId>,
-		pub corrupt_mmr_roots: bool,
+	fn create_beefy_worker(
+		peer: &BeefyPeer,
+		key: &Keyring,
+		min_block_delta: u32,
+	) -> BeefyWorker<Block, Backend, PeersFullClient, TestApi, Arc<NetworkService<Block, H256>>> {
+		let keystore = create_beefy_keystore(*key);
+
+		let (signed_commitment_sender, signed_commitment_stream) =
+			BeefySignedCommitmentStream::<Block>::channel();
+		let (beefy_best_block_sender, beefy_best_block_stream) =
+			BeefyBestBlockStream::<Block>::channel();
+		let beefy_link_half = BeefyLinkHalf { signed_commitment_stream, beefy_best_block_stream };
+		*peer.data.beefy_link_half.lock() = Some(beefy_link_half);
+
+		let api = Arc::new(TestApi {});
+		let network = peer.network_service().clone();
+		let sync_oracle = network.clone();
+		let gossip_validator = Arc::new(crate::gossip::GossipValidator::new());
+		let gossip_engine =
+			GossipEngine::new(network, BEEFY_PROTOCOL_NAME, gossip_validator.clone(), None);
+		let worker_params = crate::worker::WorkerParams {
+			client: peer.client().as_client(),
+			backend: peer.client().as_backend(),
+			runtime: api,
+			key_store: Some(keystore).into(),
+			signed_commitment_sender,
+			beefy_best_block_sender,
+			gossip_engine,
+			gossip_validator,
+			min_block_delta,
+			metrics: None,
+			sync_oracle,
+		};
+		BeefyWorker::<_, _, _, _, _>::new(worker_params)
 	}
 
 	#[test]
@@ -825,7 +825,6 @@ pub(crate) mod tests {
 		let keys = &[Keyring::Alice];
 		let validator_set = ValidatorSet::new(make_beefy_ids(keys), 0).unwrap();
 		let mut net = BeefyTestNet::new(1, 0);
-		net.peer(0).data.use_validator_set(&validator_set);
 		let mut worker = create_beefy_worker(&net.peer(0), &keys[0], 1);
 
 		// rounds not initialized -> should vote: `None`
@@ -833,8 +832,9 @@ pub(crate) mod tests {
 
 		let set_up = |worker: &mut BeefyWorker<
 			Block,
-			PeersFullClient,
 			Backend,
+			PeersFullClient,
+			TestApi,
 			Arc<NetworkService<Block, H256>>,
 		>,
 		              best_grandpa: u64,
@@ -889,7 +889,6 @@ pub(crate) mod tests {
 		let keys = &[Keyring::Alice];
 		let validator_set = ValidatorSet::new(make_beefy_ids(keys), 0).unwrap();
 		let mut net = BeefyTestNet::new(1, 0);
-		net.peer(0).data.use_validator_set(&validator_set);
 		let mut worker = create_beefy_worker(&net.peer(0), &keys[0], 1);
 
 		// keystore doesn't contain other keys than validators'
@@ -913,7 +912,6 @@ pub(crate) mod tests {
 		let keys = &[Keyring::Alice];
 		let validator_set = ValidatorSet::new(make_beefy_ids(keys), 0).unwrap();
 		let mut net = BeefyTestNet::new(1, 0);
-		net.peer(0).data.use_validator_set(&validator_set);
 		let mut worker = create_beefy_worker(&net.peer(0), &keys[0], 1);
 
 		let (mut best_block_streams, _) = get_beefy_streams(&mut net, keys);
@@ -939,7 +937,7 @@ pub(crate) mod tests {
 		// generate 2 blocks, try again expect success
 		let (mut best_block_streams, _) = get_beefy_streams(&mut net, keys);
 		let mut best_block_stream = best_block_streams.drain(..).next().unwrap();
-		net.generate_blocks(2, 10, &validator_set);
+		net.generate_blocks(2, 10, &validator_set, false);
 
 		worker.set_best_beefy_block(2);
 		assert_eq!(worker.best_beefy_block, Some(2));
@@ -961,7 +959,6 @@ pub(crate) mod tests {
 		let keys = &[Keyring::Alice];
 		let validator_set = ValidatorSet::new(make_beefy_ids(keys), 0).unwrap();
 		let mut net = BeefyTestNet::new(1, 0);
-		net.peer(0).data.use_validator_set(&validator_set);
 		let mut worker = create_beefy_worker(&net.peer(0), &keys[0], 1);
 
 		assert!(worker.rounds.is_none());