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::{
7    MKMap, MKMapNode, MKTree, MKTreeNode, MKTreeStoreInMemory, MKTreeStorer,
8};
9use crate::entities::{
10    BlockNumber, BlockRange, CardanoTransactionsSetProof, ProtocolParameters, SingleSignature,
11    SingleSignatureAuthenticationStatus, TransactionHash,
12};
13use crate::test::builder::{MithrilFixtureBuilder, StakeDistributionGenerationMethod};
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 [CardanoTransactionsSetProof]
25pub trait CardanoTransactionsSetProofTestExtension {
26    /// `TEST ONLY` - Helper to create a proof from a list of leaves
27    fn from_leaves<S: MKTreeStorer>(
28        leaves: &[(BlockNumber, TransactionHash)],
29    ) -> StdResult<CardanoTransactionsSetProof>;
30}
31
32impl CardanoTransactionsSetProofTestExtension for CardanoTransactionsSetProof {
33    fn from_leaves<S: MKTreeStorer>(
34        leaves: &[(BlockNumber, TransactionHash)],
35    ) -> StdResult<CardanoTransactionsSetProof> {
36        let transactions_hashes: Vec<TransactionHash> =
37            leaves.iter().map(|(_, t)| t.into()).collect();
38        let mut transactions_by_block_ranges: HashMap<BlockRange, Vec<TransactionHash>> =
39            HashMap::new();
40        for (block_number, transaction_hash) in leaves {
41            let block_range = BlockRange::from_block_number(*block_number);
42            transactions_by_block_ranges
43                .entry(block_range)
44                .or_default()
45                .push(transaction_hash.to_owned());
46        }
47        let mk_map = MKMap::<_, _, MKTreeStoreInMemory>::new(
48            transactions_by_block_ranges
49                .into_iter()
50                .try_fold(
51                    vec![],
52                    |mut acc,
53                     (block_range, transactions)|
54                     -> StdResult<Vec<(_, MKMapNode<_, S>)>> {
55                        acc.push((block_range, MKTree::<S>::new(&transactions)?.into()));
56                        Ok(acc)
57                    },
58                )?
59                .as_slice(),
60        )?;
61        let mk_leaves: Vec<MKTreeNode> =
62            transactions_hashes.iter().map(|h| h.to_owned().into()).collect();
63        let mk_proof = mk_map.compute_proof(&mk_leaves)?;
64        Ok(Self::new(transactions_hashes, mk_proof))
65    }
66}
67
68/// Extension trait adding test utilities to [SingleSignature]
69pub trait SingleSignatureTestExtension {
70    /// `TEST ONLY` - Create a fake [SingleSignature] with valid cryptographic data for testing purposes.
71    fn fake<TPartyId: Into<String>, TMessage: Into<String>>(
72        party_id: TPartyId,
73        message: TMessage,
74    ) -> SingleSignature;
75}
76
77impl SingleSignatureTestExtension for SingleSignature {
78    fn fake<TPartyId: Into<String>, TMessage: Into<String>>(
79        party_id: TPartyId,
80        message: TMessage,
81    ) -> SingleSignature {
82        let party_id = party_id.into();
83        let message = message.into();
84
85        let fixture = MithrilFixtureBuilder::default()
86            .with_stake_distribution(StakeDistributionGenerationMethod::Custom(
87                std::collections::BTreeMap::from([(party_id.to_string(), 100)]),
88            ))
89            .with_protocol_parameters(ProtocolParameters::new(1, 1, 1.0))
90            .build();
91        let signature = fixture.signers_fixture()[0].sign(&message).unwrap();
92
93        Self {
94            party_id,
95            signature: signature.signature,
96            won_indexes: vec![10, 15],
97            authentication_status: SingleSignatureAuthenticationStatus::Unauthenticated,
98        }
99    }
100}