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 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_chainspec::net::NodeRecord;
use reth_tokio_util::EventStream; use reth_tokio_util::EventStream;
use reth_tracing::tracing::info; use reth_tracing::tracing::info;
@ -32,13 +35,17 @@ impl NetworkTestContext {
self.network.local_node_record() self.network.local_node_record()
} }
/// Expects a session to be established /// Awaits the next event for an established session.
pub async fn expect_session(&mut self) { pub async fn next_session_established(&mut self) -> Option<PeerId> {
match self.network_events.next().await { while let Some(ev) = self.network_events.next().await {
Some(NetworkEvent::SessionEstablished { remote_addr, .. }) => { match ev {
info!(?remote_addr, "Session established") 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>) { pub async fn connect(&mut self, node: &mut NodeTestContext<Node>) {
self.network.add_peer(node.network.record()).await; self.network.add_peer(node.network.record()).await;
node.network.add_peer(self.network.record()).await; node.network.next_session_established().await;
node.network.expect_session().await; self.network.next_session_established().await;
self.network.expect_session().await;
} }
/// Advances the chain `length` blocks. /// Advances the chain `length` blocks.