solana_cli/
test_utils.rs

1use {
2    solana_clock::{Epoch, DEFAULT_MS_PER_SLOT},
3    solana_commitment_config::CommitmentConfig,
4    solana_rpc_client::rpc_client::RpcClient,
5    std::{thread::sleep, time::Duration},
6};
7
8#[macro_export]
9macro_rules! check_balance {
10    ($expected_balance:expr, $client:expr, $pubkey:expr) => {
11        (0..5).for_each(|tries| {
12            let balance = $client
13                .get_balance_with_commitment($pubkey, CommitmentConfig::processed())
14                .unwrap()
15                .value;
16            if balance == $expected_balance {
17                return;
18            }
19            if tries == 4 {
20                assert_eq!(balance, $expected_balance);
21            }
22            std::thread::sleep(std::time::Duration::from_millis(500));
23        });
24    };
25    ($expected_balance:expr, $client:expr, $pubkey:expr,) => {
26        check_balance!($expected_balance, $client, $pubkey)
27    };
28}
29
30pub fn check_ready(rpc_client: &RpcClient) {
31    while rpc_client
32        .get_slot_with_commitment(CommitmentConfig::processed())
33        .unwrap()
34        < 5
35    {
36        sleep(Duration::from_millis(DEFAULT_MS_PER_SLOT));
37    }
38}
39
40pub fn wait_n_slots(rpc_client: &RpcClient, n: u64) -> u64 {
41    let slot = rpc_client.get_slot().unwrap();
42    loop {
43        sleep(Duration::from_millis(DEFAULT_MS_PER_SLOT));
44        let new_slot = rpc_client.get_slot().unwrap();
45        if new_slot.saturating_sub(slot) >= n {
46            return new_slot;
47        }
48    }
49}
50
51pub fn wait_for_next_epoch_plus_n_slots(rpc_client: &RpcClient, n: u64) -> (Epoch, u64) {
52    let current_epoch = rpc_client.get_epoch_info().unwrap().epoch;
53    let next_epoch = current_epoch.saturating_add(1);
54    println!("waiting for epoch {next_epoch} plus {n} slots");
55    loop {
56        sleep(Duration::from_millis(DEFAULT_MS_PER_SLOT));
57
58        let next_epoch = rpc_client.get_epoch_info().unwrap().epoch;
59        if next_epoch > current_epoch {
60            let new_slot = wait_n_slots(rpc_client, n);
61            return (next_epoch, new_slot);
62        }
63    }
64}