swarm/src/lib: Improve connection counting for test_behaviour_disconnect_all (#2532)

This commit is contained in:
Max Inden
2022-03-02 12:10:57 +01:00
committed by GitHub
parent a179dbb449
commit b919d0002a
2 changed files with 36 additions and 16 deletions

View File

@ -1583,27 +1583,36 @@ mod tests {
TBehaviour: NetworkBehaviour, TBehaviour: NetworkBehaviour,
<<TBehaviour::ConnectionHandler as IntoConnectionHandler>::Handler as ConnectionHandler>::OutEvent: Clone, <<TBehaviour::ConnectionHandler as IntoConnectionHandler>::Handler as ConnectionHandler>::OutEvent: Clone,
{ {
[swarm1, swarm2] swarm1
.iter() .behaviour()
.all(|s| s.behaviour.inject_connection_established.len() == num_connections) .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>( fn swarms_disconnected<TBehaviour: NetworkBehaviour>(
swarm1: &Swarm<CallTraceBehaviour<TBehaviour>>, swarm1: &Swarm<CallTraceBehaviour<TBehaviour>>,
swarm2: &Swarm<CallTraceBehaviour<TBehaviour>>, swarm2: &Swarm<CallTraceBehaviour<TBehaviour>>,
num_connections: usize,
) -> bool ) -> bool
where where
TBehaviour: NetworkBehaviour, TBehaviour: NetworkBehaviour,
<<TBehaviour::ConnectionHandler as IntoConnectionHandler>::Handler as ConnectionHandler>::OutEvent: Clone <<TBehaviour::ConnectionHandler as IntoConnectionHandler>::Handler as ConnectionHandler>::OutEvent: Clone
{ {
[swarm1, swarm2] swarm1
.iter() .behaviour()
.all(|s| s.behaviour.inject_connection_closed.len() == num_connections) .num_connections_to_peer(*swarm2.local_peer_id())
&& [swarm1, swarm2].iter().all(|s| { == 0
let (.., last_remaining) = s.behaviour.inject_connection_closed.last().unwrap(); && swarm2
*last_remaining == 0 .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, /// Establishes multiple connections between two peers,
@ -1781,7 +1790,7 @@ mod tests {
} }
} }
State::Disconnecting => { State::Disconnecting => {
if swarms_disconnected(&swarm1, &swarm2, num_connections) { if swarms_disconnected(&swarm1, &swarm2) {
if reconnected { if reconnected {
return Poll::Ready(()); return Poll::Ready(());
} }
@ -1850,18 +1859,17 @@ mod tests {
}, },
); );
state = State::Disconnecting; state = State::Disconnecting;
continue;
} }
} }
State::Disconnecting => { State::Disconnecting => {
if swarms_disconnected(&swarm1, &swarm2, num_connections) { if swarms_disconnected(&swarm1, &swarm2) {
if reconnected {
return Poll::Ready(());
}
reconnected = true; reconnected = true;
for _ in 0..num_connections { for _ in 0..num_connections {
swarm2.dial(addr1.clone()).unwrap(); swarm2.dial(addr1.clone()).unwrap();
} }
state = State::Connecting; state = State::Connecting;
continue;
} }
} }
} }

View File

@ -159,6 +159,18 @@ where
&mut self.inner &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 /// Checks that when the expected number of closed connection notifications are received, a
/// given number of expected disconnections have been received as well. /// given number of expected disconnections have been received as well.
/// ///