tests: use future::join instead of future::select (#2813)

* tests/av-store: use future::join instead of future::select

* tests/backing: use future::join instead of future::select

* tests/provisioner: use future::join instead of future::select

* tests/av-dist: use future::join instead of future::select

* tests/av-recovery: use future::join instead of future::select

* tests/bridge: use future::join instead of future::select

* tests/collator-protocol: use future::join instead of future::select

* tests/stmt-dist: use future::join instead of future::select

* fix tests

polkadot/node/network/collator-protocol/src/collator_side.rs

@@ -1139,7 +1139,7 @@ mod tests {
 		virtual_overseer: VirtualOverseer,
 	}
 
-	fn test_harness<T: Future<Output = ()>>(
+	fn test_harness<T: Future<Output = VirtualOverseer>>(
 		local_peer_id: PeerId,
 		collator_pair: CollatorPair,
 		test: impl FnOnce(TestHarness) -> T,
@@ -1155,13 +1155,16 @@ mod tests {
 		futures::pin_mut!(test_fut);
 		futures::pin_mut!(subsystem);
 
-		executor::block_on(future::select(test_fut, subsystem));
+		executor::block_on(future::join(async move {
+			let mut overseer = test_fut.await;
+			overseer_signal(&mut overseer, OverseerSignal::Conclude).await;
+		}, subsystem)).1.unwrap();
 	}
 
 	const TIMEOUT: Duration = Duration::from_millis(100);
 
 	async fn overseer_send(
-		overseer: &mut test_helpers::TestSubsystemContextHandle<CollatorProtocolMessage>,
+		overseer: &mut VirtualOverseer,
 		msg: CollatorProtocolMessage,
 	) {
 		tracing::trace!(?msg, "sending message");
@@ -1173,7 +1176,7 @@ mod tests {
 	}
 
 	async fn overseer_recv(
-		overseer: &mut test_helpers::TestSubsystemContextHandle<CollatorProtocolMessage>,
+		overseer: &mut VirtualOverseer,
 	) -> AllMessages {
 		let msg = overseer_recv_with_timeout(overseer, TIMEOUT)
 			.await
@@ -1185,7 +1188,7 @@ mod tests {
 	}
 
 	async fn overseer_recv_with_timeout(
-		overseer: &mut test_helpers::TestSubsystemContextHandle<CollatorProtocolMessage>,
+		overseer: &mut VirtualOverseer,
 		timeout: Duration,
 	) -> Option<AllMessages> {
 		tracing::trace!("waiting for message...");
@@ -1196,7 +1199,7 @@ mod tests {
 	}
 
 	async fn overseer_signal(
-		overseer: &mut test_helpers::TestSubsystemContextHandle<CollatorProtocolMessage>,
+		overseer: &mut VirtualOverseer,
 		signal: OverseerSignal,
 	) {
 		overseer
@@ -1570,6 +1573,7 @@ mod tests {
 			).await;
 
 			expect_advertise_collation_msg(&mut virtual_overseer, &peer, test_state.relay_parent).await;
+			virtual_overseer
 		});
 	}
 
@@ -1589,6 +1593,7 @@ mod tests {
 			// A validator connected to us
 			connect_peer(&mut virtual_overseer, peer.clone()).await;
 			expect_declare_msg(&mut virtual_overseer, &test_state, &peer).await;
+			virtual_overseer
 		})
 	}
 
@@ -1632,6 +1637,7 @@ mod tests {
 
 			// After changing the view we should receive the advertisement
 			expect_advertise_collation_msg(&mut virtual_overseer, &peer, test_state.relay_parent).await;
+			virtual_overseer
 		})
 	}
 
@@ -1678,7 +1684,9 @@ mod tests {
 			expect_advertise_collation_msg(&mut virtual_overseer, &peer, old_relay_parent).await;
 
 			send_peer_view_change(&mut virtual_overseer, &peer2, vec![test_state.relay_parent]).await;
+
 			expect_advertise_collation_msg(&mut virtual_overseer, &peer2, test_state.relay_parent).await;
+			virtual_overseer
 		})
 	}
 
@@ -1714,6 +1722,7 @@ mod tests {
 			send_peer_view_change(&mut virtual_overseer, &peer, vec![test_state.relay_parent]).await;
 
 			assert!(overseer_recv_with_timeout(&mut virtual_overseer, TIMEOUT).await.is_none());
+			virtual_overseer
 		})
 	}
 
@@ -1756,6 +1765,7 @@ mod tests {
 					PeerSet::Collation,
 				)) if p == peer
 			);
+			virtual_overseer
 		})
 	}
 }

polkadot/node/network/collator-protocol/src/validator_side.rs

@@ -1325,7 +1325,7 @@ mod tests {
 		virtual_overseer: VirtualOverseer,
 	}
 
-	fn test_harness<T: Future<Output = ()>>(test: impl FnOnce(TestHarness) -> T) {
+	fn test_harness<T: Future<Output = VirtualOverseer>>(test: impl FnOnce(TestHarness) -> T) {
 		let _ = env_logger::builder()
 			.is_test(true)
 			.filter(
@@ -1363,7 +1363,10 @@ mod tests {
 		futures::pin_mut!(test_fut);
 		futures::pin_mut!(subsystem);
 
-		executor::block_on(future::select(test_fut, subsystem));
+		executor::block_on(future::join(async move {
+			let mut overseer = test_fut.await;
+			overseer_signal(&mut overseer, OverseerSignal::Conclude).await;
+		}, subsystem)).1.unwrap();
 	}
 
 	const TIMEOUT: Duration = Duration::from_millis(200);
@@ -1403,6 +1406,17 @@ mod tests {
 			.await
 	}
 
+	async fn overseer_signal(
+		overseer: &mut VirtualOverseer,
+		signal: OverseerSignal,
+	) {
+		overseer
+			.send(FromOverseer::Signal(signal))
+			.timeout(TIMEOUT)
+			.await
+			.expect(&format!("{:?} is more than enough for sending signals.", TIMEOUT));
+	}
+
 	async fn respond_to_core_info_queries(
 		virtual_overseer: &mut VirtualOverseer,
 		test_state: &TestState,
@@ -1513,6 +1527,7 @@ mod tests {
 				assert_eq!(para_id, test_state.chain_ids[0]);
 				assert_eq!(collator, pair.public());
 			});
+			virtual_overseer
 		});
 	}
 
@@ -1615,6 +1630,7 @@ mod tests {
 					assert_eq!(rep, BENEFIT_NOTIFY_GOOD);
 				}
 			);
+			virtual_overseer
 		});
 	}
 
@@ -1664,6 +1680,7 @@ mod tests {
 					assert_eq!(rep, COST_INVALID_SIGNATURE);
 				}
 			);
+			virtual_overseer
 		});
 	}
 
@@ -1883,6 +1900,7 @@ mod tests {
 
 			assert_eq!(collation_0.0, candidate_a);
 			assert_eq!(collation_1.0, candidate_b);
+			virtual_overseer
 		});
 	}
 
@@ -1970,7 +1988,8 @@ mod tests {
 					assert_eq!(peer, peer_b);
 					assert_eq!(peer_set, PeerSet::Collation);
 				}
-			)
+			);
+			virtual_overseer
 		});
 	}
 
@@ -2119,7 +2138,8 @@ mod tests {
 					assert_eq!(peer, peer_b);
 					assert_eq!(peer_set, PeerSet::Collation);
 				}
-			)
+			);
+			virtual_overseer
 		});
 	}
 
@@ -2162,7 +2182,8 @@ mod tests {
 					assert_eq!(peer, peer_b);
 					assert_eq!(peer_set, PeerSet::Collation);
 				}
-			)
+			);
+			virtual_overseer
 		})
 	}
 
@@ -2233,6 +2254,7 @@ mod tests {
 					assert_eq!(peer_set, PeerSet::Collation);
 				}
 			);
+			virtual_overseer
 		})
 	}
 
@@ -2304,6 +2326,7 @@ mod tests {
 					assert_eq!(peer_set, PeerSet::Collation);
 				}
 			);
+			virtual_overseer
 		})
 	}
 }