From c85c1cd7591f519f4e23b0ee037d803839f5fef7 Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Wed, 26 Jun 2024 09:29:10 +0200
Subject: [PATCH] test: fix flaky connect

---
 crates/e2e-test-utils/src/network.rs | 21 ++++++++++++++-------
 crates/e2e-test-utils/src/node.rs    |  6 +++---
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/crates/e2e-test-utils/src/network.rs b/crates/e2e-test-utils/src/network.rs
index 2bb69044fc42..8c1a60ba7bd9 100644
--- a/crates/e2e-test-utils/src/network.rs
+++ b/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
     }
 }
diff --git a/crates/e2e-test-utils/src/node.rs b/crates/e2e-test-utils/src/node.rs
index 524a5d5562be..684e0f401c1b 100644
--- a/crates/e2e-test-utils/src/node.rs
+++ b/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.