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::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;
}
}
}

View File

@ -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.
///