diff --git a/storage/src/merkle_tree_public/merkle_tree.rs b/storage/src/merkle_tree_public/merkle_tree.rs index ac933ea..16d185e 100644 --- a/storage/src/merkle_tree_public/merkle_tree.rs +++ b/storage/src/merkle_tree_public/merkle_tree.rs @@ -95,3 +95,128 @@ impl HashStorageMerkleTree { 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()); + } +}