Clean up tests from runtime.block_on() and moving around Runtime and Handle (#12941)

substrate/client/beefy/src/aux_schema.rs

@@ -77,7 +77,6 @@ pub(crate) mod tests {
 	use super::*;
 	use crate::tests::BeefyTestNet;
 	use sc_network_test::TestNetFactory;
-	use tokio::runtime::Runtime;
 
 	// also used in tests.rs
 	pub fn verify_persisted_version<B: BlockT, BE: Backend<B>>(backend: &BE) -> bool {
@@ -85,10 +84,9 @@ pub(crate) mod tests {
 		version == CURRENT_VERSION
 	}
 
-	#[test]
-	fn should_load_persistent_sanity_checks() {
-		let runtime = Runtime::new().unwrap();
-		let mut net = BeefyTestNet::new(runtime.handle().clone(), 1);
+	#[tokio::test]
+	async fn should_load_persistent_sanity_checks() {
+		let mut net = BeefyTestNet::new(1);
 		let backend = net.peer(0).client().as_backend();
 
 		// version not available in db -> None

substrate/client/beefy/src/tests.rs

@@ -47,7 +47,7 @@ use sc_consensus::{
 use sc_network::{config::RequestResponseConfig, ProtocolName};
 use sc_network_test::{
 	Block, BlockImportAdapter, FullPeerConfig, PassThroughVerifier, Peer, PeersClient,
-	PeersFullClient, TestNetFactory, WithRuntime,
+	PeersFullClient, TestNetFactory,
 };
 use sc_utils::notification::NotificationReceiver;
 use serde::{Deserialize, Serialize};
@@ -64,10 +64,7 @@ use sp_runtime::{
 };
 use std::{collections::HashMap, marker::PhantomData, sync::Arc, task::Poll};
 use substrate_test_runtime_client::{runtime::Header, ClientExt};
-use tokio::{
-	runtime::{Handle, Runtime},
-	time::Duration,
-};
+use tokio::time::Duration;
 
 const GENESIS_HASH: H256 = H256::zero();
 fn beefy_gossip_proto_name() -> ProtocolName {
@@ -106,23 +103,14 @@ pub(crate) struct PeerData {
 		Mutex<Option<BeefyJustifsRequestHandler<Block, PeersFullClient>>>,
 }
 
+#[derive(Default)]
 pub(crate) struct BeefyTestNet {
-	rt_handle: Handle,
 	peers: Vec<BeefyPeer>,
 }
 
-impl WithRuntime for BeefyTestNet {
-	fn with_runtime(rt_handle: Handle) -> Self {
-		BeefyTestNet { rt_handle, peers: Vec::new() }
-	}
-	fn rt_handle(&self) -> &Handle {
-		&self.rt_handle
-	}
-}
-
 impl BeefyTestNet {
-	pub(crate) fn new(rt_handle: Handle, n_authority: usize) -> Self {
-		let mut net = BeefyTestNet::with_runtime(rt_handle);
+	pub(crate) fn new(n_authority: usize) -> Self {
+		let mut net = BeefyTestNet { peers: Vec::with_capacity(n_authority) };
 
 		for i in 0..n_authority {
 			let (rx, cfg) = on_demand_justifications_protocol_config(GENESIS_HASH, None);
@@ -151,13 +139,12 @@ impl BeefyTestNet {
 		});
 	}
 
-	pub(crate) fn generate_blocks_and_sync(
+	pub(crate) async fn generate_blocks_and_sync(
 		&mut self,
 		count: usize,
 		session_length: u64,
 		validator_set: &BeefyValidatorSet,
 		include_mmr_digest: bool,
-		runtime: &mut Runtime,
 	) {
 		self.peer(0).generate_blocks(count, BlockOrigin::File, |builder| {
 			let mut block = builder.build().unwrap().block;
@@ -175,7 +162,7 @@ impl BeefyTestNet {
 
 			block
 		});
-		runtime.block_on(self.wait_until_sync());
+		self.run_until_sync().await;
 	}
 }
 
@@ -404,17 +391,16 @@ where
 	tasks.for_each(|_| async move {})
 }
 
-fn block_until(future: impl Future + Unpin, net: &Arc<Mutex<BeefyTestNet>>, runtime: &mut Runtime) {
+async fn run_until(future: impl Future + Unpin, net: &Arc<Mutex<BeefyTestNet>>) {
 	let drive_to_completion = futures::future::poll_fn(|cx| {
 		net.lock().poll(cx);
 		Poll::<()>::Pending
 	});
-	runtime.block_on(future::select(future, drive_to_completion));
+	let _ = future::select(future, drive_to_completion).await;
 }
 
-fn run_for(duration: Duration, net: &Arc<Mutex<BeefyTestNet>>, runtime: &mut Runtime) {
-	let sleep = runtime.spawn(async move { tokio::time::sleep(duration).await });
-	block_until(sleep, net, runtime);
+async fn run_for(duration: Duration, net: &Arc<Mutex<BeefyTestNet>>) {
+	run_until(Box::pin(tokio::time::sleep(duration)), net).await;
 }
 
 pub(crate) fn get_beefy_streams(
@@ -435,10 +421,9 @@ pub(crate) fn get_beefy_streams(
 	(best_block_streams, versioned_finality_proof_streams)
 }
 
-fn wait_for_best_beefy_blocks(
+async fn wait_for_best_beefy_blocks(
 	streams: Vec<NotificationReceiver<H256>>,
 	net: &Arc<Mutex<BeefyTestNet>>,
-	runtime: &mut Runtime,
 	expected_beefy_blocks: &[u64],
 ) {
 	let mut wait_for = Vec::new();
@@ -458,13 +443,12 @@ fn wait_for_best_beefy_blocks(
 		})));
 	});
 	let wait_for = futures::future::join_all(wait_for);
-	block_until(wait_for, net, runtime);
+	run_until(wait_for, net).await;
 }
 
-fn wait_for_beefy_signed_commitments(
+async fn wait_for_beefy_signed_commitments(
 	streams: Vec<NotificationReceiver<BeefyVersionedFinalityProof<Block>>>,
 	net: &Arc<Mutex<BeefyTestNet>>,
-	runtime: &mut Runtime,
 	expected_commitment_block_nums: &[u64],
 ) {
 	let mut wait_for = Vec::new();
@@ -484,34 +468,33 @@ fn wait_for_beefy_signed_commitments(
 		})));
 	});
 	let wait_for = futures::future::join_all(wait_for);
-	block_until(wait_for, net, runtime);
+	run_until(wait_for, net).await;
 }
 
-fn streams_empty_after_timeout<T>(
+async fn streams_empty_after_timeout<T>(
 	streams: Vec<NotificationReceiver<T>>,
 	net: &Arc<Mutex<BeefyTestNet>>,
-	runtime: &mut Runtime,
 	timeout: Option<Duration>,
 ) where
 	T: std::fmt::Debug,
 	T: std::cmp::PartialEq,
 {
 	if let Some(timeout) = timeout {
-		run_for(timeout, net, runtime);
+		run_for(timeout, net).await;
 	}
-	streams.into_iter().for_each(|mut stream| {
-		runtime.block_on(future::poll_fn(move |cx| {
+	for mut stream in streams.into_iter() {
+		future::poll_fn(move |cx| {
 			assert_eq!(stream.poll_next_unpin(cx), Poll::Pending);
 			Poll::Ready(())
-		}));
-	});
+		})
+		.await;
+	}
 }
 
-fn finalize_block_and_wait_for_beefy(
+async fn finalize_block_and_wait_for_beefy(
 	net: &Arc<Mutex<BeefyTestNet>>,
 	// peer index and key
 	peers: impl Iterator<Item = (usize, BeefyKeyring)> + Clone,
-	runtime: &mut Runtime,
 	finalize_targets: &[u64],
 	expected_beefy: &[u64],
 ) {
@@ -528,33 +511,32 @@ fn finalize_block_and_wait_for_beefy(
 	if expected_beefy.is_empty() {
 		// run for quarter second then verify no new best beefy block available
 		let timeout = Some(Duration::from_millis(250));
-		streams_empty_after_timeout(best_blocks, &net, runtime, timeout);
-		streams_empty_after_timeout(versioned_finality_proof, &net, runtime, None);
+		streams_empty_after_timeout(best_blocks, &net, timeout).await;
+		streams_empty_after_timeout(versioned_finality_proof, &net, None).await;
 	} else {
 		// run until expected beefy blocks are received
-		wait_for_best_beefy_blocks(best_blocks, &net, runtime, expected_beefy);
-		wait_for_beefy_signed_commitments(versioned_finality_proof, &net, runtime, expected_beefy);
+		wait_for_best_beefy_blocks(best_blocks, &net, expected_beefy).await;
+		wait_for_beefy_signed_commitments(versioned_finality_proof, &net, expected_beefy).await;
 	}
 }
 
-#[test]
-fn beefy_finalizing_blocks() {
+#[tokio::test]
+async fn beefy_finalizing_blocks() {
 	sp_tracing::try_init_simple();
 
-	let mut runtime = Runtime::new().unwrap();
 	let peers = [BeefyKeyring::Alice, BeefyKeyring::Bob];
 	let validator_set = ValidatorSet::new(make_beefy_ids(&peers), 0).unwrap();
 	let session_len = 10;
 	let min_block_delta = 4;
 
-	let mut net = BeefyTestNet::new(runtime.handle().clone(), 2);
+	let mut net = BeefyTestNet::new(2);
 
 	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));
+	tokio::spawn(initialize_beefy(&mut net, beefy_peers, min_block_delta));
 
 	// push 42 blocks including `AuthorityChange` digests every 10 blocks.
-	net.generate_blocks_and_sync(42, session_len, &validator_set, true, &mut runtime);
+	net.generate_blocks_and_sync(42, session_len, &validator_set, true).await;
 
 	let net = Arc::new(Mutex::new(net));
 
@@ -562,51 +544,44 @@ fn beefy_finalizing_blocks() {
 
 	let peers = peers.into_iter().enumerate();
 	// finalize block #5 -> BEEFY should finalize #1 (mandatory) and #5 from diff-power-of-two rule.
-	finalize_block_and_wait_for_beefy(&net, peers.clone(), &mut runtime, &[1, 5], &[1, 5]);
+	finalize_block_and_wait_for_beefy(&net, peers.clone(), &[1, 5], &[1, 5]).await;
 
 	// GRANDPA finalize #10 -> BEEFY finalize #10 (mandatory)
-	finalize_block_and_wait_for_beefy(&net, peers.clone(), &mut runtime, &[10], &[10]);
+	finalize_block_and_wait_for_beefy(&net, peers.clone(), &[10], &[10]).await;
 
 	// GRANDPA finalize #18 -> BEEFY finalize #14, then #18 (diff-power-of-two rule)
-	finalize_block_and_wait_for_beefy(&net, peers.clone(), &mut runtime, &[18], &[14, 18]);
+	finalize_block_and_wait_for_beefy(&net, peers.clone(), &[18], &[14, 18]).await;
 
 	// GRANDPA finalize #20 -> BEEFY finalize #20 (mandatory)
-	finalize_block_and_wait_for_beefy(&net, peers.clone(), &mut runtime, &[20], &[20]);
+	finalize_block_and_wait_for_beefy(&net, peers.clone(), &[20], &[20]).await;
 
 	// GRANDPA finalize #21 -> BEEFY finalize nothing (yet) because min delta is 4
-	finalize_block_and_wait_for_beefy(&net, peers, &mut runtime, &[21], &[]);
+	finalize_block_and_wait_for_beefy(&net, peers, &[21], &[]).await;
 }
 
-#[test]
-fn lagging_validators() {
+#[tokio::test]
+async fn lagging_validators() {
 	sp_tracing::try_init_simple();
 
-	let mut runtime = Runtime::new().unwrap();
 	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(runtime.handle().clone(), 2);
+	let mut net = BeefyTestNet::new(2);
 	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));
+	tokio::spawn(initialize_beefy(&mut net, beefy_peers, min_block_delta));
 
 	// push 62 blocks including `AuthorityChange` digests every 30 blocks.
-	net.generate_blocks_and_sync(62, session_len, &validator_set, true, &mut runtime);
+	net.generate_blocks_and_sync(62, session_len, &validator_set, true).await;
 
 	let net = Arc::new(Mutex::new(net));
 
 	let peers = peers.into_iter().enumerate();
 	// finalize block #15 -> BEEFY should finalize #1 (mandatory) and #9, #13, #14, #15 from
 	// diff-power-of-two rule.
-	finalize_block_and_wait_for_beefy(
-		&net,
-		peers.clone(),
-		&mut runtime,
-		&[1, 15],
-		&[1, 9, 13, 14, 15],
-	);
+	finalize_block_and_wait_for_beefy(&net, peers.clone(), &[1, 15], &[1, 9, 13, 14, 15]).await;
 
 	// Alice finalizes #25, Bob lags behind
 	let finalize = net
@@ -620,18 +595,18 @@ fn lagging_validators() {
 	net.lock().peer(0).client().as_client().finalize_block(finalize, None).unwrap();
 	// verify nothing gets finalized by BEEFY
 	let timeout = Some(Duration::from_millis(250));
-	streams_empty_after_timeout(best_blocks, &net, &mut runtime, timeout);
-	streams_empty_after_timeout(versioned_finality_proof, &net, &mut runtime, None);
+	streams_empty_after_timeout(best_blocks, &net, timeout).await;
+	streams_empty_after_timeout(versioned_finality_proof, &net, None).await;
 
 	// Bob catches up and also finalizes #25
 	let (best_blocks, versioned_finality_proof) = get_beefy_streams(&mut net.lock(), peers.clone());
 	net.lock().peer(1).client().as_client().finalize_block(finalize, None).unwrap();
 	// expected beefy finalizes block #17 from diff-power-of-two
-	wait_for_best_beefy_blocks(best_blocks, &net, &mut runtime, &[23, 24, 25]);
-	wait_for_beefy_signed_commitments(versioned_finality_proof, &net, &mut runtime, &[23, 24, 25]);
+	wait_for_best_beefy_blocks(best_blocks, &net, &[23, 24, 25]).await;
+	wait_for_beefy_signed_commitments(versioned_finality_proof, &net, &[23, 24, 25]).await;
 
 	// Both finalize #30 (mandatory session) and #32 -> BEEFY finalize #30 (mandatory), #31, #32
-	finalize_block_and_wait_for_beefy(&net, peers.clone(), &mut runtime, &[30, 32], &[30, 31, 32]);
+	finalize_block_and_wait_for_beefy(&net, peers.clone(), &[30, 32], &[30, 31, 32]).await;
 
 	// Verify that session-boundary votes get buffered by client and only processed once
 	// session-boundary block is GRANDPA-finalized (this guarantees authenticity for the new session
@@ -649,28 +624,27 @@ fn lagging_validators() {
 	net.lock().peer(0).client().as_client().finalize_block(finalize, None).unwrap();
 	// verify nothing gets finalized by BEEFY
 	let timeout = Some(Duration::from_millis(250));
-	streams_empty_after_timeout(best_blocks, &net, &mut runtime, timeout);
-	streams_empty_after_timeout(versioned_finality_proof, &net, &mut runtime, None);
+	streams_empty_after_timeout(best_blocks, &net, timeout).await;
+	streams_empty_after_timeout(versioned_finality_proof, &net, None).await;
 
 	// Bob catches up and also finalizes #60 (and should have buffered Alice's vote on #60)
 	let (best_blocks, versioned_finality_proof) = get_beefy_streams(&mut net.lock(), peers);
 	net.lock().peer(1).client().as_client().finalize_block(finalize, None).unwrap();
 	// verify beefy skips intermediary votes, and successfully finalizes mandatory block #60
-	wait_for_best_beefy_blocks(best_blocks, &net, &mut runtime, &[60]);
-	wait_for_beefy_signed_commitments(versioned_finality_proof, &net, &mut runtime, &[60]);
+	wait_for_best_beefy_blocks(best_blocks, &net, &[60]).await;
+	wait_for_beefy_signed_commitments(versioned_finality_proof, &net, &[60]).await;
 }
 
-#[test]
-fn correct_beefy_payload() {
+#[tokio::test]
+async fn correct_beefy_payload() {
 	sp_tracing::try_init_simple();
 
-	let mut runtime = Runtime::new().unwrap();
 	let peers = [BeefyKeyring::Alice, BeefyKeyring::Bob, BeefyKeyring::Charlie, BeefyKeyring::Dave];
 	let validator_set = ValidatorSet::new(make_beefy_ids(&peers), 0).unwrap();
 	let session_len = 20;
 	let min_block_delta = 2;
 
-	let mut net = BeefyTestNet::new(runtime.handle().clone(), 4);
+	let mut net = BeefyTestNet::new(4);
 
 	// Alice, Bob, Charlie will vote on good payloads
 	let good_api = Arc::new(four_validators::TestApi {});
@@ -679,20 +653,20 @@ fn correct_beefy_payload() {
 		.enumerate()
 		.map(|(id, key)| (id, key, good_api.clone()))
 		.collect();
-	runtime.spawn(initialize_beefy(&mut net, good_peers, min_block_delta));
+	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_peers = vec![(3, &BeefyKeyring::Dave, bad_api)];
-	runtime.spawn(initialize_beefy(&mut net, bad_peers, min_block_delta));
+	tokio::spawn(initialize_beefy(&mut net, bad_peers, min_block_delta));
 
 	// push 12 blocks
-	net.generate_blocks_and_sync(12, session_len, &validator_set, false, &mut runtime);
+	net.generate_blocks_and_sync(12, session_len, &validator_set, false).await;
 
 	let net = Arc::new(Mutex::new(net));
 	let peers = peers.into_iter().enumerate();
 	// with 3 good voters and 1 bad one, consensus should happen and best blocks produced.
-	finalize_block_and_wait_for_beefy(&net, peers, &mut runtime, &[1, 10], &[1, 9]);
+	finalize_block_and_wait_for_beefy(&net, peers, &[1, 10], &[1, 9]).await;
 
 	let (best_blocks, versioned_finality_proof) =
 		get_beefy_streams(&mut net.lock(), [(0, BeefyKeyring::Alice)].into_iter());
@@ -711,8 +685,8 @@ fn correct_beefy_payload() {
 
 	// verify consensus is _not_ reached
 	let timeout = Some(Duration::from_millis(250));
-	streams_empty_after_timeout(best_blocks, &net, &mut runtime, timeout);
-	streams_empty_after_timeout(versioned_finality_proof, &net, &mut runtime, None);
+	streams_empty_after_timeout(best_blocks, &net, timeout).await;
+	streams_empty_after_timeout(versioned_finality_proof, &net, None).await;
 
 	// 3rd good validator catches up and votes as well
 	let (best_blocks, versioned_finality_proof) =
@@ -720,21 +694,19 @@ fn correct_beefy_payload() {
 	net.lock().peer(2).client().as_client().finalize_block(hashof11, None).unwrap();
 
 	// verify consensus is reached
-	wait_for_best_beefy_blocks(best_blocks, &net, &mut runtime, &[11]);
-	wait_for_beefy_signed_commitments(versioned_finality_proof, &net, &mut runtime, &[11]);
+	wait_for_best_beefy_blocks(best_blocks, &net, &[11]).await;
+	wait_for_beefy_signed_commitments(versioned_finality_proof, &net, &[11]).await;
 }
 
-#[test]
-fn beefy_importing_blocks() {
+#[tokio::test]
+async fn beefy_importing_blocks() {
 	use futures::{future::poll_fn, task::Poll};
 	use sc_block_builder::BlockBuilderProvider;
 	use sc_client_api::BlockBackend;
 
 	sp_tracing::try_init_simple();
 
-	let runtime = Runtime::new().unwrap();
-
-	let mut net = BeefyTestNet::new(runtime.handle().clone(), 2);
+	let mut net = BeefyTestNet::new(2);
 
 	let client = net.peer(0).client().clone();
 	let (mut block_import, _, peer_data) = net.make_block_import(client.clone());
@@ -759,16 +731,15 @@ fn beefy_importing_blocks() {
 	// Import without justifications.
 	let mut justif_recv = justif_stream.subscribe();
 	assert_eq!(
-		runtime
-			.block_on(block_import.import_block(params(block.clone(), None), HashMap::new()))
+		block_import
+			.import_block(params(block.clone(), None), HashMap::new())
+			.await
 			.unwrap(),
 		ImportResult::Imported(ImportedAux { is_new_best: true, ..Default::default() }),
 	);
 	assert_eq!(
-		runtime
-			.block_on(block_import.import_block(params(block, None), HashMap::new()))
-			.unwrap(),
-		ImportResult::AlreadyInChain
+		block_import.import_block(params(block, None), HashMap::new()).await.unwrap(),
+		ImportResult::AlreadyInChain,
 	);
 	// Verify no BEEFY justifications present:
 	{
@@ -781,10 +752,11 @@ fn beefy_importing_blocks() {
 			None
 		);
 		// and none sent to BEEFY worker.
-		runtime.block_on(poll_fn(move |cx| {
+		poll_fn(move |cx| {
 			assert_eq!(justif_recv.poll_next_unpin(cx), Poll::Pending);
 			Poll::Ready(())
-		}));
+		})
+		.await;
 	}
 
 	// Import with valid justification.
@@ -802,9 +774,7 @@ fn beefy_importing_blocks() {
 	let hashof2 = block.header.hash();
 	let mut justif_recv = justif_stream.subscribe();
 	assert_eq!(
-		runtime
-			.block_on(block_import.import_block(params(block, justif), HashMap::new()))
-			.unwrap(),
+		block_import.import_block(params(block, justif), HashMap::new()).await.unwrap(),
 		ImportResult::Imported(ImportedAux {
 			bad_justification: false,
 			is_new_best: true,
@@ -823,13 +793,14 @@ fn beefy_importing_blocks() {
 		);
 		// but sent to BEEFY worker
 		// (worker will append it to backend when all previous mandatory justifs are there as well).
-		runtime.block_on(poll_fn(move |cx| {
+		poll_fn(move |cx| {
 			match justif_recv.poll_next_unpin(cx) {
 				Poll::Ready(Some(_justification)) => (),
 				v => panic!("unexpected value: {:?}", v),
 			}
 			Poll::Ready(())
-		}));
+		})
+		.await;
 	}
 
 	// Import with invalid justification (incorrect validator set).
@@ -847,9 +818,7 @@ fn beefy_importing_blocks() {
 	let hashof3 = block.header.hash();
 	let mut justif_recv = justif_stream.subscribe();
 	assert_eq!(
-		runtime
-			.block_on(block_import.import_block(params(block, justif), HashMap::new()))
-			.unwrap(),
+		block_import.import_block(params(block, justif), HashMap::new()).await.unwrap(),
 		ImportResult::Imported(ImportedAux {
 			// Still `false` because we don't want to fail import on bad BEEFY justifications.
 			bad_justification: false,
@@ -868,33 +837,33 @@ fn beefy_importing_blocks() {
 			None
 		);
 		// and none sent to BEEFY worker.
-		runtime.block_on(poll_fn(move |cx| {
+		poll_fn(move |cx| {
 			assert_eq!(justif_recv.poll_next_unpin(cx), Poll::Pending);
 			Poll::Ready(())
-		}));
+		})
+		.await;
 	}
 }
 
-#[test]
-fn voter_initialization() {
+#[tokio::test]
+async fn voter_initialization() {
 	sp_tracing::try_init_simple();
 	// Regression test for voter initialization where finality notifications were dropped
 	// after waiting for BEEFY pallet availability.
 
-	let mut runtime = Runtime::new().unwrap();
 	let peers = [BeefyKeyring::Alice, BeefyKeyring::Bob];
 	let validator_set = ValidatorSet::new(make_beefy_ids(&peers), 0).unwrap();
 	let session_len = 5;
 	// Should vote on all mandatory blocks no matter the `min_block_delta`.
 	let min_block_delta = 10;
 
-	let mut net = BeefyTestNet::new(runtime.handle().clone(), 2);
+	let mut net = BeefyTestNet::new(2);
 	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));
+	tokio::spawn(initialize_beefy(&mut net, beefy_peers, min_block_delta));
 
 	// push 26 blocks
-	net.generate_blocks_and_sync(26, session_len, &validator_set, false, &mut runtime);
+	net.generate_blocks_and_sync(26, session_len, &validator_set, false).await;
 	let net = Arc::new(Mutex::new(net));
 
 	// Finalize multiple blocks at once to get a burst of finality notifications right from start.
@@ -903,31 +872,30 @@ fn voter_initialization() {
 	finalize_block_and_wait_for_beefy(
 		&net,
 		peers.into_iter().enumerate(),
-		&mut runtime,
 		&[1, 6, 10, 17, 24, 26],
 		&[1, 5, 10, 15, 20, 25],
-	);
+	)
+	.await;
 }
 
-#[test]
-fn on_demand_beefy_justification_sync() {
+#[tokio::test]
+async fn on_demand_beefy_justification_sync() {
 	sp_tracing::try_init_simple();
 
-	let mut runtime = Runtime::new().unwrap();
 	let all_peers =
 		[BeefyKeyring::Alice, BeefyKeyring::Bob, BeefyKeyring::Charlie, BeefyKeyring::Dave];
 	let validator_set = ValidatorSet::new(make_beefy_ids(&all_peers), 0).unwrap();
 	let session_len = 5;
 	let min_block_delta = 5;
 
-	let mut net = BeefyTestNet::new(runtime.handle().clone(), 4);
+	let mut net = BeefyTestNet::new(4);
 
 	// Alice, Bob, Charlie start first and make progress through voting.
 	let api = Arc::new(four_validators::TestApi {});
 	let fast_peers = [BeefyKeyring::Alice, BeefyKeyring::Bob, BeefyKeyring::Charlie];
 	let voting_peers =
 		fast_peers.iter().enumerate().map(|(id, key)| (id, key, api.clone())).collect();
-	runtime.spawn(initialize_beefy(&mut net, voting_peers, min_block_delta));
+	tokio::spawn(initialize_beefy(&mut net, voting_peers, min_block_delta));
 
 	// Dave will start late and have to catch up using on-demand justification requests (since
 	// in this test there is no block import queue to automatically import justifications).
@@ -937,7 +905,7 @@ fn on_demand_beefy_justification_sync() {
 	let dave_index = 3;
 
 	// push 30 blocks
-	net.generate_blocks_and_sync(30, session_len, &validator_set, false, &mut runtime);
+	net.generate_blocks_and_sync(30, session_len, &validator_set, false).await;
 
 	let fast_peers = fast_peers.into_iter().enumerate();
 	let net = Arc::new(Mutex::new(net));
@@ -946,15 +914,15 @@ fn on_demand_beefy_justification_sync() {
 	finalize_block_and_wait_for_beefy(
 		&net,
 		fast_peers.clone(),
-		&mut runtime,
 		&[1, 6, 10, 17, 24],
 		&[1, 5, 10, 15, 20],
-	);
+	)
+	.await;
 
 	// Spawn Dave, he's now way behind voting and can only catch up through on-demand justif sync.
-	runtime.spawn(dave_task);
+	tokio::spawn(dave_task);
 	// give Dave a chance to spawn and init.
-	run_for(Duration::from_millis(400), &net, &mut runtime);
+	run_for(Duration::from_millis(400), &net).await;
 
 	let (dave_best_blocks, _) =
 		get_beefy_streams(&mut net.lock(), [(dave_index, BeefyKeyring::Dave)].into_iter());
@@ -962,41 +930,40 @@ fn on_demand_beefy_justification_sync() {
 	let hashof1 = client.expect_block_hash_from_id(&BlockId::number(1)).unwrap();
 	client.finalize_block(hashof1, None).unwrap();
 	// Give Dave task some cpu cycles to process the finality notification,
-	run_for(Duration::from_millis(100), &net, &mut runtime);
+	run_for(Duration::from_millis(100), &net).await;
 	// freshly spun up Dave now needs to listen for gossip to figure out the state of his peers.
 
 	// Have the other peers do some gossip so Dave finds out about their progress.
-	finalize_block_and_wait_for_beefy(&net, fast_peers, &mut runtime, &[25], &[25]);
+	finalize_block_and_wait_for_beefy(&net, fast_peers, &[25], &[25]).await;
 
 	// Now verify Dave successfully finalized #1 (through on-demand justification request).
-	wait_for_best_beefy_blocks(dave_best_blocks, &net, &mut runtime, &[1]);
+	wait_for_best_beefy_blocks(dave_best_blocks, &net, &[1]).await;
 
 	// Give Dave all tasks some cpu cycles to burn through their events queues,
-	run_for(Duration::from_millis(100), &net, &mut runtime);
+	run_for(Duration::from_millis(100), &net).await;
 	// then verify Dave catches up through on-demand justification requests.
 	finalize_block_and_wait_for_beefy(
 		&net,
 		[(dave_index, BeefyKeyring::Dave)].into_iter(),
-		&mut runtime,
 		&[6, 10, 17, 24, 26],
 		&[5, 10, 15, 20, 25],
-	);
+	)
+	.await;
 
 	let all_peers = all_peers.into_iter().enumerate();
 	// Now that Dave has caught up, sanity check voting works for all of them.
-	finalize_block_and_wait_for_beefy(&net, all_peers, &mut runtime, &[30], &[30]);
+	finalize_block_and_wait_for_beefy(&net, all_peers, &[30], &[30]).await;
 }
 
-#[test]
-fn should_initialize_voter_at_genesis() {
+#[tokio::test]
+async fn should_initialize_voter_at_genesis() {
 	let keys = &[BeefyKeyring::Alice];
 	let validator_set = ValidatorSet::new(make_beefy_ids(keys), 0).unwrap();
-	let mut runtime = Runtime::new().unwrap();
-	let mut net = BeefyTestNet::new(runtime.handle().clone(), 1);
+	let mut net = BeefyTestNet::new(1);
 	let backend = net.peer(0).client().as_backend();
 
 	// push 15 blocks with `AuthorityChange` digests every 10 blocks
-	net.generate_blocks_and_sync(15, 10, &validator_set, false, &mut runtime);
+	net.generate_blocks_and_sync(15, 10, &validator_set, false).await;
 
 	let mut finality = net.peer(0).client().as_client().finality_notification_stream().fuse();
 
@@ -1033,16 +1000,15 @@ fn should_initialize_voter_at_genesis() {
 	assert_eq!(state, persisted_state);
 }
 
-#[test]
-fn should_initialize_voter_when_last_final_is_session_boundary() {
+#[tokio::test]
+async fn should_initialize_voter_when_last_final_is_session_boundary() {
 	let keys = &[BeefyKeyring::Alice];
 	let validator_set = ValidatorSet::new(make_beefy_ids(keys), 0).unwrap();
-	let mut runtime = Runtime::new().unwrap();
-	let mut net = BeefyTestNet::new(runtime.handle().clone(), 1);
+	let mut net = BeefyTestNet::new(1);
 	let backend = net.peer(0).client().as_backend();
 
 	// push 15 blocks with `AuthorityChange` digests every 10 blocks
-	net.generate_blocks_and_sync(15, 10, &validator_set, false, &mut runtime);
+	net.generate_blocks_and_sync(15, 10, &validator_set, false).await;
 
 	let mut finality = net.peer(0).client().as_client().finality_notification_stream().fuse();
 
@@ -1094,16 +1060,15 @@ fn should_initialize_voter_when_last_final_is_session_boundary() {
 	assert_eq!(state, persisted_state);
 }
 
-#[test]
-fn should_initialize_voter_at_latest_finalized() {
+#[tokio::test]
+async fn should_initialize_voter_at_latest_finalized() {
 	let keys = &[BeefyKeyring::Alice];
 	let validator_set = ValidatorSet::new(make_beefy_ids(keys), 0).unwrap();
-	let mut runtime = Runtime::new().unwrap();
-	let mut net = BeefyTestNet::new(runtime.handle().clone(), 1);
+	let mut net = BeefyTestNet::new(1);
 	let backend = net.peer(0).client().as_backend();
 
 	// push 15 blocks with `AuthorityChange` digests every 10 blocks
-	net.generate_blocks_and_sync(15, 10, &validator_set, false, &mut runtime);
+	net.generate_blocks_and_sync(15, 10, &validator_set, false).await;
 
 	let mut finality = net.peer(0).client().as_client().finality_notification_stream().fuse();
 

substrate/client/beefy/src/worker.rs

@@ -983,7 +983,6 @@ pub(crate) mod tests {
 		runtime::{Block, Digest, DigestItem, Header, H256},
 		Backend,
 	};
-	use tokio::runtime::Runtime;
 
 	impl<B: super::Block> PersistedState<B> {
 		pub fn voting_oracle(&self) -> &VoterOracle<B> {
@@ -1295,12 +1294,11 @@ pub(crate) mod tests {
 		assert_eq!(extracted, Some(validator_set));
 	}
 
-	#[test]
-	fn keystore_vs_validator_set() {
+	#[tokio::test]
+	async fn keystore_vs_validator_set() {
 		let keys = &[Keyring::Alice];
 		let validator_set = ValidatorSet::new(make_beefy_ids(keys), 0).unwrap();
-		let runtime = Runtime::new().unwrap();
-		let mut net = BeefyTestNet::new(runtime.handle().clone(), 1);
+		let mut net = BeefyTestNet::new(1);
 		let mut worker = create_beefy_worker(&net.peer(0), &keys[0], 1, validator_set.clone());
 
 		// keystore doesn't contain other keys than validators'
@@ -1319,12 +1317,11 @@ pub(crate) mod tests {
 		assert_eq!(worker.verify_validator_set(&1, &validator_set), expected_err);
 	}
 
-	#[test]
-	fn should_finalize_correctly() {
+	#[tokio::test]
+	async fn should_finalize_correctly() {
 		let keys = [Keyring::Alice];
 		let validator_set = ValidatorSet::new(make_beefy_ids(&keys), 0).unwrap();
-		let runtime = Runtime::new().unwrap();
-		let mut net = BeefyTestNet::new(runtime.handle().clone(), 1);
+		let mut net = BeefyTestNet::new(1);
 		let backend = net.peer(0).client().as_backend();
 		let mut worker = create_beefy_worker(&net.peer(0), &keys[0], 1, validator_set.clone());
 		// remove default session, will manually add custom one.
@@ -1347,11 +1344,12 @@ pub(crate) mod tests {
 
 		// no 'best beefy block' or finality proofs
 		assert_eq!(worker.best_beefy_block(), 0);
-		runtime.block_on(poll_fn(move |cx| {
+		poll_fn(move |cx| {
 			assert_eq!(best_block_stream.poll_next_unpin(cx), Poll::Pending);
 			assert_eq!(finality_proof.poll_next_unpin(cx), Poll::Pending);
 			Poll::Ready(())
-		}));
+		})
+		.await;
 
 		// unknown hash for block #1
 		let (mut best_block_streams, mut finality_proofs) =
@@ -1368,7 +1366,7 @@ pub(crate) mod tests {
 		worker.finalize(justif.clone()).unwrap();
 		// verify block finalized
 		assert_eq!(worker.best_beefy_block(), 1);
-		runtime.block_on(poll_fn(move |cx| {
+		poll_fn(move |cx| {
 			// unknown hash -> nothing streamed
 			assert_eq!(best_block_stream.poll_next_unpin(cx), Poll::Pending);
 			// commitment streamed
@@ -1378,7 +1376,8 @@ pub(crate) mod tests {
 				v => panic!("unexpected value: {:?}", v),
 			}
 			Poll::Ready(())
-		}));
+		})
+		.await;
 
 		// generate 2 blocks, try again expect success
 		let (mut best_block_streams, _) = get_beefy_streams(&mut net, keys);
@@ -1400,7 +1399,7 @@ pub(crate) mod tests {
 		assert_eq!(worker.active_rounds().unwrap().session_start(), 2);
 		// verify block finalized
 		assert_eq!(worker.best_beefy_block(), 2);
-		runtime.block_on(poll_fn(move |cx| {
+		poll_fn(move |cx| {
 			match best_block_stream.poll_next_unpin(cx) {
 				// expect Some(hash-of-block-2)
 				Poll::Ready(Some(hash)) => {
@@ -1410,19 +1409,19 @@ pub(crate) mod tests {
 				v => panic!("unexpected value: {:?}", v),
 			}
 			Poll::Ready(())
-		}));
+		})
+		.await;
 
 		// check BEEFY justifications are also appended to backend
 		let justifs = backend.blockchain().justifications(hashof2).unwrap().unwrap();
 		assert!(justifs.get(BEEFY_ENGINE_ID).is_some())
 	}
 
-	#[test]
-	fn should_init_session() {
+	#[tokio::test]
+	async fn should_init_session() {
 		let keys = &[Keyring::Alice, Keyring::Bob];
 		let validator_set = ValidatorSet::new(make_beefy_ids(keys), 0).unwrap();
-		let runtime = Runtime::new().unwrap();
-		let mut net = BeefyTestNet::new(runtime.handle().clone(), 1);
+		let mut net = BeefyTestNet::new(1);
 		let mut worker = create_beefy_worker(&net.peer(0), &keys[0], 1, validator_set.clone());
 
 		let worker_rounds = worker.active_rounds().unwrap();
@@ -1449,12 +1448,11 @@ pub(crate) mod tests {
 		assert_eq!(rounds.validator_set_id(), new_validator_set.id());
 	}
 
-	#[test]
-	fn should_triage_votes_and_process_later() {
+	#[tokio::test]
+	async fn should_triage_votes_and_process_later() {
 		let keys = &[Keyring::Alice, Keyring::Bob];
 		let validator_set = ValidatorSet::new(make_beefy_ids(keys), 0).unwrap();
-		let runtime = Runtime::new().unwrap();
-		let mut net = BeefyTestNet::new(runtime.handle().clone(), 1);
+		let mut net = BeefyTestNet::new(1);
 		let mut worker = create_beefy_worker(&net.peer(0), &keys[0], 1, validator_set.clone());
 		// remove default session, will manually add custom one.
 		worker.persisted_state.voting_oracle.sessions.clear();