Setting up integration testing

Let's take a look at how to write integration tests using the solana-program-test framework and bonfida-test-utils.

The project already contains a test template at tests/functional.rs. A functional test is a very simple integration test which tests basic program functionality. In real-world applications, we would require a battery of such tests to cover edge cases. For the sake of brevity and since this program is quite straightforward as it is, let's focus our efforts on a single test.

Test scenario

We want to test for a basic test scenario. Alice wants to vest some tokens for Bob. This means that Alice will create a vesting contract with Bob set as the recipient. The vesting schedule will be composed of three unlock phases.

To set up this scenario, we need to create a token with a given mint authority, and then mint some tokens for Alice. We can then have Alice create the vesting contract. Finally, Bob will attempt to unlock tokens at different simulated times and we will check that the unlocked amounts are valid.

Writing the test

Let's start by clearing the contents of the test_01 function. We can then set up the identities (keypairs) that we will be requiring. Using named indices instead of variables can make things neater.

use bonfida_test_utils::{ProgramTestContextExt, ProgramTestExt};
use solana_program::pubkey::Pubkey;
use token_vesting::{
    entrypoint::process_instruction,
    state::vesting_contract::{VestingContract, VestingSchedule},
};

use {
    solana_program_test::{processor, ProgramTest},
    solana_sdk::signer::{keypair::Keypair, Signer},
};

#[tokio::test]
async fn test_01() {
    // Create program and test environment
    const ALICE: usize = 0;
    const BOB: usize = 1;
    const MINT_AUTHORITY: usize = 2;

    // We'll need this later when writing and testing our schedules
    const SECONDS_IN_HOUR: u64 = 3600;

    let keypairs = [Keypair::new(), Keypair::new(), Keypair::new()];

    let mut program_test = ProgramTest::new(
        "token_vesting",
        token_vesting::ID,
        processor!(process_instruction),
    );
}

Then, we can initialize a new token mint with 6 decimals (this is mostly standard). The add_mint method comes from the ProgramTestExt trait included with bonfida-test-utils.

let (mint_key, _) = program_test.add_mint(None, 6, &keypairs[MINT_AUTHORITY].pubkey());

Once program environment setup is done, we can start the test runtime.

////
// Create test context
////
let mut prg_test_ctx = program_test.start_with_context().await;

We then need to initialize Alice and Bob's token accounts. We'll also mint some tokens into Alice's account, this time thanks to the ProgramTestContextExt trait. The initialize_token_accounts method will initialize associated token accounts tied to the given keys.

// Initialize Alice and Bob's token accounts:
let ata_keys = prg_test_ctx
    .initialize_token_accounts(
        mint_key,
        &keypairs.iter().map(|k| k.pubkey()).collect::<Vec<_>>(),
    )
    .await
    .unwrap();

// Alice gets 100 tokens
prg_test_ctx
    .mint_tokens(
        &keypairs[MINT_AUTHORITY],
        &mint_key,
        &ata_keys[ALICE],
        100_000_000, // This is 100 actual tokens since the token is defined with 6 decimals
    )
    .await
    .unwrap();

We can then define the vesting schedule.

// Alice will vest 16 tokens for Bob
// We first define the schedule we want

let now = prg_test_ctx.get_current_timestamp().await.unwrap() as u64;

let schedule = vec![
    VestingSchedule {
        unlock_timestamp: now + SECONDS_IN_HOUR,
        quantity: 10_000_000,
    },
    VestingSchedule {
        unlock_timestamp: now + 2 * SECONDS_IN_HOUR,
        quantity: 5_000_000,
    },
    VestingSchedule {
        unlock_timestamp: now + 3 * SECONDS_IN_HOUR,
        quantity: 1_000_000,
    },
];

In preparation for the token vesting contract initialization, we need to allocate the contract account. This is where VestingContract's compute_allocation_size method comes in handy.

let allocation_size = VestingContract::compute_allocation_size(schedule.len());
let vesting_contract = prg_test_ctx
    .initialize_new_account(allocation_size, token_vesting::ID)
    .await
    .unwrap();

We then need to initialize the vesting contract's token account. The token account will belong to the vesting contract's vault signer, which we need to derive first.

let (vault_signer, vault_signer_nonce) =
    Pubkey::find_program_address(&[&vesting_contract.to_bytes()], &token_vesting::ID);
let vault = prg_test_ctx
    .initialize_token_accounts(mint_key, &[vault_signer])
    .await
    .unwrap()[0];

We can finally create the vesting contract.

let ix = token_vesting::instruction::create(
    token_vesting::instruction::create::Accounts {
        spl_token_program: &spl_token::ID,
        vesting_contract: &vesting_contract,
        vault: &vault,
        source_tokens: &ata_keys[ALICE],
        source_tokens_owner: &keypairs[ALICE].pubkey(),
        recipient: &keypairs[BOB].pubkey(),
    },
    token_vesting::instruction::create::Params {
        signer_nonce: &(vault_signer_nonce as u64),
        schedule: &schedule,
    },
);

prg_test_ctx
    .sign_send_instructions(&[ix], &[&keypairs[ALICE]])
    .await
    .unwrap();

Once this is done, we can check that the tokens have been successfully debited from Alice's account.

let alice_token_account_balance = prg_test_ctx
    .get_token_account(ata_keys[ALICE])
    .await
    .unwrap()
    .amount;
assert_eq!(alice_token_account_balance, 84_000_000);

Already, we can run this test using cargo test-bpf and check that it successfully passes.

The next step is to make Bob claim the schedules one by one. We're going to be using the warp_to_timestamp method to simulate the passage of time in order to test the scheduling logic.

for v in schedule {
    // We fast-forward to the unlock
    let previous_balance = prg_test_ctx
        .get_token_account(ata_keys[BOB])
        .await
        .unwrap()
        .amount;
    prg_test_ctx
        .warp_to_timestamp(v.unlock_timestamp as i64)
        .await
        .unwrap();
    let ix = token_vesting::instruction::claim(
        token_vesting::instruction::claim::Accounts {
            spl_token_program: &spl_token::ID,
            vesting_contract: &vesting_contract,
            vesting_contract_signer: &vault_signer,
            vault: &vault,
            destination_token_account: &ata_keys[BOB],
            owner: &keypairs[BOB].pubkey(),
        },
        token_vesting::instruction::claim::Params {},
    );

    prg_test_ctx
        .sign_send_instructions(&[ix], &[&keypairs[BOB]])
        .await
        .unwrap();

    // We check that the tokens have been properly unvested
    let bob_token_account_balance = prg_test_ctx
        .get_token_account(ata_keys[BOB])
        .await
        .unwrap()
        .amount;
    assert_eq!(bob_token_account_balance - previous_balance, v.quantity);
}

We can run this test and check that it indeed passes.

Conclusion

With this, we have successfully tested our program's overall functionality! The next step would be to write another test in which Bob attempts to unlock the tokens at the wrong times. We could also add in tests in which the provided accounts for each instruction are actually incorrect. The idea is to think about attack vectors and then write integration tests for those attack vectors. If at anytime in the future, a change opens up a vulnerability, it should ideally be directly caught by your test suite.

Writing tests with good coverage is a very generic and hard problem to solve in software engineering, so it is definitely out of the scope of this tutorial.