From 6072b90096102767076b24da840d1e1ccf1baabc Mon Sep 17 00:00:00 2001
From: Aaro Altonen <48052676+altonen@users.noreply.github.com>
Date: Tue, 11 Oct 2022 18:25:12 +0300
Subject: [PATCH] Fix flaky service test (#12472)

Sometimes `NotificationStreamOpenened` would be received for the
other protocol before `SyncConnected` was received so when the
connection was closed, an incorrect event was read from the event
stream.
---
 client/network/common/src/config.rs |  4 ++++
 client/network/src/service/tests.rs | 26 ++++++++++++++------------
 2 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/client/network/common/src/config.rs b/client/network/common/src/config.rs
index e4a7f04c8d6e8..96c7c11ec2696 100644
--- a/client/network/common/src/config.rs
+++ b/client/network/common/src/config.rs
@@ -244,6 +244,10 @@ pub struct NonDefaultSetConfig {
 	/// `sc_network::protocol::event::Event::NotificationStreamOpened::negotiated_fallback`
 	pub fallback_names: Vec<protocol::ProtocolName>,
 	/// Handshake of the protocol
+	///
+	/// NOTE: Currently custom handshakes are not fully supported. See issue #5685 for more
+	/// details. This field is temporarily used to allow moving the hardcoded block announcement
+	/// protocol out of `protocol.rs`.
 	pub handshake: Option<NotificationHandshake>,
 	/// Maximum allowed size of single notifications.
 	pub max_notification_size: u64,
diff --git a/client/network/src/service/tests.rs b/client/network/src/service/tests.rs
index 7c651c675b83e..b656d7a7c0ddc 100644
--- a/client/network/src/service/tests.rs
+++ b/client/network/src/service/tests.rs
@@ -637,19 +637,21 @@ async fn disconnect_sync_peer_using_block_announcement_protocol_name() {
 		..config::NetworkConfiguration::new_local()
 	});
 
-	loop {
-		match events_stream1.next().await.unwrap() {
-			Event::NotificationStreamOpened { .. } => break,
-			_ => {},
-		};
+	async fn wait_for_events(stream: &mut (impl Stream<Item = Event> + std::marker::Unpin)) {
+		let mut notif_received = false;
+		let mut sync_received = false;
+
+		while !notif_received || !sync_received {
+			match stream.next().await.unwrap() {
+				Event::NotificationStreamOpened { .. } => notif_received = true,
+				Event::SyncConnected { .. } => sync_received = true,
+				_ => {},
+			};
+		}
 	}
 
-	loop {
-		match events_stream2.next().await.unwrap() {
-			Event::NotificationStreamOpened { .. } => break,
-			_ => {},
-		};
-	}
+	wait_for_events(&mut events_stream1).await;
+	wait_for_events(&mut events_stream2).await;
 
 	// disconnect peer using `PROTOCOL_NAME`, verify `NotificationStreamClosed` event is emitted
 	node2.disconnect_peer(node1.local_peer_id(), PROTOCOL_NAME.into());
@@ -659,7 +661,7 @@ async fn disconnect_sync_peer_using_block_announcement_protocol_name() {
 	));
 	let _ = events_stream2.next().await; // ignore the reopen event
 
-	// now disconnect using the block announcement protocol, verify that `SyncDisconnected` is
+	// now disconnect using `BLOCK_ANNOUNCE_PROTO_NAME`, verify that `SyncDisconnected` is
 	// emitted
 	node2.disconnect_peer(node1.local_peer_id(), BLOCK_ANNOUNCE_PROTO_NAME.into());
 	assert!(std::matches!(events_stream2.next().await, Some(Event::SyncDisconnected { .. })));