Wait for connection handshake finished in tests

RingsNetwork · Sep 4, 2023 · 92ed3fb · 92ed3fb
1 parent 8303909
commit 92ed3fb
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 3 deletions.
diff --git a/core/src/tests/default/test_stabilization.rs b/core/src/tests/default/test_stabilization.rs
@@ -72,7 +72,6 @@ async fn test_stabilization_once() -> Result<()> {
             );
         } => { unreachable!(); }
         _ = async {
-            sleep(Duration::from_millis(1000)).await;
             assert!(swarm1.dht().successors().list()?.contains(&key2.address().into()));
             assert!(swarm2.dht().successors().list()?.contains(&key1.address().into()));
             let stabilization = Stabilization::new(Arc::clone(&swarm1), 5usize);
@@ -122,7 +121,6 @@ async fn test_stabilization() -> Result<()> {
             );
         } => { unreachable!(); }
         _ = async {
-            sleep(Duration::from_millis(1000)).await;
             assert!(swarm1.dht().successors().list()?.contains(&key2.address().into()));
             assert!(swarm2.dht().successors().list()?.contains(&key1.address().into()));
             sleep(Duration::from_millis(10000)).await;

diff --git a/core/src/tests/mod.rs b/core/src/tests/mod.rs
@@ -30,4 +30,12 @@ pub async fn manually_establish_connection(swarm1: &Swarm, swarm2: &Swarm) {
 
     assert!(swarm1.get_connection(swarm2.did()).is_some());
     assert!(swarm2.get_connection(swarm1.did()).is_some());
+
+    // Wait for connectin established
+    swarm1
+        .get_connection(swarm2.did())
+        .unwrap()
+        .webrtc_wait_for_data_channel_ready()
+        .await
+        .unwrap();
 }
diff --git a/transport/src/connections/native_webrtc/mod.rs b/transport/src/connections/native_webrtc/mod.rs
@@ -55,7 +55,7 @@ impl WebrtcConnection {
             .ok_or(Error::WebrtcLocalSdpGenerationError)
     }
 
-    async fn webrtc_wait_for_data_channel_ready(&self) -> Result<()> {
+    pub async fn webrtc_wait_for_data_channel_ready(&self) -> Result<()> {
         loop {
             if matches!(
                 self.webrtc_connection_state(),