Skip to main content

solana_cli/
test_utils.rs

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