Merge pull request #12 from vacp2p/merkle_tree_tests

Merkle tree tests

storage/src/merkle_tree_public/merkle_tree.rs

@@ -95,3 +95,128 @@ impl<Leav: TreeLeavItem + Clone> HashStorageMerkleTree<Leav> {
         self.tree.commit();
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // Mock implementation of TreeLeavItem trait for testing
+    #[derive(Debug, Clone, PartialEq, Eq, Hash)]
+    struct MockTransaction {
+        pub hash: TreeHashType,
+    }
+
+    impl TreeLeavItem for MockTransaction {
+        fn hash(&self) -> TreeHashType {
+            self.hash
+        }
+    }
+
+    fn get_first_32_bytes(s: &str) -> [u8; 32] {
+        let mut buffer = [0u8; 32];
+        let bytes = s.as_bytes();
+        let len = std::cmp::min(32, bytes.len());
+
+        buffer[..len].copy_from_slice(&bytes[..len]);
+        buffer
+    }
+
+    #[test]
+    fn test_new_merkle_tree() {
+        let tx1 = MockTransaction {
+            hash: get_first_32_bytes("tx1"),
+        };
+        let tx2 = MockTransaction {
+            hash: get_first_32_bytes("tx2"),
+        };
+
+        let tree = HashStorageMerkleTree::new(vec![tx1.clone(), tx2.clone()]);
+
+        assert_eq!(tree.leaves.len(), 2);
+        assert!(tree.get_root().is_some());
+    }
+
+    #[test]
+    fn test_get_tx() {
+        let tx1 = MockTransaction {
+            hash: get_first_32_bytes("tx1"),
+        };
+        let tx2 = MockTransaction {
+            hash: get_first_32_bytes("tx2"),
+        };
+
+        let tree = HashStorageMerkleTree::new(vec![tx1.clone(), tx2.clone()]);
+
+        assert_eq!(tree.get_tx(tx1.hash()), Some(&tx1));
+        assert_eq!(tree.get_tx(tx2.hash()), Some(&tx2));
+    }
+
+    #[test]
+    fn test_get_proof() {
+        let tx1 = MockTransaction {
+            hash: get_first_32_bytes("tx1"),
+        };
+        let tx2 = MockTransaction {
+            hash: get_first_32_bytes("tx2"),
+        };
+
+        let tree = HashStorageMerkleTree::new(vec![tx1.clone(), tx2.clone()]);
+
+        let proof = tree.get_proof(tx1.hash());
+        assert!(proof.is_some());
+    }
+
+    #[test]
+    fn test_add_tx() {
+        let tx1 = MockTransaction {
+            hash: get_first_32_bytes("tx1"),
+        };
+        let tx2 = MockTransaction {
+            hash: get_first_32_bytes("tx2"),
+        };
+
+        let mut tree = HashStorageMerkleTree::new(vec![tx1.clone()]);
+
+        tree.add_tx(tx2.clone());
+        assert_eq!(tree.leaves.len(), 2);
+        assert_eq!(tree.get_tx(tx2.hash()), Some(&tx2));
+    }
+
+    #[test]
+    fn test_add_tx_multiple() {
+        let tx1 = MockTransaction {
+            hash: get_first_32_bytes("tx1"),
+        };
+        let tx2 = MockTransaction {
+            hash: get_first_32_bytes("tx2"),
+        };
+        let tx3 = MockTransaction {
+            hash: get_first_32_bytes("tx3"),
+        };
+
+        let mut tree = HashStorageMerkleTree::new(vec![tx1.clone()]);
+        tree.add_tx_multiple(vec![tx2.clone(), tx3.clone()]);
+
+        assert_eq!(tree.leaves.len(), 3);
+        assert_eq!(tree.get_tx(tx2.hash()), Some(&tx2));
+        assert_eq!(tree.get_tx(tx3.hash()), Some(&tx3));
+    }
+
+    #[test]
+    fn test_get_proof_multiple() {
+        let tx1 = MockTransaction {
+            hash: get_first_32_bytes("tx1"),
+        };
+        let tx2 = MockTransaction {
+            hash: get_first_32_bytes("tx2"),
+        };
+        let tx3 = MockTransaction {
+            hash: get_first_32_bytes("tx3"),
+        };
+
+        let tree = HashStorageMerkleTree::new(vec![tx1.clone(), tx2.clone(), tx3.clone()]);
+        let proof = tree.get_proof_multiple(&[tx1.hash(), tx2.hash()]);
+
+        assert!(proof.is_some());
+    }
+}