From 239a62c7640681290c4b2737c6067cf49aba828b Mon Sep 17 00:00:00 2001
From: Max Inden <mail@max-inden.de>
Date: Tue, 14 Feb 2023 21:51:27 +0100
Subject: [PATCH] test(swarm): Wait for swarm1 to disconnect (#3465)

This commit does two things:

- Check that swarm1, given that it has no ban, establishes 21 connections and swarm2, given that it has a ban, establishes 20 connections.

- Wait for swarm1 to notice disconnect, given that swarm2 closed the connection to the banned swarm1.

Fixes flake introduced in https://github.com/libp2p/rust-libp2p/commit/caed1fe2c717ba1688a4eb0549284cddba8c9ea6.
---
 swarm/src/lib.rs  | 35 ++++++++++++++++++++++-------------
 swarm/src/test.rs |  4 ++--
 2 files changed, 24 insertions(+), 15 deletions(-)

diff --git a/swarm/src/lib.rs b/swarm/src/lib.rs
index f2cacec67ff..4adeb0b9794 100644
--- a/swarm/src/lib.rs
+++ b/swarm/src/lib.rs
@@ -1891,8 +1891,10 @@ mod tests {
     /// calls should be registered.
     #[test]
     fn test_connect_disconnect_ban() {
+        let _ = env_logger::try_init();
+
         // Since the test does not try to open any substreams, we can
-        // use the dummy protocols handler.
+        // use keep alive protocols handler.
         let handler_proto = keep_alive::ConnectionHandler;
 
         let mut swarm1 = new_test_swarm::<_, ()>(handler_proto.clone()).build();
@@ -1906,6 +1908,7 @@ mod tests {
 
         let swarm1_id = *swarm1.local_peer_id();
 
+        #[derive(Debug)]
         enum Stage {
             /// Waiting for the peers to connect. Banning has not occurred.
             Connecting,
@@ -1950,25 +1953,31 @@ mod tests {
                     {
                         // Setup to test that new connections of banned peers are not reported.
                         swarm1.dial(addr2.clone()).unwrap();
+                        s1_expected_conns += 1;
                         stage = Stage::BannedDial;
                     }
                 }
                 Stage::BannedDial => {
-                    // The banned connection was established. Check that it was not reported to
-                    // the behaviour of the banning swarm.
-                    assert_eq!(
-                        swarm2.behaviour.on_connection_established.len(),
-                        s2_expected_conns,
-                        "No additional closed connections should be reported for the banned peer"
-                    );
+                    if swarm1.behaviour.assert_disconnected(s1_expected_conns, 2) {
+                        // The banned connection was established. Given the ban, swarm2 closed the
+                        // connection. Check that it was not reported to the behaviour of the
+                        // banning swarm.
+                        assert_eq!(
+                            swarm2.behaviour.on_connection_established.len(),
+                            s2_expected_conns,
+                            "No additional closed connections should be reported for the banned peer"
+                        );
 
-                    // Setup to test that the banned connection is not reported upon closing
-                    // even if the peer is unbanned.
-                    swarm2.unban_peer_id(swarm1_id);
-                    stage = Stage::Unbanned;
+                        // Setup to test that the banned connection is not reported upon closing
+                        // even if the peer is unbanned.
+                        swarm2.unban_peer_id(swarm1_id);
+                        stage = Stage::Unbanned;
+                    }
                 }
                 Stage::Unbanned => {
-                    if swarm2.network_info().num_peers() == 0 {
+                    if swarm1.network_info().num_peers() == 0
+                        && swarm2.network_info().num_peers() == 0
+                    {
                         // The banned connection has closed. Check that it was not reported.
                         assert_eq!(
                             swarm2.behaviour.on_connection_closed.len(), s2_expected_conns,
diff --git a/swarm/src/test.rs b/swarm/src/test.rs
index 9c8a780fd0c..f8f3daaa067 100644
--- a/swarm/src/test.rs
+++ b/swarm/src/test.rs
@@ -232,8 +232,8 @@ where
             assert_eq!(
                 self.on_connection_established
                     .iter()
-                    .filter(|(.., reported_aditional_connections)| {
-                        *reported_aditional_connections == 0
+                    .filter(|(.., reported_additional_connections)| {
+                        *reported_additional_connections == 0
                     })
                     .count(),
                 expected_connections