From 1e5f19376563b55ef5b3f9aa7a6a7d4c0d6880fe Mon Sep 17 00:00:00 2001
From: Robert Klotzner <eskimor@users.noreply.github.com>
Date: Mon, 10 May 2021 17:16:08 +0200
Subject: [PATCH] Fix flaky test (#3002)

* Reporting peer might come first

before trying to request data, depending on scheduler.

* Better test.
---
 .../network/statement-distribution/src/lib.rs | 61 +++++++++++--------
 1 file changed, 34 insertions(+), 27 deletions(-)

diff --git a/polkadot/node/network/statement-distribution/src/lib.rs b/polkadot/node/network/statement-distribution/src/lib.rs
index 74ba2e978a..ccef05229f 100644
--- a/polkadot/node/network/statement-distribution/src/lib.rs
+++ b/polkadot/node/network/statement-distribution/src/lib.rs
@@ -3663,34 +3663,41 @@ mod tests {
 				}).await;
 			}
 
-			// We should try to fetch the data:
-			assert_matches!(
-				handle.recv().await,
-				AllMessages::NetworkBridge(
-					NetworkBridgeMessage::SendRequests(
-						mut reqs, IfDisconnected::ImmediateError
-					)
-				) => {
-					let reqs = reqs.pop().unwrap();
-					let outgoing = match reqs {
-						Requests::StatementFetching(outgoing) => outgoing,
-						_ => panic!("Unexpected request"),
-					};
-					let req = outgoing.payload;
-					assert_eq!(req.relay_parent, metadata.relay_parent);
-					assert_eq!(req.candidate_hash, metadata.candidate_hash);
-					assert_eq!(outgoing.peer, Recipient::Peer(peer_a));
-					// Just drop request - should trigger error.
-				}
-			);
+			// We should try to fetch the data and punish the peer (but we don't know what comes
+			// first):
+			let mut requested = false;
+			let mut punished = false;
+			for _ in 0..2 {
+				match handle.recv().await {
+					AllMessages::NetworkBridge(
+						NetworkBridgeMessage::SendRequests(
+							mut reqs, IfDisconnected::ImmediateError
+						)
+					) => {
+						let reqs = reqs.pop().unwrap();
+						let outgoing = match reqs {
+							Requests::StatementFetching(outgoing) => outgoing,
+							_ => panic!("Unexpected request"),
+						};
+						let req = outgoing.payload;
+						assert_eq!(req.relay_parent, metadata.relay_parent);
+						assert_eq!(req.candidate_hash, metadata.candidate_hash);
+						assert_eq!(outgoing.peer, Recipient::Peer(peer_a));
+						// Just drop request - should trigger error.
+						requested = true;
+					}
 
-			// Then we should punish peer:
-			assert_matches!(
-				handle.recv().await,
-				AllMessages::NetworkBridge(
-					NetworkBridgeMessage::ReportPeer(p, r)
-				) if p == peer_a && r == COST_APPARENT_FLOOD => {}
-			);
+					AllMessages::NetworkBridge(
+						NetworkBridgeMessage::ReportPeer(p, r)
+					) if p == peer_a && r == COST_APPARENT_FLOOD => {
+						punished = true;
+					}
+
+					m => panic!("Unexpected message: {:?}", m),
+				}
+			}
+			assert!(requested, "large data has not been requested.");
+			assert!(punished, "Peer should have been punished for flooding.");
 
 			handle.send(FromOverseer::Signal(OverseerSignal::Conclude)).await;
 		};