From b0538afcc160e5fbda3f4a29862f5e60594f3d93 Mon Sep 17 00:00:00 2001
From: qima <qi.ma@maidsafe.net>
Date: Wed, 21 Aug 2024 20:29:57 +0800
Subject: [PATCH] test(kad): test get_closest_with_different_num_results

---
 protocols/kad/src/behaviour.rs      |  1 +
 protocols/kad/src/behaviour/test.rs | 73 +++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+)

diff --git a/protocols/kad/src/behaviour.rs b/protocols/kad/src/behaviour.rs
index cb1e29360568..94d26cdd6baa 100644
--- a/protocols/kad/src/behaviour.rs
+++ b/protocols/kad/src/behaviour.rs
@@ -744,6 +744,7 @@ where
     /// [`Event::OutboundQueryProgressed{QueryResult::GetClosestPeers}`].
     ///
     /// The expected responding peers is specified by `num_results`
+    /// Note that the result is capped after exceeds K_VALUE
     pub fn get_closest_peers_num_results<K>(&mut self, key: K, num_results: NonZeroUsize) -> QueryId
     where
         K: Into<kbucket::Key<K>> + Into<Vec<u8>> + Clone,
diff --git a/protocols/kad/src/behaviour/test.rs b/protocols/kad/src/behaviour/test.rs
index 276b156c9ae1..026be2aedd9f 100644
--- a/protocols/kad/src/behaviour/test.rs
+++ b/protocols/kad/src/behaviour/test.rs
@@ -425,6 +425,79 @@ fn unresponsive_not_returned_indirect() {
     }))
 }
 
+#[test]
+// Test the result of get_closest_peers with different num_results
+// Note that the result is capped after exceeds K_VALUE
+fn get_closest_with_different_num_results() {
+    let k_value = K_VALUE.get();
+    for num_results in k_value / 2..k_value * 2 {
+        get_closest_with_different_num_results_inner(num_results)
+    }
+}
+
+fn get_closest_with_different_num_results_inner(num_results: usize) {
+    let k_value = K_VALUE.get();
+    let num_of_nodes = 3 * k_value;
+    let mut swarms = build_nodes(num_of_nodes);
+
+    // Connect second to first.
+    let second_peer_id = *swarms[1].1.local_peer_id();
+    let second_address = swarms[1].0.clone();
+    swarms[0]
+        .1
+        .behaviour_mut()
+        .add_address(&second_peer_id, second_address);
+
+    // Connect others to second.
+    for index in 2..(num_of_nodes) {
+        let peer_id = *swarms[index].1.local_peer_id();
+        let address = swarms[index].0.clone();
+        swarms[1].1.behaviour_mut().add_address(&peer_id, address);
+    }
+
+    let mut swarms = swarms
+        .into_iter()
+        .map(|(_addr, swarm)| swarm)
+        .collect::<Vec<_>>();
+
+    // Ask first to search a random value.
+    let search_target = PeerId::random();
+    let Some(num_results_nonzero) = std::num::NonZeroUsize::new(num_results) else {
+        panic!("Unexpected NonZeroUsize val of {num_results}");
+    };
+    swarms[0]
+        .behaviour_mut()
+        .get_closest_peers_num_results(search_target, num_results_nonzero);
+
+    block_on(poll_fn(move |ctx| {
+        for swarm in &mut swarms {
+            loop {
+                match swarm.poll_next_unpin(ctx) {
+                    Poll::Ready(Some(SwarmEvent::Behaviour(Event::OutboundQueryProgressed {
+                        result: QueryResult::GetClosestPeers(Ok(ok)),
+                        ..
+                    }))) => {
+                        assert_eq!(&ok.key[..], search_target.to_bytes().as_slice());
+                        if num_results > k_value {
+                            assert_eq!(ok.peers.len(), k_value + 1);
+                        } else {
+                            assert_eq!(ok.peers.len(), num_results);
+                        }
+
+                        return Poll::Ready(());
+                    }
+                    // Ignore any other event.
+                    Poll::Ready(Some(_)) => (),
+                    e @ Poll::Ready(_) => panic!("Unexpected return value: {e:?}"),
+                    Poll::Pending => break,
+                }
+            }
+        }
+
+        Poll::Pending
+    }))
+}
+
 #[test]
 fn get_record_not_found() {
     let mut swarms = build_nodes(3);