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

substrate/client/network/src/service/tests/chain_sync.rs

@@ -44,7 +44,6 @@ use std::{
 	time::Duration,
 };
 use substrate_test_runtime_client::{TestClientBuilder, TestClientBuilderExt as _};
-use tokio::runtime::Handle;
 
 fn set_default_expecations_no_peers(
 	chain_sync: &mut MockChainSync<substrate_test_runtime_client::runtime::Block>,
@@ -72,7 +71,7 @@ async fn normal_network_poll_no_peers() {
 	let chain_sync_service =
 		Box::new(MockChainSyncInterface::<substrate_test_runtime_client::runtime::Block>::new());
 
-	let mut network = TestNetworkBuilder::new(Handle::current())
+	let mut network = TestNetworkBuilder::new()
 		.with_chain_sync((chain_sync, chain_sync_service))
 		.build();
 
@@ -104,7 +103,7 @@ async fn request_justification() {
 	let mut chain_sync = MockChainSync::<substrate_test_runtime_client::runtime::Block>::new();
 
 	set_default_expecations_no_peers(&mut chain_sync);
-	let mut network = TestNetworkBuilder::new(Handle::current())
+	let mut network = TestNetworkBuilder::new()
 		.with_chain_sync((Box::new(chain_sync), chain_sync_service))
 		.build();
 
@@ -135,7 +134,7 @@ async fn clear_justification_requests() {
 		Box::new(MockChainSync::<substrate_test_runtime_client::runtime::Block>::new());
 
 	set_default_expecations_no_peers(&mut chain_sync);
-	let mut network = TestNetworkBuilder::new(Handle::current())
+	let mut network = TestNetworkBuilder::new()
 		.with_chain_sync((chain_sync, chain_sync_service))
 		.build();
 
@@ -174,7 +173,7 @@ async fn set_sync_fork_request() {
 		.once()
 		.returning(|_, _, _| ());
 
-	let mut network = TestNetworkBuilder::new(Handle::current())
+	let mut network = TestNetworkBuilder::new()
 		.with_chain_sync((chain_sync, Box::new(chain_sync_service)))
 		.build();
 
@@ -218,7 +217,7 @@ async fn on_block_finalized() {
 		.returning(|_, _| ());
 
 	set_default_expecations_no_peers(&mut chain_sync);
-	let mut network = TestNetworkBuilder::new(Handle::current())
+	let mut network = TestNetworkBuilder::new()
 		.with_client(client)
 		.with_chain_sync((chain_sync, chain_sync_service))
 		.build();
@@ -316,7 +315,7 @@ async fn invalid_justification_imported() {
 	let justification_info = Arc::new(RwLock::new(None));
 	let listen_addr = config::build_multiaddr![Memory(rand::random::<u64>())];
 
-	let (service1, mut event_stream1) = TestNetworkBuilder::new(Handle::current())
+	let (service1, mut event_stream1) = TestNetworkBuilder::new()
 		.with_import_queue(Box::new(DummyImportQueue(
 			justification_info.clone(),
 			DummyImportQueueHandle {},
@@ -325,7 +324,7 @@ async fn invalid_justification_imported() {
 		.build()
 		.start_network();
 
-	let (service2, mut event_stream2) = TestNetworkBuilder::new(Handle::current())
+	let (service2, mut event_stream2) = TestNetworkBuilder::new()
 		.with_set_config(SetConfig {
 			reserved_nodes: vec![MultiaddrWithPeerId {
 				multiaddr: listen_addr,
@@ -393,7 +392,7 @@ async fn disconnect_peer_using_chain_sync_handle() {
 	)
 	.unwrap();
 
-	let (node1, mut event_stream1) = TestNetworkBuilder::new(Handle::current())
+	let (node1, mut event_stream1) = TestNetworkBuilder::new()
 		.with_listen_addresses(vec![listen_addr.clone()])
 		.with_chain_sync((Box::new(chain_sync), Box::new(chain_sync_service)))
 		.with_chain_sync_network((chain_sync_network_provider, chain_sync_network_handle))
@@ -401,7 +400,7 @@ async fn disconnect_peer_using_chain_sync_handle() {
 		.build()
 		.start_network();
 
-	let (node2, mut event_stream2) = TestNetworkBuilder::new(Handle::current())
+	let (node2, mut event_stream2) = TestNetworkBuilder::new()
 		.with_set_config(SetConfig {
 			reserved_nodes: vec![MultiaddrWithPeerId {
 				multiaddr: listen_addr,

substrate/client/network/src/service/tests/mod.rs

@@ -44,7 +44,6 @@ use substrate_test_runtime_client::{
 	runtime::{Block as TestBlock, Hash as TestHash},
 	TestClient, TestClientBuilder, TestClientBuilderExt as _,
 };
-use tokio::runtime::Handle;
 
 #[cfg(test)]
 mod chain_sync;
@@ -59,12 +58,11 @@ const PROTOCOL_NAME: &str = "/foo";
 
 struct TestNetwork {
 	network: TestNetworkWorker,
-	rt_handle: Handle,
 }
 
 impl TestNetwork {
-	pub fn new(network: TestNetworkWorker, rt_handle: Handle) -> Self {
-		Self { network, rt_handle }
+	pub fn new(network: TestNetworkWorker) -> Self {
+		Self { network }
 	}
 
 	pub fn service(&self) -> &Arc<TestNetworkService> {
@@ -82,7 +80,7 @@ impl TestNetwork {
 		let service = worker.service().clone();
 		let event_stream = service.event_stream("test");
 
-		self.rt_handle.spawn(async move {
+		tokio::spawn(async move {
 			futures::pin_mut!(worker);
 			let _ = worker.await;
 		});
@@ -100,11 +98,10 @@ struct TestNetworkBuilder {
 	chain_sync: Option<(Box<dyn ChainSyncT<TestBlock>>, Box<dyn ChainSyncInterface<TestBlock>>)>,
 	chain_sync_network: Option<(NetworkServiceProvider, NetworkServiceHandle)>,
 	config: Option<config::NetworkConfiguration>,
-	rt_handle: Handle,
 }
 
 impl TestNetworkBuilder {
-	pub fn new(rt_handle: Handle) -> Self {
+	pub fn new() -> Self {
 		Self {
 			import_queue: None,
 			link: None,
@@ -114,7 +111,6 @@ impl TestNetworkBuilder {
 			chain_sync: None,
 			chain_sync_network: None,
 			config: None,
-			rt_handle,
 		}
 	}
 
@@ -229,21 +225,21 @@ impl TestNetworkBuilder {
 		let block_request_protocol_config = {
 			let (handler, protocol_config) =
 				BlockRequestHandler::new(&protocol_id, None, client.clone(), 50);
-			self.rt_handle.spawn(handler.run().boxed());
+			tokio::spawn(handler.run().boxed());
 			protocol_config
 		};
 
 		let state_request_protocol_config = {
 			let (handler, protocol_config) =
 				StateRequestHandler::new(&protocol_id, None, client.clone(), 50);
-			self.rt_handle.spawn(handler.run().boxed());
+			tokio::spawn(handler.run().boxed());
 			protocol_config
 		};
 
 		let light_client_request_protocol_config = {
 			let (handler, protocol_config) =
 				LightClientRequestHandler::new(&protocol_id, None, client.clone());
-			self.rt_handle.spawn(handler.run().boxed());
+			tokio::spawn(handler.run().boxed());
 			protocol_config
 		};
 
@@ -310,11 +306,6 @@ impl TestNetworkBuilder {
 			.link
 			.unwrap_or(Box::new(sc_network_sync::service::mock::MockChainSyncInterface::new()));
 
-		let handle = self.rt_handle.clone();
-		let executor = move |f| {
-			handle.spawn(f);
-		};
-
 		let worker = NetworkWorker::<
 			substrate_test_runtime_client::runtime::Block,
 			substrate_test_runtime_client::runtime::Hash,
@@ -322,7 +313,9 @@ impl TestNetworkBuilder {
 		>::new(config::Params {
 			block_announce_config,
 			role: config::Role::Full,
-			executor: Box::new(executor),
+			executor: Box::new(|f| {
+				tokio::spawn(f);
+			}),
 			network_config,
 			chain: client.clone(),
 			protocol_id,
@@ -340,10 +333,10 @@ impl TestNetworkBuilder {
 		.unwrap();
 
 		let service = worker.service().clone();
-		self.rt_handle.spawn(async move {
+		tokio::spawn(async move {
 			let _ = chain_sync_network_provider.run(service).await;
 		});
-		self.rt_handle.spawn(async move {
+		tokio::spawn(async move {
 			loop {
 				futures::future::poll_fn(|cx| {
 					import_queue.poll_actions(cx, &mut *link);
@@ -354,6 +347,6 @@ impl TestNetworkBuilder {
 			}
 		});
 
-		TestNetwork::new(worker, self.rt_handle)
+		TestNetwork::new(worker)
 	}
 }

substrate/client/network/src/service/tests/service.rs

@@ -26,7 +26,6 @@ use sc_network_common::{
 	service::{NetworkNotification, NetworkPeers, NetworkStateInfo},
 };
 use std::{sync::Arc, time::Duration};
-use tokio::runtime::Handle;
 
 type TestNetworkService = NetworkService<
 	substrate_test_runtime_client::runtime::Block,
@@ -38,9 +37,7 @@ const PROTOCOL_NAME: &str = "/foo";
 
 /// Builds two nodes and their associated events stream.
 /// The nodes are connected together and have the `PROTOCOL_NAME` protocol registered.
-fn build_nodes_one_proto(
-	rt_handle: &Handle,
-) -> (
+fn build_nodes_one_proto() -> (
 	Arc<TestNetworkService>,
 	impl Stream<Item = Event>,
 	Arc<TestNetworkService>,
@@ -48,12 +45,12 @@ fn build_nodes_one_proto(
 ) {
 	let listen_addr = config::build_multiaddr![Memory(rand::random::<u64>())];
 
-	let (node1, events_stream1) = TestNetworkBuilder::new(rt_handle.clone())
+	let (node1, events_stream1) = TestNetworkBuilder::new()
 		.with_listen_addresses(vec![listen_addr.clone()])
 		.build()
 		.start_network();
 
-	let (node2, events_stream2) = TestNetworkBuilder::new(rt_handle.clone())
+	let (node2, events_stream2) = TestNetworkBuilder::new()
 		.with_set_config(SetConfig {
 			reserved_nodes: vec![MultiaddrWithPeerId {
 				multiaddr: listen_addr,
@@ -67,15 +64,12 @@ fn build_nodes_one_proto(
 	(node1, events_stream1, node2, events_stream2)
 }
 
-#[test]
-fn notifications_state_consistent() {
+#[tokio::test]
+async fn notifications_state_consistent() {
 	// Runs two nodes and ensures that events are propagated out of the API in a consistent
 	// correct order, which means no notification received on a closed substream.
 
-	let runtime = tokio::runtime::Runtime::new().unwrap();
-
-	let (node1, mut events_stream1, node2, mut events_stream2) =
-		build_nodes_one_proto(runtime.handle());
+	let (node1, mut events_stream1, node2, mut events_stream2) = build_nodes_one_proto();
 
 	// Write some initial notifications that shouldn't get through.
 	for _ in 0..(rand::random::<u8>() % 5) {
@@ -93,140 +87,130 @@ fn notifications_state_consistent() {
 		);
 	}
 
-	runtime.block_on(async move {
-		// True if we have an active substream from node1 to node2.
-		let mut node1_to_node2_open = false;
-		// True if we have an active substream from node2 to node1.
-		let mut node2_to_node1_open = false;
-		// We stop the test after a certain number of iterations.
-		let mut iterations = 0;
-		// Safe guard because we don't want the test to pass if no substream has been open.
-		let mut something_happened = false;
+	// True if we have an active substream from node1 to node2.
+	let mut node1_to_node2_open = false;
+	// True if we have an active substream from node2 to node1.
+	let mut node2_to_node1_open = false;
+	// We stop the test after a certain number of iterations.
+	let mut iterations = 0;
+	// Safe guard because we don't want the test to pass if no substream has been open.
+	let mut something_happened = false;
 
-		loop {
-			iterations += 1;
-			if iterations >= 1_000 {
-				assert!(something_happened);
-				break
-			}
-
-			// Start by sending a notification from node1 to node2 and vice-versa. Part of the
-			// test consists in ensuring that notifications get ignored if the stream isn't open.
-			if rand::random::<u8>() % 5 >= 3 {
-				node1.write_notification(
-					node2.local_peer_id(),
-					PROTOCOL_NAME.into(),
-					b"hello world".to_vec(),
-				);
-			}
-			if rand::random::<u8>() % 5 >= 3 {
-				node2.write_notification(
-					node1.local_peer_id(),
-					PROTOCOL_NAME.into(),
-					b"hello world".to_vec(),
-				);
-			}
-
-			// Also randomly disconnect the two nodes from time to time.
-			if rand::random::<u8>() % 20 == 0 {
-				node1.disconnect_peer(node2.local_peer_id(), PROTOCOL_NAME.into());
-			}
-			if rand::random::<u8>() % 20 == 0 {
-				node2.disconnect_peer(node1.local_peer_id(), PROTOCOL_NAME.into());
-			}
-
-			// Grab next event from either `events_stream1` or `events_stream2`.
-			let next_event = {
-				let next1 = events_stream1.next();
-				let next2 = events_stream2.next();
-				// We also await on a small timer, otherwise it is possible for the test to wait
-				// forever while nothing at all happens on the network.
-				let continue_test = futures_timer::Delay::new(Duration::from_millis(20));
-				match future::select(future::select(next1, next2), continue_test).await {
-					future::Either::Left((future::Either::Left((Some(ev), _)), _)) =>
-						future::Either::Left(ev),
-					future::Either::Left((future::Either::Right((Some(ev), _)), _)) =>
-						future::Either::Right(ev),
-					future::Either::Right(_) => continue,
-					_ => break,
-				}
-			};
-
-			match next_event {
-				future::Either::Left(Event::NotificationStreamOpened {
-					remote, protocol, ..
-				}) =>
-					if protocol == PROTOCOL_NAME.into() {
-						something_happened = true;
-						assert!(!node1_to_node2_open);
-						node1_to_node2_open = true;
-						assert_eq!(remote, node2.local_peer_id());
-					},
-				future::Either::Right(Event::NotificationStreamOpened {
-					remote, protocol, ..
-				}) =>
-					if protocol == PROTOCOL_NAME.into() {
-						something_happened = true;
-						assert!(!node2_to_node1_open);
-						node2_to_node1_open = true;
-						assert_eq!(remote, node1.local_peer_id());
-					},
-				future::Either::Left(Event::NotificationStreamClosed {
-					remote, protocol, ..
-				}) =>
-					if protocol == PROTOCOL_NAME.into() {
-						assert!(node1_to_node2_open);
-						node1_to_node2_open = false;
-						assert_eq!(remote, node2.local_peer_id());
-					},
-				future::Either::Right(Event::NotificationStreamClosed {
-					remote, protocol, ..
-				}) =>
-					if protocol == PROTOCOL_NAME.into() {
-						assert!(node2_to_node1_open);
-						node2_to_node1_open = false;
-						assert_eq!(remote, node1.local_peer_id());
-					},
-				future::Either::Left(Event::NotificationsReceived { remote, .. }) => {
-					assert!(node1_to_node2_open);
-					assert_eq!(remote, node2.local_peer_id());
-					if rand::random::<u8>() % 5 >= 4 {
-						node1.write_notification(
-							node2.local_peer_id(),
-							PROTOCOL_NAME.into(),
-							b"hello world".to_vec(),
-						);
-					}
-				},
-				future::Either::Right(Event::NotificationsReceived { remote, .. }) => {
-					assert!(node2_to_node1_open);
-					assert_eq!(remote, node1.local_peer_id());
-					if rand::random::<u8>() % 5 >= 4 {
-						node2.write_notification(
-							node1.local_peer_id(),
-							PROTOCOL_NAME.into(),
-							b"hello world".to_vec(),
-						);
-					}
-				},
-
-				// Add new events here.
-				future::Either::Left(Event::SyncConnected { .. }) => {},
-				future::Either::Right(Event::SyncConnected { .. }) => {},
-				future::Either::Left(Event::SyncDisconnected { .. }) => {},
-				future::Either::Right(Event::SyncDisconnected { .. }) => {},
-				future::Either::Left(Event::Dht(_)) => {},
-				future::Either::Right(Event::Dht(_)) => {},
-			};
+	loop {
+		iterations += 1;
+		if iterations >= 1_000 {
+			assert!(something_happened);
+			break
 		}
-	});
+
+		// Start by sending a notification from node1 to node2 and vice-versa. Part of the
+		// test consists in ensuring that notifications get ignored if the stream isn't open.
+		if rand::random::<u8>() % 5 >= 3 {
+			node1.write_notification(
+				node2.local_peer_id(),
+				PROTOCOL_NAME.into(),
+				b"hello world".to_vec(),
+			);
+		}
+		if rand::random::<u8>() % 5 >= 3 {
+			node2.write_notification(
+				node1.local_peer_id(),
+				PROTOCOL_NAME.into(),
+				b"hello world".to_vec(),
+			);
+		}
+
+		// Also randomly disconnect the two nodes from time to time.
+		if rand::random::<u8>() % 20 == 0 {
+			node1.disconnect_peer(node2.local_peer_id(), PROTOCOL_NAME.into());
+		}
+		if rand::random::<u8>() % 20 == 0 {
+			node2.disconnect_peer(node1.local_peer_id(), PROTOCOL_NAME.into());
+		}
+
+		// Grab next event from either `events_stream1` or `events_stream2`.
+		let next_event = {
+			let next1 = events_stream1.next();
+			let next2 = events_stream2.next();
+			// We also await on a small timer, otherwise it is possible for the test to wait
+			// forever while nothing at all happens on the network.
+			let continue_test = futures_timer::Delay::new(Duration::from_millis(20));
+			match future::select(future::select(next1, next2), continue_test).await {
+				future::Either::Left((future::Either::Left((Some(ev), _)), _)) =>
+					future::Either::Left(ev),
+				future::Either::Left((future::Either::Right((Some(ev), _)), _)) =>
+					future::Either::Right(ev),
+				future::Either::Right(_) => continue,
+				_ => break,
+			}
+		};
+
+		match next_event {
+			future::Either::Left(Event::NotificationStreamOpened { remote, protocol, .. }) =>
+				if protocol == PROTOCOL_NAME.into() {
+					something_happened = true;
+					assert!(!node1_to_node2_open);
+					node1_to_node2_open = true;
+					assert_eq!(remote, node2.local_peer_id());
+				},
+			future::Either::Right(Event::NotificationStreamOpened { remote, protocol, .. }) =>
+				if protocol == PROTOCOL_NAME.into() {
+					something_happened = true;
+					assert!(!node2_to_node1_open);
+					node2_to_node1_open = true;
+					assert_eq!(remote, node1.local_peer_id());
+				},
+			future::Either::Left(Event::NotificationStreamClosed { remote, protocol, .. }) =>
+				if protocol == PROTOCOL_NAME.into() {
+					assert!(node1_to_node2_open);
+					node1_to_node2_open = false;
+					assert_eq!(remote, node2.local_peer_id());
+				},
+			future::Either::Right(Event::NotificationStreamClosed { remote, protocol, .. }) =>
+				if protocol == PROTOCOL_NAME.into() {
+					assert!(node2_to_node1_open);
+					node2_to_node1_open = false;
+					assert_eq!(remote, node1.local_peer_id());
+				},
+			future::Either::Left(Event::NotificationsReceived { remote, .. }) => {
+				assert!(node1_to_node2_open);
+				assert_eq!(remote, node2.local_peer_id());
+				if rand::random::<u8>() % 5 >= 4 {
+					node1.write_notification(
+						node2.local_peer_id(),
+						PROTOCOL_NAME.into(),
+						b"hello world".to_vec(),
+					);
+				}
+			},
+			future::Either::Right(Event::NotificationsReceived { remote, .. }) => {
+				assert!(node2_to_node1_open);
+				assert_eq!(remote, node1.local_peer_id());
+				if rand::random::<u8>() % 5 >= 4 {
+					node2.write_notification(
+						node1.local_peer_id(),
+						PROTOCOL_NAME.into(),
+						b"hello world".to_vec(),
+					);
+				}
+			},
+
+			// Add new events here.
+			future::Either::Left(Event::SyncConnected { .. }) => {},
+			future::Either::Right(Event::SyncConnected { .. }) => {},
+			future::Either::Left(Event::SyncDisconnected { .. }) => {},
+			future::Either::Right(Event::SyncDisconnected { .. }) => {},
+			future::Either::Left(Event::Dht(_)) => {},
+			future::Either::Right(Event::Dht(_)) => {},
+		};
+	}
 }
 
 #[tokio::test]
 async fn lots_of_incoming_peers_works() {
 	let listen_addr = config::build_multiaddr![Memory(rand::random::<u64>())];
 
-	let (main_node, _) = TestNetworkBuilder::new(Handle::current())
+	let (main_node, _) = TestNetworkBuilder::new()
 		.with_listen_addresses(vec![listen_addr.clone()])
 		.with_set_config(SetConfig { in_peers: u32::MAX, ..Default::default() })
 		.build()
@@ -239,7 +223,7 @@ async fn lots_of_incoming_peers_works() {
 	let mut background_tasks_to_wait = Vec::new();
 
 	for _ in 0..32 {
-		let (_dialing_node, event_stream) = TestNetworkBuilder::new(Handle::current())
+		let (_dialing_node, event_stream) = TestNetworkBuilder::new()
 			.with_set_config(SetConfig {
 				reserved_nodes: vec![MultiaddrWithPeerId {
 					multiaddr: listen_addr.clone(),
@@ -286,20 +270,17 @@ async fn lots_of_incoming_peers_works() {
 	future::join_all(background_tasks_to_wait).await;
 }
 
-#[test]
-fn notifications_back_pressure() {
+#[tokio::test]
+async fn notifications_back_pressure() {
 	// Node 1 floods node 2 with notifications. Random sleeps are done on node 2 to simulate the
 	// node being busy. We make sure that all notifications are received.
 
 	const TOTAL_NOTIFS: usize = 10_000;
 
-	let runtime = tokio::runtime::Runtime::new().unwrap();
-
-	let (node1, mut events_stream1, node2, mut events_stream2) =
-		build_nodes_one_proto(runtime.handle());
+	let (node1, mut events_stream1, node2, mut events_stream2) = build_nodes_one_proto();
 	let node2_id = node2.local_peer_id();
 
-	let receiver = runtime.spawn(async move {
+	let receiver = tokio::spawn(async move {
 		let mut received_notifications = 0;
 
 		while received_notifications < TOTAL_NOTIFS {
@@ -320,40 +301,36 @@ fn notifications_back_pressure() {
 		}
 	});
 
-	runtime.block_on(async move {
-		// Wait for the `NotificationStreamOpened`.
-		loop {
-			match events_stream1.next().await.unwrap() {
-				Event::NotificationStreamOpened { .. } => break,
-				_ => {},
-			};
-		}
+	// Wait for the `NotificationStreamOpened`.
+	loop {
+		match events_stream1.next().await.unwrap() {
+			Event::NotificationStreamOpened { .. } => break,
+			_ => {},
+		};
+	}
 
-		// Sending!
-		for num in 0..TOTAL_NOTIFS {
-			let notif = node1.notification_sender(node2_id, PROTOCOL_NAME.into()).unwrap();
-			notif
-				.ready()
-				.await
-				.unwrap()
-				.send(format!("hello #{}", num).into_bytes())
-				.unwrap();
-		}
+	// Sending!
+	for num in 0..TOTAL_NOTIFS {
+		let notif = node1.notification_sender(node2_id, PROTOCOL_NAME.into()).unwrap();
+		notif
+			.ready()
+			.await
+			.unwrap()
+			.send(format!("hello #{}", num).into_bytes())
+			.unwrap();
+	}
 
-		receiver.await.unwrap();
-	});
+	receiver.await.unwrap();
 }
 
-#[test]
-fn fallback_name_working() {
+#[tokio::test]
+async fn fallback_name_working() {
 	// Node 1 supports the protocols "new" and "old". Node 2 only supports "old". Checks whether
 	// they can connect.
 	const NEW_PROTOCOL_NAME: &str = "/new-shiny-protocol-that-isnt-PROTOCOL_NAME";
 
-	let runtime = tokio::runtime::Runtime::new().unwrap();
-
 	let listen_addr = config::build_multiaddr![Memory(rand::random::<u64>())];
-	let (node1, mut events_stream1) = TestNetworkBuilder::new(runtime.handle().clone())
+	let (node1, mut events_stream1) = TestNetworkBuilder::new()
 		.with_config(config::NetworkConfiguration {
 			extra_sets: vec![NonDefaultSetConfig {
 				notifications_protocol: NEW_PROTOCOL_NAME.into(),
@@ -369,7 +346,7 @@ fn fallback_name_working() {
 		.build()
 		.start_network();
 
-	let (_, mut events_stream2) = TestNetworkBuilder::new(runtime.handle().clone())
+	let (_, mut events_stream2) = TestNetworkBuilder::new()
 		.with_set_config(SetConfig {
 			reserved_nodes: vec![MultiaddrWithPeerId {
 				multiaddr: listen_addr,
@@ -380,7 +357,7 @@ fn fallback_name_working() {
 		.build()
 		.start_network();
 
-	let receiver = runtime.spawn(async move {
+	let receiver = tokio::spawn(async move {
 		// Wait for the `NotificationStreamOpened`.
 		loop {
 			match events_stream2.next().await.unwrap() {
@@ -394,30 +371,27 @@ fn fallback_name_working() {
 		}
 	});
 
-	runtime.block_on(async move {
-		// Wait for the `NotificationStreamOpened`.
-		loop {
-			match events_stream1.next().await.unwrap() {
-				Event::NotificationStreamOpened { protocol, negotiated_fallback, .. }
-					if protocol == NEW_PROTOCOL_NAME.into() =>
-				{
-					assert_eq!(negotiated_fallback, Some(PROTOCOL_NAME.into()));
-					break
-				},
-				_ => {},
-			};
-		}
+	// Wait for the `NotificationStreamOpened`.
+	loop {
+		match events_stream1.next().await.unwrap() {
+			Event::NotificationStreamOpened { protocol, negotiated_fallback, .. }
+				if protocol == NEW_PROTOCOL_NAME.into() =>
+			{
+				assert_eq!(negotiated_fallback, Some(PROTOCOL_NAME.into()));
+				break
+			},
+			_ => {},
+		};
+	}
 
-		receiver.await.unwrap();
-	});
+	receiver.await.unwrap();
 }
 
 // Disconnect peer by calling `Protocol::disconnect_peer()` with the supplied block announcement
 // protocol name and verify that `SyncDisconnected` event is emitted
 #[tokio::test]
 async fn disconnect_sync_peer_using_block_announcement_protocol_name() {
-	let (node1, mut events_stream1, node2, mut events_stream2) =
-		build_nodes_one_proto(&Handle::current());
+	let (node1, mut events_stream1, node2, mut events_stream2) = build_nodes_one_proto();
 
 	async fn wait_for_events(stream: &mut (impl Stream<Item = Event> + std::marker::Unpin)) {
 		let mut notif_received = false;
@@ -454,7 +428,7 @@ async fn disconnect_sync_peer_using_block_announcement_protocol_name() {
 async fn ensure_listen_addresses_consistent_with_transport_memory() {
 	let listen_addr = config::build_multiaddr![Ip4([127, 0, 0, 1]), Tcp(0_u16)];
 
-	let _ = TestNetworkBuilder::new(Handle::current())
+	let _ = TestNetworkBuilder::new()
 		.with_config(config::NetworkConfiguration {
 			listen_addresses: vec![listen_addr.clone()],
 			transport: TransportConfig::MemoryOnly,
@@ -474,7 +448,7 @@ async fn ensure_listen_addresses_consistent_with_transport_memory() {
 async fn ensure_listen_addresses_consistent_with_transport_not_memory() {
 	let listen_addr = config::build_multiaddr![Memory(rand::random::<u64>())];
 
-	let _ = TestNetworkBuilder::new(Handle::current())
+	let _ = TestNetworkBuilder::new()
 		.with_config(config::NetworkConfiguration {
 			listen_addresses: vec![listen_addr.clone()],
 			..config::NetworkConfiguration::new(
@@ -497,7 +471,7 @@ async fn ensure_boot_node_addresses_consistent_with_transport_memory() {
 		peer_id: PeerId::random(),
 	};
 
-	let _ = TestNetworkBuilder::new(Handle::current())
+	let _ = TestNetworkBuilder::new()
 		.with_config(config::NetworkConfiguration {
 			listen_addresses: vec![listen_addr.clone()],
 			transport: TransportConfig::MemoryOnly,
@@ -522,7 +496,7 @@ async fn ensure_boot_node_addresses_consistent_with_transport_not_memory() {
 		peer_id: PeerId::random(),
 	};
 
-	let _ = TestNetworkBuilder::new(Handle::current())
+	let _ = TestNetworkBuilder::new()
 		.with_config(config::NetworkConfiguration {
 			listen_addresses: vec![listen_addr.clone()],
 			boot_nodes: vec![boot_node],
@@ -546,7 +520,7 @@ async fn ensure_reserved_node_addresses_consistent_with_transport_memory() {
 		peer_id: PeerId::random(),
 	};
 
-	let _ = TestNetworkBuilder::new(Handle::current())
+	let _ = TestNetworkBuilder::new()
 		.with_config(config::NetworkConfiguration {
 			listen_addresses: vec![listen_addr.clone()],
 			transport: TransportConfig::MemoryOnly,
@@ -574,7 +548,7 @@ async fn ensure_reserved_node_addresses_consistent_with_transport_not_memory() {
 		peer_id: PeerId::random(),
 	};
 
-	let _ = TestNetworkBuilder::new(Handle::current())
+	let _ = TestNetworkBuilder::new()
 		.with_config(config::NetworkConfiguration {
 			listen_addresses: vec![listen_addr.clone()],
 			default_peers_set: SetConfig {
@@ -598,7 +572,7 @@ async fn ensure_public_addresses_consistent_with_transport_memory() {
 	let listen_addr = config::build_multiaddr![Memory(rand::random::<u64>())];
 	let public_address = config::build_multiaddr![Ip4([127, 0, 0, 1]), Tcp(0_u16)];
 
-	let _ = TestNetworkBuilder::new(Handle::current())
+	let _ = TestNetworkBuilder::new()
 		.with_config(config::NetworkConfiguration {
 			listen_addresses: vec![listen_addr.clone()],
 			transport: TransportConfig::MemoryOnly,
@@ -620,7 +594,7 @@ async fn ensure_public_addresses_consistent_with_transport_not_memory() {
 	let listen_addr = config::build_multiaddr![Ip4([127, 0, 0, 1]), Tcp(0_u16)];
 	let public_address = config::build_multiaddr![Memory(rand::random::<u64>())];
 
-	let _ = TestNetworkBuilder::new(Handle::current())
+	let _ = TestNetworkBuilder::new()
 		.with_config(config::NetworkConfiguration {
 			listen_addresses: vec![listen_addr.clone()],
 			public_addresses: vec![public_address],