zksync_consensus_roles/validator/messages/
testonly.rs

1use rand::{
2    distributions::{Distribution, Standard},
3    Rng,
4};
5use zksync_concurrency::time;
6use zksync_consensus_utils::enum_util::Variant;
7
8use super::{
9    Block, BlockNumber, ChainId, ConsensusMsg, ForkNumber, Genesis, GenesisHash, GenesisRaw,
10    Justification, LeaderSelection, LeaderSelectionMode, Msg, MsgHash, NetAddress, Payload,
11    PayloadHash, PreGenesisBlock, Proposal, ProtocolVersion, ReplicaState, Schedule, Signed,
12    ValidatorInfo, ViewNumber,
13};
14use crate::validator::SecretKey;
15
16impl Distribution<ViewNumber> for Standard {
17    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> ViewNumber {
18        ViewNumber(rng.gen())
19    }
20}
21
22impl Distribution<BlockNumber> for Standard {
23    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> BlockNumber {
24        BlockNumber(rng.gen())
25    }
26}
27
28impl Distribution<ProtocolVersion> for Standard {
29    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> ProtocolVersion {
30        let range = ProtocolVersion::supported_range();
31        ProtocolVersion(rng.gen_range(range))
32    }
33}
34
35impl Distribution<ForkNumber> for Standard {
36    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> ForkNumber {
37        ForkNumber(rng.gen())
38    }
39}
40
41impl Distribution<ValidatorInfo> for Standard {
42    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> ValidatorInfo {
43        ValidatorInfo {
44            key: rng.gen::<SecretKey>().public(),
45            weight: rng.gen_range(1..100),
46            leader: rng.gen_bool(0.7), // 70% chance to be a leader
47        }
48    }
49}
50
51impl Distribution<ChainId> for Standard {
52    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> ChainId {
53        ChainId(rng.gen())
54    }
55}
56
57impl Distribution<Genesis> for Standard {
58    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Genesis {
59        rng.gen::<GenesisRaw>().with_hash()
60    }
61}
62
63impl Distribution<GenesisHash> for Standard {
64    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> GenesisHash {
65        GenesisHash(rng.gen())
66    }
67}
68
69impl Distribution<GenesisRaw> for Standard {
70    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> GenesisRaw {
71        GenesisRaw {
72            chain_id: rng.gen(),
73            fork_number: rng.gen(),
74            first_block: rng.gen(),
75            protocol_version: rng.gen(),
76            validators_schedule: rng.gen(),
77        }
78    }
79}
80
81impl Distribution<Schedule> for Standard {
82    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Schedule {
83        // Create random validators (between 1 and 5)
84        let num_validators = rng.gen_range(1..6);
85
86        // Create validators with 70% chance to be leaders
87        let mut validators: Vec<ValidatorInfo> = (0..num_validators).map(|_| rng.gen()).collect();
88
89        // Ensure at least one validator is a leader
90        if !validators.iter().any(|v| v.leader) {
91            // Make a random validator a leader
92            let leader_idx = rng.gen_range(0..validators.len());
93            validators[leader_idx].leader = true;
94        }
95
96        // This should never fail since we ensure at least one leader
97        Schedule::new(validators, rng.gen()).unwrap()
98    }
99}
100
101impl Distribution<LeaderSelection> for Standard {
102    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> LeaderSelection {
103        LeaderSelection {
104            frequency: rng.gen_range(1..10),
105            mode: if rng.gen_bool(0.5) {
106                LeaderSelectionMode::RoundRobin
107            } else {
108                LeaderSelectionMode::Weighted
109            },
110        }
111    }
112}
113
114impl Distribution<LeaderSelectionMode> for Standard {
115    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> LeaderSelectionMode {
116        match rng.gen_range(0..=2) {
117            0 => LeaderSelectionMode::RoundRobin,
118            1 => LeaderSelectionMode::Sticky(rng.gen()),
119            _ => LeaderSelectionMode::Weighted,
120        }
121    }
122}
123
124impl Distribution<PayloadHash> for Standard {
125    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> PayloadHash {
126        PayloadHash(rng.gen())
127    }
128}
129
130impl Distribution<Payload> for Standard {
131    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Payload {
132        let size: usize = rng.gen_range(500..1000);
133        Payload((0..size).map(|_| rng.gen()).collect())
134    }
135}
136
137impl Distribution<Proposal> for Standard {
138    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Proposal {
139        Proposal {
140            number: rng.gen(),
141            payload: rng.gen(),
142        }
143    }
144}
145
146impl Distribution<Justification> for Standard {
147    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Justification {
148        let size: usize = rng.gen_range(500..1000);
149        Justification((0..size).map(|_| rng.gen()).collect())
150    }
151}
152
153impl Distribution<PreGenesisBlock> for Standard {
154    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> PreGenesisBlock {
155        PreGenesisBlock {
156            number: rng.gen(),
157            payload: rng.gen(),
158            justification: rng.gen(),
159        }
160    }
161}
162
163impl Distribution<Block> for Standard {
164    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Block {
165        match rng.gen_range(0..2) {
166            0 => Block::PreGenesis(rng.gen()),
167            _ => Block::FinalV1(rng.gen()),
168        }
169    }
170}
171
172impl Distribution<NetAddress> for Standard {
173    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> NetAddress {
174        NetAddress {
175            addr: std::net::SocketAddr::new(
176                std::net::IpAddr::from(rng.gen::<[u8; 16]>()),
177                rng.gen(),
178            ),
179            version: rng.gen(),
180            timestamp: time::UNIX_EPOCH + time::Duration::seconds(rng.gen_range(0..1000000000)),
181        }
182    }
183}
184
185impl Distribution<Msg> for Standard {
186    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Msg {
187        match rng.gen_range(0..3) {
188            0 => Msg::Consensus(rng.gen()),
189            1 => Msg::SessionId(rng.gen()),
190            2 => Msg::NetAddress(rng.gen()),
191            _ => unreachable!(),
192        }
193    }
194}
195
196impl Distribution<MsgHash> for Standard {
197    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> MsgHash {
198        MsgHash(rng.gen())
199    }
200}
201
202impl<V: Variant<Msg>> Distribution<Signed<V>> for Standard
203where
204    Standard: Distribution<V>,
205{
206    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Signed<V> {
207        rng.gen::<SecretKey>().sign_msg(rng.gen())
208    }
209}
210
211impl Distribution<ConsensusMsg> for Standard {
212    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> ConsensusMsg {
213        match rng.gen_range(0..4) {
214            0 => ConsensusMsg::LeaderProposal(rng.gen()),
215            1 => ConsensusMsg::ReplicaCommit(rng.gen()),
216            2 => ConsensusMsg::ReplicaNewView(rng.gen()),
217            3 => ConsensusMsg::ReplicaTimeout(rng.gen()),
218            _ => unreachable!(),
219        }
220    }
221}
222
223impl Distribution<ReplicaState> for Standard {
224    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> ReplicaState {
225        ReplicaState {
226            view: rng.gen(),
227            phase: rng.gen(),
228            high_vote: rng.gen(),
229            high_commit_qc: rng.gen(),
230            high_timeout_qc: rng.gen(),
231            proposals: (0..rng.gen_range(1..11)).map(|_| rng.gen()).collect(),
232            v2: rng.gen(),
233        }
234    }
235}
zksync_consensus_roles/validator/messages/testonly.rs

zksync_consensus_roles/validator/messages/
testonly.rs
