fix test initializtion & fn on_connection_established

libp2p · Oct 17, 2023 · dd730da · dd730da
1 parent 26257f7
commit dd730da
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/protocols/kad/src/behaviour.rs b/protocols/kad/src/behaviour.rs
@@ -1884,12 +1884,18 @@ where
         ConnectionEstablished {
             peer_id,
             failed_addresses,
+            other_established,
             ..
         }: ConnectionEstablished,
     ) {
         for addr in failed_addresses {
             self.address_failed(peer_id, addr);
         }
+
+        // Peer's first connection.
+        if other_established == 0 {
+            self.connected_peers.insert(peer_id);
+        }
     }
 
     fn on_address_change(
@@ -2025,8 +2031,6 @@ where
         }) {
             handler.on_behaviour_event(event)
         }
-
-        self.connected_peers.insert(peer);
     }
 }
 

diff --git a/protocols/kad/src/behaviour/test.rs b/protocols/kad/src/behaviour/test.rs
@@ -72,7 +72,8 @@ fn build_node_with_config(cfg: Config) -> (Multiaddr, TestSwarm) {
         transport,
         behaviour,
         local_id,
-        swarm::Config::with_async_std_executor(),
+        swarm::Config::with_async_std_executor()
+            .with_idle_connection_timeout(Duration::from_secs(60 * 5)),
     );
 
     let address: Multiaddr = Protocol::Memory(random::<u64>()).into();