From d3958f7608940bc0dfa818b7e806e7eb3bb54f9f Mon Sep 17 00:00:00 2001
From: magine <magine233@gmail.com>
Date: Mon, 25 Mar 2024 11:25:35 +0800
Subject: [PATCH] Fix receiver dropping bug in test_stabilization_final_dht

---
 crates/core/src/tests/default/mod.rs                | 6 +++++-
 crates/core/src/tests/default/test_stabilization.rs | 6 +++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/crates/core/src/tests/default/mod.rs b/crates/core/src/tests/default/mod.rs
index 68dbe146..a414dcea 100644
--- a/crates/core/src/tests/default/mod.rs
+++ b/crates/core/src/tests/default/mod.rs
@@ -78,7 +78,11 @@ impl SwarmCallback for NodeCallback {
         &self,
         payload: &MessagePayload,
     ) -> Result<(), Box<dyn std::error::Error>> {
-        self.message_tx.send(payload.clone()).map_err(|e| e.into())
+        // Here we are using on_validate to record messages.
+        // When on_validate return error, the message will be ignored, which is not on purpose.
+        // To prevent returning errors when sending fails, we choose to panic instead.
+        self.message_tx.send(payload.clone()).unwrap();
+        Ok(())
     }
 }
 
diff --git a/crates/core/src/tests/default/test_stabilization.rs b/crates/core/src/tests/default/test_stabilization.rs
index ab24a55b..18125ac0 100644
--- a/crates/core/src/tests/default/test_stabilization.rs
+++ b/crates/core/src/tests/default/test_stabilization.rs
@@ -96,10 +96,13 @@ async fn test_stabilization() -> Result<()> {
 }
 
 #[tokio::test]
-#[ignore]
 async fn test_stabilization_final_dht() -> Result<()> {
     let mut swarms = vec![];
 
+    // Store the nodes to prevent the receiver from being dropped.
+    // That will cause the callback monitor panic when record message.
+    let mut nodes = vec![];
+
     let keys = vec![
         "af3543cde0c40fd217c536a358fb5f3c609eb1135f68daf7e2f2fbd51f164221", // 0xfe81c75f0ef75d7436b84089c5be31b692518d73
         "d842f7143beac06ab8d81589c3c53cffd7eb8e07dadae8fdcb3ed1e1319ab477", // 0x8d4300b4df3c85ee107009e354c1b95717ab1c17
@@ -111,6 +114,7 @@ async fn test_stabilization_final_dht() -> Result<()> {
         let node = prepare_node(key).await;
         swarms.push(node.swarm.clone());
         let stabilization = Arc::new(Stabilization::new(node.swarm.clone(), 3));
+        nodes.push(node);
         tokio::spawn(stabilization.wait());
     }