solana_runtime/
test_utils.rs

1#[cfg(feature = "dev-context-only-utils")]
2use {
3    rand::Rng,
4    solana_account::{AccountSharedData, ReadableAccount, 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<F>(
21    lamports: u64,
22    node_pubkey: Option<Pubkey>,
23    set_bls_pubkey: bool,
24    pending_delegator_rewards_from_account_lamports: F,
25) -> AccountSharedData
26where
27    F: Fn(u64) -> u64,
28{
29    let owner = solana_sdk_ids::vote::id();
30    let mut account = AccountSharedData::new(lamports, VoteStateV4::size_of(), &owner);
31
32    let bls_pubkey_compressed = set_bls_pubkey.then(|| {
33        let bls_pubkey: BLSPubkeyCompressed = (*BLSKeypair::new().public).into();
34        let bls_pubkey_buffer = wincode::serialize(&bls_pubkey).unwrap();
35        bls_pubkey_buffer.try_into().unwrap()
36    });
37    let vote_state = VoteStateV4 {
38        node_pubkey: node_pubkey.unwrap_or(Pubkey::new_unique()),
39        authorized_voters: AuthorizedVoters::new(0, Pubkey::new_unique()),
40        authorized_withdrawer: Pubkey::new_unique(),
41        bls_pubkey_compressed,
42        pending_delegator_rewards: pending_delegator_rewards_from_account_lamports(
43            account.lamports(),
44        ),
45        ..VoteStateV4::default()
46    };
47
48    VoteStateV4::serialize(
49        &VoteStateVersions::V4(Box::new(vote_state)),
50        account.data_as_mut_slice(),
51    )
52    .unwrap();
53    account
54}
55
56#[cfg(feature = "dev-context-only-utils")]
57pub fn new_rand_vote_accounts<R: Rng>(
58    rng: &mut R,
59    num_nodes: usize,
60    max_stake_for_staked_account: u64,
61) -> impl Iterator<Item = (Pubkey, (/*stake:*/ u64, VoteAccount))> + '_ {
62    let nodes: Vec<_> = repeat_with(Pubkey::new_unique).take(num_nodes).collect();
63    repeat_with(move || {
64        let node = nodes[rng.random_range(0..nodes.len())];
65        let account = new_rand_vote_account(rng.random(), Some(node), true, |_| 0);
66        let stake = rng.random_range(0..max_stake_for_staked_account);
67        let vote_account = VoteAccount::try_from(account).unwrap();
68        (Pubkey::new_unique(), (stake, vote_account))
69    })
70}
71
72/// Creates `num_nodes` random vote accounts with the specified stake.
73/// The first `num_nodes_with_bls_pubkeys` have the bls_pubkeys set while the rest are unset.
74/// If `stake_per_node` is specified, then each node will have that stake, otherwise a random amount
75/// between `min_stake_for_staked_account` and `max_stake_for_staked_account` is chosen.
76#[cfg(feature = "dev-context-only-utils")]
77pub fn new_staked_vote_accounts<R: Rng, F, G, H>(
78    rng: &mut R,
79    num_nodes: usize,
80    num_nodes_with_bls_pubkeys: usize,
81    stake_per_node: Option<u64>,
82    min_stake_for_staked_account: u64,
83    max_stake_for_staked_account: u64,
84    lamports_per_node: F,
85    pending_delegator_rewards_from_account_lamports_generator: G,
86) -> VoteAccounts
87where
88    F: Fn(usize) -> u64,
89    G: Fn(usize) -> H, // generator to create a function based on index
90    H: Fn(u64) -> u64, // function returned by generator to set pending rewards
91{
92    let mut vote_accounts = VoteAccounts::default();
93    for index in 0..num_nodes {
94        let pubkey = Pubkey::new_unique();
95        let stake = stake_per_node.unwrap_or_else(|| {
96            rng.random_range(min_stake_for_staked_account..max_stake_for_staked_account)
97        });
98        let node_pubkey = Pubkey::new_unique();
99        let set_bls_pubkey = index < num_nodes_with_bls_pubkeys;
100        let pending_delegator_rewards_fn =
101            pending_delegator_rewards_from_account_lamports_generator(index);
102        let account = new_rand_vote_account(
103            lamports_per_node(index),
104            Some(node_pubkey),
105            set_bls_pubkey,
106            pending_delegator_rewards_fn,
107        );
108        vote_accounts.insert(pubkey, VoteAccount::try_from(account).unwrap(), || stake);
109    }
110    vote_accounts
111}
112
113#[cfg(feature = "dev-context-only-utils")]
114pub fn staked_nodes<'a, I>(vote_accounts: I) -> HashMap<Pubkey, u64>
115where
116    I: IntoIterator<Item = &'a (Pubkey, (u64, VoteAccount))>,
117{
118    let mut staked_nodes = HashMap::new();
119    for (_, (stake, vote_account)) in vote_accounts
120        .into_iter()
121        .filter(|(_, (stake, _))| *stake != 0)
122    {
123        staked_nodes
124            .entry(*vote_account.node_pubkey())
125            .and_modify(|s: &mut u64| *s = s.saturating_add(*stake))
126            .or_insert(*stake);
127    }
128    staked_nodes
129}
solana_runtime/test_utils.rs

solana_runtime/
test_utils.rs
