Fix flaky test

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

@@ -85,12 +85,16 @@ const BENEFIT_NOTIFY_GOOD: Rep =
 /// to finish on time.
 ///
 /// There is debug logging output, so we can adjust this value based on production results.
+#[cfg(not(test))]
 const MAX_UNSHARED_DOWNLOAD_TIME: Duration = Duration::from_millis(400);
 
 // How often to check all peers with activity.
 #[cfg(not(test))]
 const ACTIVITY_POLL: Duration = Duration::from_secs(1);
 
+#[cfg(test)]
+const MAX_UNSHARED_DOWNLOAD_TIME: Duration = Duration::from_millis(100);
+
 #[cfg(test)]
 const ACTIVITY_POLL: Duration = Duration::from_millis(10);
 

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

@@ -20,7 +20,7 @@ use futures::{executor, future, Future};
 use sp_core::{crypto::Pair, Encode};
 use sp_keyring::Sr25519Keyring;
 use sp_keystore::{testing::KeyStore as TestKeyStore, SyncCryptoStore};
-use std::{iter, sync::Arc, time::Duration};
+use std::{iter, sync::Arc, task::Poll, time::Duration};
 
 use polkadot_node_network_protocol::{
 	our_view,
@@ -493,17 +493,11 @@ fn collator_authentication_verification_works() {
 	});
 }
 
-// A test scenario that takes the following steps
-//  - Two collators connect, declare themselves and advertise a collation relevant to
-//	our view.
-//	- Collation protocol should request one PoV.
-//	- Collation protocol should disconnect both collators after having received the collation.
-//	- The same collators plus an additional collator connect again and send `PoV`s for a different relay parent.
-//	- Collation protocol will request one PoV, but we will cancel it.
-//	- Collation protocol should request the second PoV which does not succeed in time.
-//	- Collation protocol should request third PoV.
+/// Tests that a validator fetches only one collation at any moment of time
+/// per relay parent and ignores other advertisements once a candidate gets
+/// seconded.
 #[test]
-fn fetch_collations_works() {
+fn fetch_one_collation_at_a_time() {
 	let test_state = TestState::default();
 
 	test_harness(|test_harness| async move {
@@ -575,22 +569,38 @@ fn fetch_collations_works() {
 		)
 		.await;
 
-		overseer_send(
-			&mut virtual_overseer,
-			CollatorProtocolMessage::NetworkBridgeUpdate(NetworkBridgeEvent::PeerDisconnected(
-				peer_b.clone(),
-			)),
-		)
-		.await;
+		// Ensure the subsystem is polled.
+		test_helpers::Yield::new().await;
+
+		// Second collation is not requested since there's already seconded one.
+		assert_matches!(futures::poll!(virtual_overseer.recv().boxed()), Poll::Pending);
+
+		virtual_overseer
+	})
+}
+
+/// Tests that a validator starts fetching next queued collations on [`MAX_UNSHARED_DOWNLOAD_TIME`]
+/// timeout and in case of an error.
+#[test]
+fn fetches_next_collation() {
+	let test_state = TestState::default();
+
+	test_harness(|test_harness| async move {
+		let TestHarness { mut virtual_overseer } = test_harness;
+
+		let second = Hash::random();
 
 		overseer_send(
 			&mut virtual_overseer,
-			CollatorProtocolMessage::NetworkBridgeUpdate(NetworkBridgeEvent::PeerDisconnected(
-				peer_c.clone(),
+			CollatorProtocolMessage::NetworkBridgeUpdate(NetworkBridgeEvent::OurViewChange(
+				our_view![test_state.relay_parent, second],
 			)),
 		)
 		.await;
 
+		respond_to_core_info_queries(&mut virtual_overseer, &test_state).await;
+		respond_to_core_info_queries(&mut virtual_overseer, &test_state).await;
+
 		let peer_b = PeerId::random();
 		let peer_c = PeerId::random();
 		let peer_d = PeerId::random();

node/subsystem-test-helpers/src/lib.rs

@@ -30,6 +30,7 @@ use sp_core::testing::TaskExecutor;
 
 use std::{
 	convert::Infallible,
+	future::Future,
 	pin::Pin,
 	sync::Arc,
 	task::{Context, Poll, Waker},
@@ -391,6 +392,34 @@ macro_rules! arbitrary_order {
 	};
 }
 
+/// Future that yields the execution once and resolves
+/// immediately after.
+///
+/// Useful when one wants to poll the background task to completion
+/// before sending messages to it in order to avoid races.
+pub struct Yield(bool);
+
+impl Yield {
+	/// Returns new `Yield` future.
+	pub fn new() -> Self {
+		Self(false)
+	}
+}
+
+impl Future for Yield {
+	type Output = ();
+
+	fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+		if !self.0 {
+			self.0 = true;
+			cx.waker().wake_by_ref();
+			Poll::Pending
+		} else {
+			Poll::Ready(())
+		}
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	use super::*;