test: fix flaky connect (#9113)

This commit is contained in:
Matthias Seitz
2024-06-26 12:07:49 +02:00
committed by GitHub
parent ce5a191af6
commit 832f7a5170
2 changed files with 17 additions and 10 deletions

View File

@ -1,5 +1,8 @@
use futures_util::StreamExt;
use reth::network::{NetworkEvent, NetworkEvents, NetworkHandle, PeersInfo};
use reth::{
network::{NetworkEvent, NetworkEvents, NetworkHandle, PeersInfo},
rpc::types::PeerId,
};
use reth_chainspec::net::NodeRecord;
use reth_tokio_util::EventStream;
use reth_tracing::tracing::info;
@ -32,13 +35,17 @@ impl NetworkTestContext {
self.network.local_node_record()
}
/// Expects a session to be established
pub async fn expect_session(&mut self) {
match self.network_events.next().await {
Some(NetworkEvent::SessionEstablished { remote_addr, .. }) => {
info!(?remote_addr, "Session established")
/// Awaits the next event for an established session.
pub async fn next_session_established(&mut self) -> Option<PeerId> {
while let Some(ev) = self.network_events.next().await {
match ev {
NetworkEvent::SessionEstablished { peer_id, .. } => {
info!("Session established with peer: {:?}", peer_id);
return Some(peer_id)
}
_ => continue,
}
ev => panic!("Expected session established event, got: {ev:?}"),
}
None
}
}

View File

@ -52,11 +52,11 @@ where
})
}
/// Establish a connection to the node
pub async fn connect(&mut self, node: &mut NodeTestContext<Node>) {
self.network.add_peer(node.network.record()).await;
node.network.add_peer(self.network.record()).await;
node.network.expect_session().await;
self.network.expect_session().await;
node.network.next_session_established().await;
self.network.next_session_established().await;
}
/// Advances the chain `length` blocks.