iota_swarm_config/
test_utils.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{collections::HashMap, sync::Arc};
6
7use iota_config::genesis::Genesis;
8use iota_sdk_types::{
9    CheckpointDigest, CheckpointSummary, EndOfEpochData,
10    crypto::{Intent, IntentMessage, IntentScope},
11};
12use iota_types::{
13    base_types::{AuthorityName, ExecutionData},
14    committee::{Committee, EpochId, StakeUnit},
15    crypto::{
16        AuthorityKeyPair, AuthoritySignInfo, AuthoritySignature, IotaAuthoritySignature,
17        KeypairTraits,
18    },
19    messages_checkpoint::{
20        CertifiedCheckpointSummary, CheckpointSequenceNumber, CheckpointTimestamp,
21        CheckpointVersionSpecificData, FullCheckpointContents, VerifiedCheckpoint,
22        VerifiedCheckpointContents,
23    },
24};
25
26use crate::network_config::NetworkConfig;
27
28pub struct CommitteeFixture {
29    epoch: EpochId,
30    validators: HashMap<AuthorityName, (AuthorityKeyPair, StakeUnit)>,
31    committee: Committee,
32    genesis: Option<Arc<Genesis>>,
33}
34
35pub type MakeCheckpointResults = (
36    Vec<VerifiedCheckpoint>,
37    Vec<VerifiedCheckpointContents>,
38    HashMap<CheckpointSequenceNumber, CheckpointDigest>,
39    HashMap<CheckpointDigest, VerifiedCheckpoint>,
40);
41
42impl CommitteeFixture {
43    pub fn generate<R: ::rand::RngCore + ::rand::CryptoRng>(
44        mut rng: R,
45        epoch: EpochId,
46        committee_size: usize,
47    ) -> Self {
48        let validators = (0..committee_size)
49            .map(|_| iota_types::crypto::get_key_pair_from_rng::<AuthorityKeyPair, _>(&mut rng).1)
50            .map(|keypair| (keypair.public().into(), (keypair, 1)))
51            .collect::<HashMap<_, _>>();
52
53        let committee = Committee::new_for_testing_with_normalized_voting_power(
54            epoch,
55            validators
56                .iter()
57                .map(|(name, (_, stake))| (*name, *stake))
58                .collect(),
59        );
60
61        Self {
62            epoch,
63            validators,
64            committee,
65            genesis: None,
66        }
67    }
68
69    pub fn from_network_config(network_config: &NetworkConfig) -> Self {
70        let committee = network_config.genesis.committee().unwrap();
71        Self {
72            epoch: committee.epoch,
73            validators: committee
74                .members()
75                .map(|(name, stake)| {
76                    (
77                        *name,
78                        (
79                            network_config
80                                .validator_configs()
81                                .iter()
82                                .find(|config| config.authority_public_key() == *name)
83                                .unwrap()
84                                .authority_key_pair()
85                                .copy(),
86                            *stake,
87                        ),
88                    )
89                })
90                .collect(),
91            committee,
92            genesis: Some(Arc::new(network_config.genesis.clone())),
93        }
94    }
95
96    pub fn committee(&self) -> &Committee {
97        &self.committee
98    }
99
100    fn create_root_checkpoint(&self) -> (VerifiedCheckpoint, VerifiedCheckpointContents) {
101        assert_eq!(self.epoch, 0, "root checkpoint must be epoch 0");
102
103        let mut contents = empty_contents();
104        if let Some(genesis) = &self.genesis {
105            // if genesis is provided, create checkpoint contents with the genesis
106            // transaction
107            let tx = genesis.transaction().clone();
108            let effects = genesis.effects().clone();
109            let execution_data = ExecutionData::new(tx, effects);
110
111            contents = VerifiedCheckpointContents::new_unchecked(
112                FullCheckpointContents::new_with_causally_ordered_transactions(std::iter::once(
113                    execution_data,
114                )),
115            );
116        }
117
118        let contents_digest = contents
119            .clone()
120            .into_inner()
121            .into_checkpoint_contents()
122            .digest();
123
124        let checkpoint = CheckpointSummary {
125            epoch: 0,
126            sequence_number: 0,
127            network_total_transactions: contents.num_of_transactions() as u64,
128            contents_digest,
129            previous_digest: None,
130            epoch_rolling_gas_cost_summary: Default::default(),
131            end_of_epoch_data: None,
132            timestamp_ms: 0,
133            version_specific_data: bcs::to_bytes(&CheckpointVersionSpecificData::empty_for_tests())
134                .unwrap(),
135            checkpoint_commitments: Default::default(),
136        };
137
138        (self.create_certified_checkpoint(checkpoint), contents)
139    }
140
141    fn create_certified_checkpoint(&self, checkpoint: CheckpointSummary) -> VerifiedCheckpoint {
142        let signatures = self
143            .validators
144            .iter()
145            .map(|(name, (key, _))| {
146                let intent_msg = IntentMessage::new(
147                    Intent::iota_app(IntentScope::CheckpointSummary),
148                    checkpoint.clone(),
149                );
150                let signature = AuthoritySignature::new_secure(&intent_msg, &checkpoint.epoch, key);
151                AuthoritySignInfo {
152                    epoch: checkpoint.epoch,
153                    authority: *name,
154                    signature,
155                }
156            })
157            .collect();
158
159        let checkpoint = CertifiedCheckpointSummary::new(checkpoint, signatures, self.committee())
160            .unwrap()
161            .try_into_verified(self.committee())
162            .unwrap();
163
164        checkpoint
165    }
166
167    pub fn make_random_checkpoints(
168        &self,
169        number_of_checkpoints: usize,
170        previous_checkpoint: Option<VerifiedCheckpoint>,
171    ) -> MakeCheckpointResults {
172        self.make_checkpoints(number_of_checkpoints, previous_checkpoint, random_contents)
173    }
174
175    pub fn make_empty_checkpoints(
176        &self,
177        number_of_checkpoints: usize,
178        previous_checkpoint: Option<VerifiedCheckpoint>,
179    ) -> MakeCheckpointResults {
180        self.make_checkpoints(number_of_checkpoints, previous_checkpoint, empty_contents)
181    }
182
183    pub fn make_checkpoints<F: Fn() -> VerifiedCheckpointContents>(
184        &self,
185        number_of_checkpoints: usize,
186        previous_checkpoint: Option<VerifiedCheckpoint>,
187        content_generator: F,
188    ) -> MakeCheckpointResults {
189        // Only skip the first one if it was supplied
190        let skip = previous_checkpoint.is_some() as usize;
191        let first = previous_checkpoint
192            .map(|c| (c, empty_contents()))
193            .unwrap_or_else(|| self.create_root_checkpoint());
194
195        let (ordered_checkpoints, contents): (Vec<_>, Vec<_>) =
196            std::iter::successors(Some(first), |prev| {
197                let contents = content_generator();
198                let contents_digest = contents
199                    .clone()
200                    .into_inner()
201                    .into_checkpoint_contents()
202                    .digest();
203                let summary = CheckpointSummary {
204                    epoch: self.epoch,
205                    sequence_number: prev.0.sequence_number + 1,
206                    network_total_transactions: prev.0.network_total_transactions
207                        + contents.num_of_transactions() as u64,
208                    contents_digest,
209                    previous_digest: Some(*prev.0.digest()),
210                    epoch_rolling_gas_cost_summary: Default::default(),
211                    end_of_epoch_data: None,
212                    timestamp_ms: 0,
213                    version_specific_data: bcs::to_bytes(
214                        &CheckpointVersionSpecificData::empty_for_tests(),
215                    )
216                    .unwrap(),
217                    checkpoint_commitments: Default::default(),
218                };
219
220                let checkpoint = self.create_certified_checkpoint(summary);
221
222                Some((checkpoint, contents))
223            })
224            .skip(skip)
225            .take(number_of_checkpoints)
226            .unzip();
227
228        let (sequence_number_to_digest, checkpoints) = ordered_checkpoints
229            .iter()
230            .cloned()
231            .map(|checkpoint| {
232                let digest = *checkpoint.digest();
233                ((checkpoint.sequence_number, digest), (digest, checkpoint))
234            })
235            .unzip();
236
237        (
238            ordered_checkpoints,
239            contents,
240            sequence_number_to_digest,
241            checkpoints,
242        )
243    }
244
245    /// Builds a chain of empty checkpoints assigning each the given timestamp
246    /// (in order), for tests that need control over checkpoint timestamps.
247    pub fn make_checkpoints_with_timestamps(
248        &self,
249        timestamps_ms: &[CheckpointTimestamp],
250        previous_checkpoint: Option<VerifiedCheckpoint>,
251    ) -> Vec<VerifiedCheckpoint> {
252        let mut prev = previous_checkpoint.unwrap_or_else(|| self.create_root_checkpoint().0);
253        let mut checkpoints = Vec::with_capacity(timestamps_ms.len());
254        for &timestamp_ms in timestamps_ms {
255            let contents_digest = empty_contents()
256                .into_inner()
257                .into_checkpoint_contents()
258                .digest();
259            let summary = CheckpointSummary {
260                epoch: self.epoch,
261                sequence_number: prev.sequence_number + 1,
262                network_total_transactions: prev.network_total_transactions,
263                contents_digest,
264                previous_digest: Some(*prev.digest()),
265                epoch_rolling_gas_cost_summary: Default::default(),
266                end_of_epoch_data: None,
267                timestamp_ms,
268                version_specific_data: bcs::to_bytes(
269                    &CheckpointVersionSpecificData::empty_for_tests(),
270                )
271                .unwrap(),
272                checkpoint_commitments: Default::default(),
273            };
274            let checkpoint = self.create_certified_checkpoint(summary);
275            prev = checkpoint.clone();
276            checkpoints.push(checkpoint);
277        }
278        checkpoints
279    }
280
281    pub fn make_end_of_epoch_checkpoint(
282        &self,
283        previous_checkpoint: VerifiedCheckpoint,
284        end_of_epoch_data: Option<EndOfEpochData>,
285    ) -> (
286        CheckpointSequenceNumber,
287        CheckpointDigest,
288        VerifiedCheckpoint,
289    ) {
290        let summary = CheckpointSummary {
291            epoch: self.epoch,
292            sequence_number: previous_checkpoint.sequence_number + 1,
293            network_total_transactions: 0,
294            contents_digest: empty_contents()
295                .into_inner()
296                .into_checkpoint_contents()
297                .digest(),
298            previous_digest: Some(*previous_checkpoint.digest()),
299            epoch_rolling_gas_cost_summary: Default::default(),
300            end_of_epoch_data,
301            timestamp_ms: 0,
302            version_specific_data: bcs::to_bytes(&CheckpointVersionSpecificData::empty_for_tests())
303                .unwrap(),
304            checkpoint_commitments: Default::default(),
305        };
306
307        let checkpoint = self.create_certified_checkpoint(summary);
308
309        (checkpoint.sequence_number, *checkpoint.digest(), checkpoint)
310    }
311}
312
313pub fn empty_contents() -> VerifiedCheckpointContents {
314    VerifiedCheckpointContents::new_unchecked(
315        FullCheckpointContents::new_with_causally_ordered_transactions(std::iter::empty()),
316    )
317}
318
319pub fn random_contents() -> VerifiedCheckpointContents {
320    VerifiedCheckpointContents::new_unchecked(FullCheckpointContents::random_for_testing())
321}
iota_swarm_config/test_utils.rs

iota_swarm_config/
test_utils.rs
