test: fix flaky connect

crates/e2e-test-utils/src/network.rs

@@ -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
     }
 }

crates/e2e-test-utils/src/node.rs

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