solana_runtime/
test_utils.rs

1#[cfg(feature = "dev-context-only-utils")]
2use {
3    rand::Rng,
4    solana_account::{AccountSharedData, WritableAccount},
5    solana_bls_signatures::{
6        keypair::Keypair as BLSKeypair, pubkey::PubkeyCompressed as BLSPubkeyCompressed,
7    },
8    solana_pubkey::Pubkey,
9    solana_vote::vote_account::{VoteAccount, VoteAccounts},
10    solana_vote_interface::{
11        authorized_voters::AuthorizedVoters,
12        state::{VoteStateV4, VoteStateVersions},
13    },
14    std::{collections::HashMap, iter::repeat_with},
15};
16
17/// Creates a vote account
18/// `set_bls_pubkey`: controls whether the bls pubkey is None or Some
19#[cfg(feature = "dev-context-only-utils")]
20pub fn new_rand_vote_account<R: Rng>(
21    rng: &mut R,
22    node_pubkey: Option<Pubkey>,
23    set_bls_pubkey: bool,
24) -> AccountSharedData {
25    let owner = solana_sdk_ids::vote::id();
26    let mut account = AccountSharedData::new(rng.random(), VoteStateV4::size_of(), &owner);
27
28    let bls_pubkey_compressed = set_bls_pubkey.then(|| {
29        let bls_pubkey: BLSPubkeyCompressed = (*BLSKeypair::new().public).into();
30        let bls_pubkey_buffer = bincode::serialize(&bls_pubkey).unwrap();
31        bls_pubkey_buffer.try_into().unwrap()
32    });
33    let vote_state = VoteStateV4 {
34        node_pubkey: node_pubkey.unwrap_or(Pubkey::new_unique()),
35        authorized_voters: AuthorizedVoters::new(0, Pubkey::new_unique()),
36        authorized_withdrawer: Pubkey::new_unique(),
37        bls_pubkey_compressed,
38        ..VoteStateV4::default()
39    };
40
41    VoteStateV4::serialize(
42        &VoteStateVersions::V4(Box::new(vote_state)),
43        account.data_as_mut_slice(),
44    )
45    .unwrap();
46    account
47}
48
49#[cfg(feature = "dev-context-only-utils")]
50pub fn new_rand_vote_accounts<R: Rng>(
51    rng: &mut R,
52    num_nodes: usize,
53    max_stake_for_staked_account: u64,
54) -> impl Iterator<Item = (Pubkey, (/*stake:*/ u64, VoteAccount))> + '_ {
55    let nodes: Vec<_> = repeat_with(Pubkey::new_unique).take(num_nodes).collect();
56    repeat_with(move || {
57        let node = nodes[rng.random_range(0..nodes.len())];
58        let account = new_rand_vote_account(rng, Some(node), true);
59        let stake = rng.random_range(0..max_stake_for_staked_account);
60        let vote_account = VoteAccount::try_from(account).unwrap();
61        (Pubkey::new_unique(), (stake, vote_account))
62    })
63}
64
65/// Creates `num_nodes` random vote accounts with the specified stake.
66/// The first `num_nodes_with_bls_pubkeys` have the bls_pubkeys set while the rest are unset.
67/// If `stake_per_node` is specified, then each node will have that stake, otherwise a random amount
68/// between `min_stake_for_staked_account` and `max_stake_for_staked_account` is chosen.
69#[cfg(feature = "dev-context-only-utils")]
70pub fn new_staked_vote_accounts<R: Rng, F>(
71    rng: &mut R,
72    num_nodes: usize,
73    num_nodes_with_bls_pubkeys: usize,
74    stake_per_node: Option<u64>,
75    min_stake_for_staked_account: u64,
76    max_stake_for_staked_account: u64,
77    lamports_per_node: F,
78) -> VoteAccounts
79where
80    F: Fn(usize) -> u64,
81{
82    let mut vote_accounts = VoteAccounts::default();
83    for index in 0..num_nodes {
84        let pubkey = Pubkey::new_unique();
85        let stake = stake_per_node.unwrap_or_else(|| {
86            rng.random_range(min_stake_for_staked_account..max_stake_for_staked_account)
87        });
88        let node_pubkey = Pubkey::new_unique();
89        let set_bls_pubkey = index < num_nodes_with_bls_pubkeys;
90        let mut account = new_rand_vote_account(rng, Some(node_pubkey), set_bls_pubkey);
91        account.set_lamports(lamports_per_node(index));
92        vote_accounts.insert(pubkey, VoteAccount::try_from(account).unwrap(), || stake);
93    }
94    vote_accounts
95}
96
97#[cfg(feature = "dev-context-only-utils")]
98pub fn staked_nodes<'a, I>(vote_accounts: I) -> HashMap<Pubkey, u64>
99where
100    I: IntoIterator<Item = &'a (Pubkey, (u64, VoteAccount))>,
101{
102    let mut staked_nodes = HashMap::new();
103    for (_, (stake, vote_account)) in vote_accounts
104        .into_iter()
105        .filter(|(_, (stake, _))| *stake != 0)
106    {
107        staked_nodes
108            .entry(*vote_account.node_pubkey())
109            .and_modify(|s: &mut u64| *s = s.saturating_add(*stake))
110            .or_insert(*stake);
111    }
112    staked_nodes
113}
solana_runtime/test_utils.rs

solana_runtime/
test_utils.rs
