From 857e6352be56608e90cbe09621da6c0b4652b213 Mon Sep 17 00:00:00 2001
From: Chris Sosnin <48099298+slumber@users.noreply.github.com>
Date: Mon, 10 Oct 2022 10:06:44 +0400
Subject: [PATCH] Fix flaky test (#6131)

* Split test + decrease test timeout

* fmt

* spellcheck
---
 .../src/validator_side/mod.rs                 |  4 ++
 .../src/validator_side/tests.rs               | 50 +++++++++++--------
 node/subsystem-test-helpers/src/lib.rs        | 29 +++++++++++
 3 files changed, 63 insertions(+), 20 deletions(-)

diff --git a/node/network/collator-protocol/src/validator_side/mod.rs b/node/network/collator-protocol/src/validator_side/mod.rs
index b74c1d5b5a4f..b2b3dc4824b5 100644
--- a/node/network/collator-protocol/src/validator_side/mod.rs
+++ b/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);
 
diff --git a/node/network/collator-protocol/src/validator_side/tests.rs b/node/network/collator-protocol/src/validator_side/tests.rs
index 15740e5d5efa..ae8644cea521 100644
--- a/node/network/collator-protocol/src/validator_side/tests.rs
+++ b/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();
diff --git a/node/subsystem-test-helpers/src/lib.rs b/node/subsystem-test-helpers/src/lib.rs
index e2e61c2006d8..79f833b7558c 100644
--- a/node/subsystem-test-helpers/src/lib.rs
+++ b/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::*;