fix(kad): rewrite tests to be less flaky

Not sure why but these tests have recently been quite flaky and are blocking several PRs.

Pull-Request: #4634.

protocols/kad/tests/client_mode.rs

@@ -2,8 +2,10 @@ use libp2p_identify as identify;
 use libp2p_identity as identity;
 use libp2p_kad::store::MemoryStore;
 use libp2p_kad::{Behaviour, Config, Event, Mode};
-use libp2p_swarm::Swarm;
+use libp2p_swarm::{Swarm, SwarmEvent};
 use libp2p_swarm_test::SwarmExt;
+use Event::*;
+use MyBehaviourEvent::*;
 
 #[async_std::test]
 async fn server_gets_added_to_routing_table_by_client() {
@@ -16,16 +18,16 @@ async fn server_gets_added_to_routing_table_by_client() {
     client.connect(&mut server).await;
 
     let server_peer_id = *server.local_peer_id();
+    async_std::task::spawn(server.loop_on_next());
 
-    match libp2p_swarm_test::drive(&mut client, &mut server).await {
+    let peer = client
-        (
+        .wait(|e| match e {
-            [MyBehaviourEvent::Identify(_), MyBehaviourEvent::Identify(_), MyBehaviourEvent::Kad(Event::RoutingUpdated { peer, .. })],
+            SwarmEvent::Behaviour(Kad(RoutingUpdated { peer, .. })) => Some(peer),
-            [MyBehaviourEvent::Identify(_), MyBehaviourEvent::Identify(_)],
+            _ => None,
-        ) => {
+        })
-            assert_eq!(peer, server_peer_id)
+        .await;
-        }
+
-        other => panic!("Unexpected events: {other:?}"),
+    assert_eq!(peer, server_peer_id);
-    }
 }
 
 #[async_std::test]
@@ -41,12 +43,10 @@ async fn two_servers_add_each_other_to_routing_table() {
     let server1_peer_id = *server1.local_peer_id();
     let server2_peer_id = *server2.local_peer_id();
 
-    use Event::*;
-    use MyBehaviourEvent::*;
-
     match libp2p_swarm_test::drive(&mut server1, &mut server2).await {
         (
-            [Identify(_), Identify(_), Kad(RoutingUpdated { peer: peer1, .. })],
+            [Identify(_), Identify(_), Kad(RoutingUpdated { peer: peer1, .. })]
+            | [Identify(_), Kad(RoutingUpdated { peer: peer1, .. }), Identify(_)],
             [Identify(_), Identify(_)],
         ) => {
             assert_eq!(peer1, server2_peer_id);
@@ -57,19 +57,16 @@ async fn two_servers_add_each_other_to_routing_table() {
     server1.listen().await;
     server2.connect(&mut server1).await;
 
-    match libp2p_swarm_test::drive(&mut server2, &mut server1).await {
+    async_std::task::spawn(server1.loop_on_next());
-        (
+
-            [Identify(_), Kad(RoutingUpdated { peer: peer2, .. }), Identify(_)],
+    let peer = server2
-            [Identify(_), Identify(_)],
+        .wait(|e| match e {
-        )
+            SwarmEvent::Behaviour(Kad(RoutingUpdated { peer, .. })) => Some(peer),
-        | (
+            _ => None,
-            [Identify(_), Identify(_), Kad(RoutingUpdated { peer: peer2, .. })],
+        })
-            [Identify(_), Identify(_)],
+        .await;
-        ) => {
+
-            assert_eq!(peer2, server1_peer_id);
+    assert_eq!(peer, server1_peer_id);
-        }
-        other => panic!("Unexpected events: {other:?}"),
-    }
 }
 
 #[async_std::test]
@@ -85,17 +82,12 @@ async fn adding_an_external_addresses_activates_server_mode_on_existing_connecti
     // Remove memory address to simulate a server that doesn't know its external address.
     server.remove_external_address(&memory_addr);
     client.dial(memory_addr.clone()).unwrap();
-
-    use MyBehaviourEvent::*;
-
     // Do the usual identify send/receive dance.
     match libp2p_swarm_test::drive(&mut client, &mut server).await {
         ([Identify(_), Identify(_)], [Identify(_), Identify(_)]) => {}
         other => panic!("Unexpected events: {other:?}"),
     }
 
-    use Event::*;
-
     // Server learns its external address (this could be through AutoNAT or some other mechanism).
     server.add_external_address(memory_addr);
 
@@ -125,34 +117,38 @@ async fn set_client_to_server_mode() {
 
     let server_peer_id = *server.local_peer_id();
 
-    match libp2p_swarm_test::drive(&mut client, &mut server).await {
+    let client_event = client.wait(|e| match e {
-        (
+        SwarmEvent::Behaviour(Kad(RoutingUpdated { peer, .. })) => Some(peer),
-            [MyBehaviourEvent::Identify(_), MyBehaviourEvent::Identify(_), MyBehaviourEvent::Kad(Event::RoutingUpdated { peer, .. })],
+        _ => None,
-            [MyBehaviourEvent::Identify(_), MyBehaviourEvent::Identify(identify::Event::Received { info, .. })],
+    });
-        ) => {
+    let server_event = server.wait(|e| match e {
-            assert_eq!(peer, server_peer_id);
+        SwarmEvent::Behaviour(Identify(identify::Event::Received { info, .. })) => Some(info),
-            assert!(info
+        _ => None,
-                .protocols
+    });
-                .iter()
+
-                .all(|proto| libp2p_kad::PROTOCOL_NAME.ne(proto)))
+    let (peer, info) = futures::future::join(client_event, server_event).await;
-        }
+
-        other => panic!("Unexpected events: {other:?}"),
+    assert_eq!(peer, server_peer_id);
-    }
+    assert!(info
+        .protocols
+        .iter()
+        .all(|proto| libp2p_kad::PROTOCOL_NAME.ne(proto)));
 
     client.behaviour_mut().kad.set_mode(Some(Mode::Server));
 
-    match libp2p_swarm_test::drive(&mut client, &mut server).await {
+    async_std::task::spawn(client.loop_on_next());
-        (
+
-            [MyBehaviourEvent::Identify(_)],
+    let info = server
-            [MyBehaviourEvent::Identify(identify::Event::Received { info, .. }), MyBehaviourEvent::Kad(_)],
+        .wait(|e| match e {
-        ) => {
+            SwarmEvent::Behaviour(Identify(identify::Event::Received { info, .. })) => Some(info),
-            assert!(info
+            _ => None,
-                .protocols
+        })
-                .iter()
+        .await;
-                .any(|proto| libp2p_kad::PROTOCOL_NAME.eq(proto)))
+
-        }
+    assert!(info
-        other => panic!("Unexpected events: {other:?}"),
+        .protocols
-    }
+        .iter()
+        .any(|proto| libp2p_kad::PROTOCOL_NAME.eq(proto)));
 }
 
 #[derive(libp2p_swarm::NetworkBehaviour)]