swarm/src/lib: Improve connection counting for test_behaviour_disconn…

…ect_all (#2532)
libp2p · Mar 2, 2022 · b919d00 · b919d00
1 parent a179dbb
commit b919d00
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 16 deletions.
diff --git a/swarm/src/lib.rs b/swarm/src/lib.rs
@@ -1583,27 +1583,36 @@ mod tests {
         TBehaviour: NetworkBehaviour,
         <<TBehaviour::ConnectionHandler as IntoConnectionHandler>::Handler as ConnectionHandler>::OutEvent: Clone,
     {
-        [swarm1, swarm2]
-            .iter()
-            .all(|s| s.behaviour.inject_connection_established.len() == num_connections)
+        swarm1
+            .behaviour()
+            .num_connections_to_peer(*swarm2.local_peer_id())
+            == num_connections
+            && swarm2
+                .behaviour()
+                .num_connections_to_peer(*swarm1.local_peer_id())
+                == num_connections
+            && swarm1.is_connected(swarm2.local_peer_id())
+            && swarm2.is_connected(swarm1.local_peer_id())
     }
 
     fn swarms_disconnected<TBehaviour: NetworkBehaviour>(
         swarm1: &Swarm<CallTraceBehaviour<TBehaviour>>,
         swarm2: &Swarm<CallTraceBehaviour<TBehaviour>>,
-        num_connections: usize,
     ) -> bool
     where
         TBehaviour: NetworkBehaviour,
         <<TBehaviour::ConnectionHandler as IntoConnectionHandler>::Handler as ConnectionHandler>::OutEvent: Clone
     {
-        [swarm1, swarm2]
-            .iter()
-            .all(|s| s.behaviour.inject_connection_closed.len() == num_connections)
-            && [swarm1, swarm2].iter().all(|s| {
-                let (.., last_remaining) = s.behaviour.inject_connection_closed.last().unwrap();
-                *last_remaining == 0
-            })
+        swarm1
+            .behaviour()
+            .num_connections_to_peer(*swarm2.local_peer_id())
+            == 0
+            && swarm2
+                .behaviour()
+                .num_connections_to_peer(*swarm1.local_peer_id())
+                == 0
+            && !swarm1.is_connected(swarm2.local_peer_id())
+            && !swarm2.is_connected(swarm1.local_peer_id())
     }
 
     /// Establishes multiple connections between two peers,
@@ -1781,7 +1790,7 @@ mod tests {
                     }
                 }
                 State::Disconnecting => {
-                    if swarms_disconnected(&swarm1, &swarm2, num_connections) {
+                    if swarms_disconnected(&swarm1, &swarm2) {
                         if reconnected {
                             return Poll::Ready(());
                         }
@@ -1850,18 +1859,17 @@ mod tests {
                             },
                         );
                         state = State::Disconnecting;
+                        continue;
                     }
                 }
                 State::Disconnecting => {
-                    if swarms_disconnected(&swarm1, &swarm2, num_connections) {
-                        if reconnected {
-                            return Poll::Ready(());
-                        }
+                    if swarms_disconnected(&swarm1, &swarm2) {
                         reconnected = true;
                         for _ in 0..num_connections {
                             swarm2.dial(addr1.clone()).unwrap();
                         }
                         state = State::Connecting;
+                        continue;
                     }
                 }
             }

diff --git a/swarm/src/test.rs b/swarm/src/test.rs
@@ -159,6 +159,18 @@ where
         &mut self.inner
     }
 
+    pub fn num_connections_to_peer(&self, peer: PeerId) -> usize {
+        self.inject_connection_established
+            .iter()
+            .filter(|(peer_id, _, _, _)| *peer_id == peer)
+            .count()
+            - self
+                .inject_connection_closed
+                .iter()
+                .filter(|(peer_id, _, _, _)| *peer_id == peer)
+                .count()
+    }
+
     /// Checks that when the expected number of closed connection notifications are received, a
     /// given number of expected disconnections have been received as well.
     ///