From 4f661a7c428139aa17b77acb34f42d32b7512633 Mon Sep 17 00:00:00 2001
From: yukang <moorekang@gmail.com>
Date: Wed, 24 Jan 2024 21:19:11 +0800
Subject: [PATCH] add unit test for links

---
 tx-pool/src/component/links.rs       |  6 +++++-
 tx-pool/src/component/pool_map.rs    | 14 ++++++------
 tx-pool/src/component/tests/links.rs | 32 ++++++++++++++++++++++++++++
 tx-pool/src/component/tests/mod.rs   |  1 +
 4 files changed, 46 insertions(+), 7 deletions(-)
 create mode 100644 tx-pool/src/component/tests/links.rs

diff --git a/tx-pool/src/component/links.rs b/tx-pool/src/component/links.rs
index 106f9a216d..428b1fb1dd 100644
--- a/tx-pool/src/component/links.rs
+++ b/tx-pool/src/component/links.rs
@@ -72,13 +72,17 @@ impl TxLinksMap {
             relation_ids.insert(id);
         }
         // for direct parents, we don't store children in links map
-        // so filter those not in links map now, they maybe removed from tx-pool now
+        // so filter those not in links map, they maybe removed from tx-pool now
         if relation == Relation::DirectParents {
             relation_ids.retain(|id| self.inner.contains_key(id));
         }
         relation_ids
     }
 
+    pub fn add_link(&mut self, short_id: ProposalShortId, links: TxLinks) {
+        self.inner.insert(short_id, links);
+    }
+
     pub fn calc_ancestors(&self, short_id: &ProposalShortId) -> HashSet<ProposalShortId> {
         self.calc_relative_ids(short_id, Relation::Parents)
     }
diff --git a/tx-pool/src/component/pool_map.rs b/tx-pool/src/component/pool_map.rs
index 36adb4f7d7..e3daee1615 100644
--- a/tx-pool/src/component/pool_map.rs
+++ b/tx-pool/src/component/pool_map.rs
@@ -511,12 +511,14 @@ impl PoolMap {
             self.links.add_child(parent, short_id.clone());
         }
 
-        let links = TxLinks {
-            parents,
-            direct_parents,
-            children: Default::default(),
-        };
-        self.links.inner.insert(short_id, links);
+        self.links.add_link(
+            short_id,
+            TxLinks {
+                parents,
+                direct_parents,
+                children: Default::default(),
+            },
+        );
 
         Ok(true)
     }
diff --git a/tx-pool/src/component/tests/links.rs b/tx-pool/src/component/tests/links.rs
new file mode 100644
index 0000000000..e8e82414c9
--- /dev/null
+++ b/tx-pool/src/component/tests/links.rs
@@ -0,0 +1,32 @@
+use crate::component::links::{Relation, TxLinks, TxLinksMap};
+use ckb_types::packed::ProposalShortId;
+use ckb_types::prelude::Entity;
+use std::collections::HashSet;
+
+#[test]
+fn test_link_map() {
+    let mut map = TxLinksMap::default();
+    let id1 = ProposalShortId::from_slice(&[1; 10]).unwrap();
+    let id2 = ProposalShortId::from_slice(&[2; 10]).unwrap();
+    let id3 = ProposalShortId::from_slice(&[3; 10]).unwrap();
+    let id4 = ProposalShortId::from_slice(&[4; 10]).unwrap();
+
+    map.add_link(id1.clone(), TxLinks::default());
+    map.add_link(id2.clone(), TxLinks::default());
+    map.add_link(id3.clone(), TxLinks::default());
+    map.add_link(id4.clone(), TxLinks::default());
+
+    map.add_parent(&id1, id2.clone());
+    let expect: HashSet<ProposalShortId> = vec![id2.clone()].into_iter().collect();
+    assert_eq!(map.get_parents(&id1).unwrap(), &expect);
+
+    map.add_direct_parent(&id1, id2.clone());
+    map.add_direct_parent(&id2, id3.clone());
+    map.add_direct_parent(&id3, id4.clone());
+    let direct_parents = map.calc_relation_ids([id1.clone()].into(), Relation::DirectParents);
+    assert_eq!(direct_parents.len(), 4);
+
+    map.remove(&id3);
+    let direct_parents = map.calc_relation_ids([id1.clone()].into(), Relation::DirectParents);
+    assert_eq!(direct_parents.len(), 2);
+}
diff --git a/tx-pool/src/component/tests/mod.rs b/tx-pool/src/component/tests/mod.rs
index 13ba389e82..ac625eede3 100644
--- a/tx-pool/src/component/tests/mod.rs
+++ b/tx-pool/src/component/tests/mod.rs
@@ -1,5 +1,6 @@
 mod chunk;
 mod entry;
+mod links;
 mod orphan;
 mod pending;
 mod proposed;