mithril_common/test/
entities_extensions.rs

1//! A set of extension traits to add test utilities to this crate `entities`
2
3use std::collections::HashMap;
4
5use crate::StdResult;
6use crate::crypto_helper::MKTreeStorer;
7use crate::entities::{
8    BlockNumber, BlockRange, CardanoBlock, CardanoTransaction, CardanoTransactionsSetProof,
9    IntoMKTreeNode, MkSetProof, ProtocolParameters, SingleSignature,
10    SingleSignatureAuthenticationStatus, TransactionHash,
11};
12use crate::test::builder::{MithrilFixtureBuilder, StakeDistributionGenerationMethod};
13use crate::test::crypto_helper::mkmap_helpers;
14
15/// Extension trait adding test utilities to [BlockRange]
16pub trait BlockRangeTestExtension {
17    /// `TEST ONLY` - BlockRange factory
18    fn new(start: u64, end: u64) -> Self;
19
20    /// `TEST ONLY` - Try to add two BlockRanges
21    fn try_add(&self, other: &BlockRange) -> StdResult<BlockRange>;
22}
23
24/// Extension trait adding test utilities to [BlockNumber]
25pub trait BlockNumberTestExtension {
26    /// `TEST ONLY` - Group an iterator over a tuple of block number and item
27    fn group_items_by_block_range<N, I>(iter: I) -> HashMap<BlockRange, Vec<N>>
28    where
29        I: Iterator<Item = (BlockNumber, N)>;
30}
31
32impl BlockNumberTestExtension for BlockNumber {
33    fn group_items_by_block_range<N, I>(iter: I) -> HashMap<BlockRange, Vec<N>>
34    where
35        I: Iterator<Item = (BlockNumber, N)>,
36    {
37        let mut result: HashMap<BlockRange, Vec<N>> = HashMap::new();
38
39        for (block_number, item) in iter.into_iter() {
40            let block_range = BlockRange::from_block_number(block_number);
41            result.entry(block_range).or_default().push(item);
42        }
43
44        result
45    }
46}
47
48/// Extension trait adding test utilities to [CardanoTransactionsSetProof]
49pub trait CardanoTransactionsSetProofTestExtension {
50    /// `TEST ONLY` - Helper to create a proof from a list of leaves
51    fn from_leaves<S: MKTreeStorer>(
52        leaves: &[(BlockNumber, TransactionHash)],
53    ) -> StdResult<CardanoTransactionsSetProof>;
54}
55
56impl CardanoTransactionsSetProofTestExtension for CardanoTransactionsSetProof {
57    fn from_leaves<S: MKTreeStorer>(
58        leaves: &[(BlockNumber, TransactionHash)],
59    ) -> StdResult<CardanoTransactionsSetProof> {
60        let transactions_hashes: Vec<TransactionHash> =
61            leaves.iter().map(|(_, t)| t.into()).collect();
62        let transactions_by_block_ranges =
63            BlockNumber::group_items_by_block_range(leaves.iter().cloned());
64        let mk_map = mkmap_helpers::fold_nodes_per_block_range_into_mkmap::<_, _, S>(
65            transactions_by_block_ranges,
66        )?;
67        let mk_proof = mk_map.compute_proof(&transactions_hashes)?;
68        Ok(Self::new(transactions_hashes, mk_proof))
69    }
70}
71
72/// Extension trait adding test utilities to [MkSetProof]
73pub trait MkSetProofTestExtension<L: IntoMKTreeNode + Clone> {
74    /// `TEST ONLY` - Helper to create a proof from a list of leaves
75    fn from_leaves<S: MKTreeStorer>(leaves: &[L]) -> StdResult<MkSetProof<L>>;
76}
77
78impl MkSetProofTestExtension<CardanoBlock> for MkSetProof<CardanoBlock> {
79    fn from_leaves<S: MKTreeStorer>(
80        leaves: &[CardanoBlock],
81    ) -> StdResult<MkSetProof<CardanoBlock>> {
82        let node_per_block_range = BlockNumber::group_items_by_block_range(
83            leaves.iter().map(|l| (l.block_number, l.clone().into_mk_tree_node())),
84        );
85        let all_nodes = node_per_block_range.values().flatten().cloned().collect::<Vec<_>>();
86        let mk_map =
87            mkmap_helpers::fold_nodes_per_block_range_into_mkmap::<_, _, S>(node_per_block_range)?;
88        let proof = mk_map.compute_proof(&all_nodes)?;
89
90        Ok(MkSetProof::new(leaves.to_vec(), proof))
91    }
92}
93
94impl MkSetProofTestExtension<CardanoTransaction> for MkSetProof<CardanoTransaction> {
95    fn from_leaves<S: MKTreeStorer>(
96        leaves: &[CardanoTransaction],
97    ) -> StdResult<MkSetProof<CardanoTransaction>> {
98        let node_per_block_range = BlockNumber::group_items_by_block_range(
99            leaves.iter().map(|l| (l.block_number, l.clone().into_mk_tree_node())),
100        );
101        let all_nodes = node_per_block_range.values().flatten().cloned().collect::<Vec<_>>();
102        let mk_map =
103            mkmap_helpers::fold_nodes_per_block_range_into_mkmap::<_, _, S>(node_per_block_range)?;
104        let proof = mk_map.compute_proof(&all_nodes)?;
105
106        Ok(MkSetProof::new(leaves.to_vec(), proof))
107    }
108}
109
110/// Extension trait adding test utilities to [SingleSignature]
111pub trait SingleSignatureTestExtension {
112    /// `TEST ONLY` - Create a fake [SingleSignature] with valid cryptographic data for testing purposes.
113    fn fake<TPartyId: Into<String>, TMessage: Into<String>>(
114        party_id: TPartyId,
115        message: TMessage,
116    ) -> SingleSignature;
117}
118
119impl SingleSignatureTestExtension for SingleSignature {
120    fn fake<TPartyId: Into<String>, TMessage: Into<String>>(
121        party_id: TPartyId,
122        message: TMessage,
123    ) -> SingleSignature {
124        let party_id = party_id.into();
125        let message = message.into();
126
127        let fixture = MithrilFixtureBuilder::default()
128            .with_stake_distribution(StakeDistributionGenerationMethod::Custom(
129                std::collections::BTreeMap::from([(party_id.to_string(), 100)]),
130            ))
131            .with_protocol_parameters(ProtocolParameters::new(1, 1, 1.0))
132            .build();
133        let signature = fixture.signers_fixture()[0].sign(&message).unwrap();
134
135        Self {
136            party_id,
137            signature: signature.signature,
138            won_indexes: vec![10, 15],
139            authentication_status: SingleSignatureAuthenticationStatus::Unauthenticated,
140        }
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147
148    #[test]
149    fn group_list_of_blocknumber_and_string_tuple_by_block_range() {
150        let input = [
151            (BlockNumber(1), "item_1"),
152            (BlockNumber(2), "item_2"),
153            (BlockNumber(3), "item_3"),
154            (BlockNumber(16), "item_16"),
155            (BlockNumber(17), "item_17"),
156        ];
157
158        let grouped_items = BlockNumber::group_items_by_block_range(input.iter().cloned());
159
160        assert_eq!(
161            HashMap::from([
162                (
163                    BlockRange::from_block_number(BlockNumber(1)),
164                    vec!["item_1", "item_2", "item_3"]
165                ),
166                (
167                    BlockRange::from_block_number(BlockNumber(16)),
168                    vec!["item_16", "item_17"]
169                ),
170            ]),
171            grouped_items
172        );
173    }
174}
mithril_common/test/entities_extensions.rs

mithril_common/test/
entities_extensions.rs
