Make NetworkWorker::run() consume self

client/network/src/service.rs

@@ -1360,7 +1360,8 @@ where
 	/// Perform one action on the network.
 	///
 	/// Returns `false` when the worker should be shutdown.
-	async fn next_action(&mut self) -> bool {
+	/// Use in tests only.
+	pub async fn next_action(&mut self) -> bool {
 		futures::select! {
 			// Next message from the service.
 			msg = self.from_service.next() => {

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

@@ -75,8 +75,8 @@ async fn normal_network_poll_no_peers() {
 		.with_chain_sync((chain_sync, chain_sync_service))
 		.build();
 
-	// poll the network once
-	let _ = network.network().run().now_or_never();
+	// perform one action on network
+	let _ = network.network().next_action();
 }
 
 #[tokio::test]
@@ -106,7 +106,8 @@ async fn request_justification() {
 	// send "request justifiction" message and poll the network
 	network.service().request_justification(&hash, number);
 
-	let _ = network.network().run().now_or_never();
+	// perform one action on network
+	let _ = network.network().next_action();
 }
 
 #[tokio::test]
@@ -133,7 +134,8 @@ async fn clear_justification_requests() {
 	// send "request justifiction" message and poll the network
 	network.service().clear_justification_requests();
 
-	let _ = network.network().run().now_or_never();
+	// perform one action on network
+	let _ = network.network().next_action();
 }
 
 #[tokio::test]
@@ -168,7 +170,8 @@ async fn set_sync_fork_request() {
 	// send "set sync fork request" message and poll the network
 	network.service().set_sync_fork_request(copy_peers, hash, number);
 
-	let _ = network.network().run().now_or_never();
+	// perform one action on network
+	let _ = network.network().next_action();
 }
 
 #[tokio::test]
@@ -209,7 +212,8 @@ async fn on_block_finalized() {
 	// send "set sync fork request" message and poll the network
 	network.network().on_block_finalized(hash, header);
 
-	let _ = network.network().run().now_or_never();
+	// perform one action on network
+	let _ = network.network().next_action();
 }
 
 // report from mock import queue that importing a justification was not successful

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

@@ -76,13 +76,13 @@ impl TestNetwork {
 	pub fn start_network(
 		self,
 	) -> (Arc<TestNetworkService>, (impl Stream<Item = Event> + std::marker::Unpin)) {
-		let mut worker = self.network;
+		let worker = self.network;
 		let service = worker.service().clone();
 		let event_stream = service.event_stream("test");
 
-		tokio::spawn(
-			worker.run()
-		);
+		tokio::spawn(async move {
+			worker.run().await;
+		});
 
 		(service, event_stream)
 	}

client/network/test/src/lib.rs

@@ -1065,7 +1065,7 @@ where
 			for (i, peer) in peers.iter_mut().enumerate() {
 				trace!(target: "sync", "-- Polling {}: {}", i, peer.id());
 				{
-					let net_poll_future = peer.network.run();
+					let net_poll_future = peer.network.next_action();
 					pin_mut!(net_poll_future);
 					let _ = net_poll_future.poll(cx);
 				}

client/service/src/lib.rs

@@ -153,7 +153,7 @@ async fn build_network_future<
 		+ 'static,
 	H: sc_network_common::ExHashT,
 >(
-	mut network: sc_network::NetworkWorker<B, H, C>,
+	network: sc_network::NetworkWorker<B, H, C>,
 	client: Arc<C>,
 	announce_imported_blocks: bool,
 ) {