Get rid of runtime.block_on() in Aura tests

client/consensus/aura/src/lib.rs

@@ -658,7 +658,7 @@ mod tests {
 		runtime::{Header, H256},
 		TestClient,
 	};
-	use tokio::runtime::{Handle, Runtime};
+	use tokio::runtime::Handle;
 
 	const SLOT_DURATION_MS: u64 = 1000;
 
@@ -778,11 +778,10 @@ mod tests {
 		}
 	}
 
-	#[test]
-	fn authoring_blocks() {
+	#[tokio::test]
+	async fn authoring_blocks() {
 		sp_tracing::try_init_simple();
-		let runtime = Runtime::new().unwrap();
-		let net = AuraTestNet::new(runtime.handle().clone(), 3);
+		let net = AuraTestNet::new(Handle::current(), 3);
 
 		let peers = &[(0, Keyring::Alice), (1, Keyring::Bob), (2, Keyring::Charlie)];
 
@@ -848,13 +847,14 @@ mod tests {
 			);
 		}
 
-		runtime.block_on(future::select(
+		future::select(
 			future::poll_fn(move |cx| {
 				net.lock().poll(cx);
 				Poll::<()>::Pending
 			}),
 			future::select(future::join_all(aura_futures), future::join_all(import_notifications)),
-		));
+		)
+		.await;
 	}
 
 	#[test]
@@ -873,10 +873,9 @@ mod tests {
 		);
 	}
 
-	#[test]
-	fn current_node_authority_should_claim_slot() {
-		let runtime = Runtime::new().unwrap();
-		let net = AuraTestNet::new(runtime.handle().clone(), 4);
+	#[tokio::test]
+	async fn current_node_authority_should_claim_slot() {
+		let net = AuraTestNet::new(Handle::current(), 4);
 
 		let mut authorities = vec![
 			Keyring::Alice.public().into(),
@@ -920,20 +919,19 @@ mod tests {
 			Default::default(),
 			Default::default(),
 		);
-		assert!(runtime.block_on(worker.claim_slot(&head, 0.into(), &authorities)).is_none());
-		assert!(runtime.block_on(worker.claim_slot(&head, 1.into(), &authorities)).is_none());
-		assert!(runtime.block_on(worker.claim_slot(&head, 2.into(), &authorities)).is_none());
-		assert!(runtime.block_on(worker.claim_slot(&head, 3.into(), &authorities)).is_some());
-		assert!(runtime.block_on(worker.claim_slot(&head, 4.into(), &authorities)).is_none());
-		assert!(runtime.block_on(worker.claim_slot(&head, 5.into(), &authorities)).is_none());
-		assert!(runtime.block_on(worker.claim_slot(&head, 6.into(), &authorities)).is_none());
-		assert!(runtime.block_on(worker.claim_slot(&head, 7.into(), &authorities)).is_some());
+		assert!(worker.claim_slot(&head, 0.into(), &authorities).await.is_none());
+		assert!(worker.claim_slot(&head, 1.into(), &authorities).await.is_none());
+		assert!(worker.claim_slot(&head, 2.into(), &authorities).await.is_none());
+		assert!(worker.claim_slot(&head, 3.into(), &authorities).await.is_some());
+		assert!(worker.claim_slot(&head, 4.into(), &authorities).await.is_none());
+		assert!(worker.claim_slot(&head, 5.into(), &authorities).await.is_none());
+		assert!(worker.claim_slot(&head, 6.into(), &authorities).await.is_none());
+		assert!(worker.claim_slot(&head, 7.into(), &authorities).await.is_some());
 	}
 
-	#[test]
-	fn on_slot_returns_correct_block() {
-		let runtime = Runtime::new().unwrap();
-		let net = AuraTestNet::new(runtime.handle().clone(), 4);
+	#[tokio::test]
+	async fn on_slot_returns_correct_block() {
+		let net = AuraTestNet::new(Handle::current(), 4);
 
 		let keystore_path = tempfile::tempdir().expect("Creates keystore path");
 		let keystore = LocalKeystore::open(keystore_path.path(), None).expect("Creates keystore.");
@@ -969,15 +967,16 @@ mod tests {
 
 		let head = client.header(&BlockId::Number(0)).unwrap().unwrap();
 
-		let res = runtime
-			.block_on(worker.on_slot(SlotInfo {
+		let res = worker
+			.on_slot(SlotInfo {
 				slot: 0.into(),
 				ends_at: Instant::now() + Duration::from_secs(100),
 				create_inherent_data: Box::new(()),
 				duration: Duration::from_millis(1000),
 				chain_head: head,
 				block_size_limit: None,
-			}))
+			})
+			.await
 			.unwrap();
 
 		// The returned block should be imported and we should be able to get its header by now.