refactor: enhance readability CardanoTransactionsSetProof tests

Author: jpraynaud

mithril-common/src/entities/cardano_transactions_set_proof.rs

@@ -5,6 +5,12 @@ use crate::{StdError, StdResult};
 
 use super::BlockRange;
 
+cfg_test_tools! {
+    use crate::crypto_helper::{MKMap, MKTree, MKTreeNode, MKMapNode};
+    use crate::entities::BlockNumber;
+    use std::collections::HashMap;
+}
+
 /// A cryptographic proof of a set of Cardano transactions is included in the global Cardano transactions set
 #[derive(Clone, Debug, PartialEq)]
 pub struct CardanoTransactionsSetProof {
@@ -50,42 +56,51 @@ impl CardanoTransactionsSetProof {
     cfg_test_tools! {
         /// Retrieve a dummy proof (for test only)
         pub fn dummy() -> Self {
-            use crate::crypto_helper::{MKMap, MKTree, MKMapNode};
-
-            let transaction_hashes_by_block_range = vec![
-                (
-                    (0..10).into(),
-                    vec!["tx-1".to_string(), "tx-2".to_string(), "tx-3".to_string()],
-                ),
-                (
-                    (10..20).into(),
-                    vec!["tx-4".to_string()],
-                ),
-                (
-                    (20..30).into(),
-                    vec!["tx-5".to_string(), "tx-6".to_string()],
-                ),
+            let leaves = vec![
+                (0, "tx-1".to_string()),
+                (1, "tx-2".to_string()),
+                (1, "tx-3".to_string()),
+                (10, "tx-4".to_string()),
+                (20, "tx-5".to_string()),
+                (22, "tx-6".to_string()),
             ];
-            let transaction_hashes_to_certify = transaction_hashes_by_block_range
-                .iter()
-                .flat_map(|(_, h)| h.to_owned())
-                .collect::<Vec<TransactionHash>>();
 
-            let mk_hash_map: MKMap<_, MKMapNode<_>> = MKMap::new(
-                transaction_hashes_by_block_range
+            Self::from_leaves(&leaves).unwrap()
+        }
+
+        /// Helper to create a proof from a list of leaves
+        pub fn from_leaves(leaves: &[(BlockNumber, TransactionHash)]) -> StdResult<Self> {
+            let transactions_hashes: Vec<TransactionHash> =
+                leaves.iter().map(|(_, t)| t.into()).collect();
+            let mut transactions_by_block_ranges: HashMap<BlockRange, Vec<TransactionHash>> =
+                HashMap::new();
+            for (block_number, transaction_hash) in leaves {
+                let block_range = BlockRange::from_block_number(*block_number);
+                transactions_by_block_ranges
+                    .entry(block_range)
+                    .or_default()
+                    .push(transaction_hash.to_owned());
+            }
+            let mk_map = MKMap::new(
+                transactions_by_block_ranges
                     .into_iter()
-                    .map(|(block_range, transactions)| {
-                        (block_range, MKTree::new(&transactions).unwrap().into())
-                    })
-                    .collect::<Vec<_>>()
+                    .try_fold(
+                        vec![],
+                        |mut acc, (block_range, transactions)| -> StdResult<Vec<(_, MKMapNode<_>)>> {
+                            acc.push((block_range, MKTree::new(&transactions)?.into()));
+                            Ok(acc)
+                        },
+                    )?
                     .as_slice(),
-            )
-            .unwrap();
-            let proof = mk_hash_map
-                .compute_proof(&transaction_hashes_to_certify)
-                .unwrap();
-            Self::new(transaction_hashes_to_certify, proof)
+            )?;
+            let mk_leaves: Vec<MKTreeNode> = transactions_hashes
+                .iter()
+                .map(|h| h.to_owned().into())
+                .collect();
+            let mk_proof = mk_map.compute_proof(&mk_leaves)?;
+            Ok(Self::new(transactions_hashes, mk_proof))
         }
+
     }
 }
 
@@ -117,14 +132,30 @@ mod tests {
 
     #[test]
     fn should_verify_where_all_hashes_are_contained_in_the_proof() {
-        let proof = CardanoTransactionsSetProof::dummy();
+        let leaves = vec![
+            (0, "tx-1".to_string()),
+            (1, "tx-2".to_string()),
+            (1, "tx-3".to_string()),
+            (10, "tx-4".to_string()),
+            (20, "tx-5".to_string()),
+            (22, "tx-6".to_string()),
+        ];
+        let proof = CardanoTransactionsSetProof::from_leaves(&leaves).unwrap();
 
         proof.verify().expect("The proof should be valid");
     }
 
     #[test]
     fn shouldnt_verify_where_at_least_one_hash_is_not_contained_in_the_proof() {
-        let proof = CardanoTransactionsSetProof::dummy();
+        let leaves = vec![
+            (0, "tx-1".to_string()),
+            (1, "tx-2".to_string()),
+            (1, "tx-3".to_string()),
+            (10, "tx-4".to_string()),
+            (20, "tx-5".to_string()),
+            (22, "tx-6".to_string()),
+        ];
+        let proof = CardanoTransactionsSetProof::from_leaves(&leaves).unwrap();
         let mut transactions_hashes_tampered = proof.transactions_hashes().to_vec();
         transactions_hashes_tampered.push("tx-123".to_string());
         let proof = CardanoTransactionsSetProof {