spl-single-pool 6.0.0

#![allow(clippy::arithmetic_side_effects)]

mod helpers;

use {
    helpers::*,
    solana_clock::Clock,
    solana_program_test::*,
    solana_signer::Signer,
    solana_stake_interface::{
        instruction as stake_instruction, stake_history::StakeHistory, state::StakeStateV2,
    },
    solana_transaction::Transaction,
    spl_single_pool::{error::SinglePoolError, id, instruction},
    test_case::test_matrix,
};

#[test_matrix(
    [StakeProgramVersion::Stable, StakeProgramVersion::Beta, StakeProgramVersion::Edge],
    [false, true],
    [false, true]
)]
#[tokio::test]
async fn reactivate_success(
    stake_version: StakeProgramVersion,
    reactivate_pool: bool,
    fund_onramp: bool,
) {
    let Some(program_test) = program_test(stake_version) else {
        return;
    };
    let mut context = program_test.start_with_context().await;

    let accounts = SinglePoolAccounts::default();
    accounts
        .initialize_for_deposit(&mut context, TEST_STAKE_AMOUNT, Some(TEST_STAKE_AMOUNT))
        .await;

    let transaction = Transaction::new_signed_with_payer(
        &[stake_instruction::deactivate_stake(
            &accounts.bob_stake.pubkey(),
            &accounts.bob.pubkey(),
        )],
        Some(&context.payer.pubkey()),
        &[&context.payer, &accounts.bob],
        context.last_blockhash,
    );

    context
        .banks_client
        .process_transaction(transaction)
        .await
        .unwrap();

    advance_epoch(&mut context).await;

    if reactivate_pool {
        // deactivate the pool stake account
        force_deactivate_stake_account(&mut context, &accounts.stake_account).await;

        // active deposit into deactivated pool fails
        let instructions = instruction::deposit(
            &id(),
            &accounts.pool,
            &accounts.alice_stake.pubkey(),
            &accounts.alice_token,
            &accounts.alice.pubkey(),
            &accounts.alice.pubkey(),
        );
        let transaction = Transaction::new_signed_with_payer(
            &instructions,
            Some(&context.payer.pubkey()),
            &[&context.payer, &accounts.alice],
            context.last_blockhash,
        );

        let e = context
            .banks_client
            .process_transaction(transaction)
            .await
            .unwrap_err();
        check_error(e, SinglePoolError::ReplenishRequired);

        // inactive deposit into deactivated pool fails
        let instructions = instruction::deposit(
            &id(),
            &accounts.pool,
            &accounts.bob_stake.pubkey(),
            &accounts.bob_token,
            &accounts.bob.pubkey(),
            &accounts.bob.pubkey(),
        );
        let transaction = Transaction::new_signed_with_payer(
            &instructions,
            Some(&context.payer.pubkey()),
            &[&context.payer, &accounts.bob],
            context.last_blockhash,
        );

        let e = context
            .banks_client
            .process_transaction(transaction)
            .await
            .unwrap_err();
        check_error(e, SinglePoolError::ReplenishRequired);
    }

    // onramp is already inactive but it doesnt have lamports for delegation
    if fund_onramp {
        let minimum_delegation = get_minimum_delegation(
            &mut context.banks_client,
            &context.payer,
            &context.last_blockhash,
        )
        .await;

        transfer(
            &mut context.banks_client,
            &context.payer,
            &context.last_blockhash,
            &accounts.onramp_account,
            minimum_delegation,
        )
        .await;
    }

    // replenish and advance
    replenish(&mut context, &accounts.vote_account.pubkey()).await;
    advance_epoch(&mut context).await;

    // deposit works in all cases
    let instructions = instruction::deposit(
        &id(),
        &accounts.pool,
        &accounts.alice_stake.pubkey(),
        &accounts.alice_token,
        &accounts.alice.pubkey(),
        &accounts.alice.pubkey(),
    );
    let transaction = Transaction::new_signed_with_payer(
        &instructions,
        Some(&context.payer.pubkey()),
        &[&context.payer, &accounts.alice],
        context.last_blockhash,
    );

    context
        .banks_client
        .process_transaction(transaction)
        .await
        .unwrap();

    assert!(context
        .banks_client
        .get_account(accounts.alice_stake.pubkey())
        .await
        .expect("get_account")
        .is_none());

    // onramp is active now if we transfered to it
    let clock = context.banks_client.get_sysvar::<Clock>().await.unwrap();
    let (_, onramp_stake, _) =
        get_stake_account(&mut context.banks_client, &accounts.onramp_account).await;

    // we require a fully active pool for any onramp state change to reduce complexity
    // the pool for a healthy validator is never unstaked, and a fresh pool you can wait an epoch
    // NOTE we might relax this for DepositSol, in which case this test would change
    if fund_onramp && !reactivate_pool {
        let stake = onramp_stake.unwrap();
        assert_eq!(stake.delegation.activation_epoch, clock.epoch - 1);
        assert_eq!(stake.delegation.deactivation_epoch, u64::MAX);
    } else {
        assert_eq!(onramp_stake, None);
    }
}

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
enum OnRampState {
    Initialized,
    Activating,
    Active,
    Deactive,
}

#[derive(Clone, Copy, PartialEq, Eq)]
enum MoveLamportsCase {
    None,
    SubMinumumDelegation,
    MinimumDelegation,
}

// initialized/move: onramp is fresh from its last replenish, move lamports from pool, and delegate if minimum
// activating/no move: onramp has loose lamports, add them to its delegation unconditionally
// activating/move: pool has loose lamports, move and add them to onramp delegation unconditionally
// active/no move: move onramp stake to pool, onramp ends initialized with rent exemption only
// active/move: move onramp stake to pool, move lamports from pool, and delegate if minimum
// deactive/move: same as the initialized/move case, but if onramp was deactivated
// initialized/no move and deactive/no move: nothing happens. successfully!
// ALL variants return Ok(()). replenish is a "magic" instruction that only performs stake operations which should succeed
#[test_matrix(
    [StakeProgramVersion::Stable, StakeProgramVersion::Beta, StakeProgramVersion::Edge],
    [OnRampState::Initialized, OnRampState::Activating, OnRampState::Active, OnRampState::Deactive],
    [MoveLamportsCase::None, MoveLamportsCase::SubMinumumDelegation, MoveLamportsCase::MinimumDelegation]
)]
#[tokio::test]
async fn move_value_success(
    stake_version: StakeProgramVersion,
    onramp_state: OnRampState,
    move_lamports_case: MoveLamportsCase,
) {
    let Some(program_test) = program_test(stake_version) else {
        return;
    };
    let mut context = program_test.start_with_context().await;

    let rent = context.banks_client.get_rent().await.unwrap();
    let pool_rent = rent.minimum_balance(StakeStateV2::size_of());
    let onramp_rent = pool_rent;

    let accounts = SinglePoolAccounts::default();
    accounts
        .initialize_for_deposit(&mut context, TEST_STAKE_AMOUNT, None)
        .await;
    advance_epoch(&mut context).await;

    // active onramp can be as low as minimum_delegation
    let minimum_delegation = get_minimum_delegation(
        &mut context.banks_client,
        &context.payer,
        &context.last_blockhash,
    )
    .await;

    // if we dont move, we dont move; instruction succeeds
    // if we move less than minimum, we move and dont delegate; instruction succeeds
    // if we move the minimum, we move and delegate; instruction succeeds
    let move_lamports_to_onramp = match move_lamports_case {
        MoveLamportsCase::None => None,
        MoveLamportsCase::SubMinumumDelegation => Some(minimum_delegation - 1),
        MoveLamportsCase::MinimumDelegation => Some(minimum_delegation),
    };

    let minimum_pool_balance = get_minimum_pool_balance(
        &mut context.banks_client,
        &context.payer,
        &context.last_blockhash,
    )
    .await;

    // set up an activating onramp
    if onramp_state >= OnRampState::Activating {
        transfer(
            &mut context.banks_client,
            &context.payer,
            &context.last_blockhash,
            &accounts.onramp_account,
            minimum_delegation,
        )
        .await;

        replenish(&mut context, &accounts.vote_account.pubkey()).await;
    }

    // allow the delegation to activate
    if onramp_state >= OnRampState::Active {
        advance_epoch(&mut context).await;
    }

    // move it over; this case is inactive and behaves identical to Initialized
    if onramp_state == OnRampState::Deactive {
        replenish(&mut context, &accounts.vote_account.pubkey()).await;
    }

    // if we are testing the pool -> onramp leg, add lamports for it
    if let Some(move_lamports_amount) = move_lamports_to_onramp {
        transfer(
            &mut context.banks_client,
            &context.payer,
            &context.last_blockhash,
            &accounts.stake_account,
            move_lamports_amount,
        )
        .await;
    };

    // this one case is to test reupping an activating delegation in place
    let existing_onramp_lamports =
        if onramp_state == OnRampState::Activating && move_lamports_to_onramp.is_none() {
            // NOTE we want a unique value to be sure our test cases are right and this didnt come from pool.
            // we also prefer sub-minimal, so we know we can add sub-minimal lamports to an already sufficient activation.
            // however stake v4 minimum is 1. this comment can be removed when stake v5 is Stable,
            // and the value can be changed to simply `minimum_delegation - 2`
            let existing_onramp_lamports = if minimum_delegation == 1 {
                minimum_delegation + 2
            } else {
                minimum_delegation - 2
            };

            transfer(
                &mut context.banks_client,
                &context.payer,
                &context.last_blockhash,
                &accounts.onramp_account,
                existing_onramp_lamports,
            )
            .await;

            existing_onramp_lamports
        } else {
            0
        };

    // this is the replenish we test
    replenish(&mut context, &accounts.vote_account.pubkey()).await;

    let clock = context.banks_client.get_sysvar::<Clock>().await.unwrap();
    let stake_history = context
        .banks_client
        .get_sysvar::<StakeHistory>()
        .await
        .unwrap();

    let (_, pool_stake, pool_lamports) =
        get_stake_account(&mut context.banks_client, &accounts.stake_account).await;
    let pool_status = pool_stake
        .unwrap()
        .delegation
        .stake_activating_and_deactivating(clock.epoch, &stake_history, Some(0));

    let (_, onramp_stake, onramp_lamports) =
        get_stake_account(&mut context.banks_client, &accounts.onramp_account).await;
    let onramp_status = onramp_stake
        .map(|stake| {
            stake
                .delegation
                .stake_activating_and_deactivating(clock.epoch, &stake_history, Some(0))
        })
        .unwrap_or_default();

    match (onramp_state, move_lamports_to_onramp) {
        // stake moved already before test or because of test, new lamports were added to onramp
        (OnRampState::Deactive, Some(move_lamports_amount))
        | (OnRampState::Active, Some(move_lamports_amount)) => {
            assert_eq!(
                pool_status.effective,
                minimum_pool_balance + minimum_delegation
            );
            assert_eq!(
                pool_lamports,
                minimum_pool_balance + minimum_delegation + pool_rent
            );

            assert_eq!(onramp_status.effective, 0);
            assert_eq!(
                onramp_status.activating,
                if move_lamports_amount >= minimum_delegation {
                    move_lamports_amount
                } else {
                    0
                }
            );
            assert_eq!(onramp_lamports, move_lamports_amount + onramp_rent);
        }
        // no stake moved, but lamports did
        (OnRampState::Initialized, Some(move_lamports_amount)) => {
            assert_eq!(pool_status.effective, minimum_pool_balance);
            assert_eq!(pool_lamports, minimum_pool_balance + pool_rent);

            assert_eq!(onramp_status.effective, 0);
            assert_eq!(
                onramp_status.activating,
                if move_lamports_amount >= minimum_delegation {
                    move_lamports_amount
                } else {
                    0
                }
            );
            assert_eq!(onramp_lamports, move_lamports_amount + onramp_rent);
        }
        // no excess lamports moved, just stake
        (OnRampState::Active, None) => {
            assert_eq!(
                pool_status.effective,
                minimum_pool_balance + minimum_delegation
            );
            assert_eq!(
                pool_lamports,
                minimum_pool_balance + minimum_delegation + pool_rent
            );

            assert_eq!(onramp_status.effective, 0);
            assert_eq!(onramp_status.activating, 0);
            assert_eq!(onramp_lamports, onramp_rent);
        }
        // topped up an existing activation by moving from pool
        (OnRampState::Activating, Some(move_lamports_amount)) => {
            assert_eq!(pool_status.effective, minimum_pool_balance);
            assert_eq!(pool_lamports, minimum_pool_balance + pool_rent);

            assert_eq!(onramp_status.effective, 0);
            assert_eq!(
                onramp_status.activating,
                move_lamports_amount + minimum_delegation
            );
            assert_eq!(
                onramp_lamports,
                move_lamports_amount + minimum_delegation + onramp_rent
            );
        }
        // topped up an existing activation with lamports already in onramp
        (OnRampState::Activating, None) => {
            assert!(existing_onramp_lamports > 0);

            assert_eq!(pool_status.effective, minimum_pool_balance);
            assert_eq!(pool_lamports, minimum_pool_balance + pool_rent);

            assert_eq!(onramp_status.effective, 0);
            assert_eq!(
                onramp_status.activating,
                existing_onramp_lamports + minimum_delegation
            );
            assert_eq!(
                onramp_lamports,
                existing_onramp_lamports + minimum_delegation + onramp_rent
            );
        }
        // nothing happened
        (OnRampState::Initialized, None) => {
            assert_eq!(pool_status.effective, minimum_pool_balance);
            assert_eq!(pool_lamports, minimum_pool_balance + pool_rent);

            assert_eq!(onramp_status.effective, 0);
            assert_eq!(onramp_status.activating, 0);
            assert_eq!(onramp_lamports, onramp_rent);
        }
        // still nothing happened but setting up a deactive onramp added stake to pool already
        (OnRampState::Deactive, None) => {
            assert_eq!(
                pool_status.effective,
                minimum_delegation + minimum_pool_balance
            );
            assert_eq!(
                pool_lamports,
                minimum_delegation + minimum_pool_balance + pool_rent
            );

            assert_eq!(onramp_status.effective, 0);
            assert_eq!(onramp_status.activating, 0);
            assert_eq!(onramp_lamports, onramp_rent);
        }
    }
}

#[test_matrix(
    [StakeProgramVersion::Stable, StakeProgramVersion::Beta, StakeProgramVersion::Edge]
)]
#[tokio::test]
async fn move_value_deactivating(stake_version: StakeProgramVersion) {
    let Some(program_test) = program_test(stake_version) else {
        return;
    };
    let mut context = program_test.start_with_context().await;

    let accounts = SinglePoolAccounts::default();
    accounts
        .initialize_for_deposit(&mut context, TEST_STAKE_AMOUNT, None)
        .await;
    advance_epoch(&mut context).await;

    let minimum_delegation = get_minimum_delegation(
        &mut context.banks_client,
        &context.payer,
        &context.last_blockhash,
    )
    .await;

    transfer(
        &mut context.banks_client,
        &context.payer,
        &context.last_blockhash,
        &accounts.onramp_account,
        minimum_delegation,
    )
    .await;

    // set up a real active onramp
    replenish(&mut context, &accounts.vote_account.pubkey()).await;
    advance_epoch(&mut context).await;

    // edit the account to be deactivating instead
    let clock = context.banks_client.get_sysvar::<Clock>().await.unwrap();
    let mut onramp_account = get_account(&mut context.banks_client, &accounts.onramp_account).await;
    let mut onramp_data: StakeStateV2 = bincode::deserialize(&onramp_account.data).unwrap();

    match onramp_data {
        StakeStateV2::Stake(_, ref mut stake, _) => {
            stake.delegation.deactivation_epoch = clock.epoch
        }
        _ => unreachable!(),
    }

    onramp_account.data = bincode::serialize(&onramp_data).unwrap();
    context.set_account(&accounts.onramp_account, &onramp_account.into());

    // this replenish call reactivates it in the same epoch
    replenish(&mut context, &accounts.vote_account.pubkey()).await;

    let (_, Some(stake), _) =
        get_stake_account(&mut context.banks_client, &accounts.onramp_account).await
    else {
        unreachable!()
    };

    assert_eq!(stake.delegation.deactivation_epoch, u64::MAX);
}

#[test_matrix(
    [StakeProgramVersion::Stable, StakeProgramVersion::Beta, StakeProgramVersion::Edge],
    [false, true]
)]
#[tokio::test]
async fn fail_onramp_doesnt_exist(stake_version: StakeProgramVersion, activate: bool) {
    let Some(program_test) = program_test(stake_version) else {
        return;
    };
    let mut context = program_test.start_with_context().await;

    let accounts = SinglePoolAccounts::default();

    let slot = context.genesis_config().epoch_schedule.first_normal_slot + 1;
    context.warp_to_slot(slot).unwrap();

    create_vote(
        &mut context.banks_client,
        &context.payer,
        &context.last_blockhash,
        &accounts.validator,
        &accounts.voter.pubkey(),
        &accounts.withdrawer.pubkey(),
        &accounts.vote_account,
    )
    .await;

    let rent = context.banks_client.get_rent().await.unwrap();
    let minimum_pool_balance = get_minimum_pool_balance(
        &mut context.banks_client,
        &context.payer,
        &context.last_blockhash,
    )
    .await;

    let mut instructions = instruction::initialize(
        &id(),
        &accounts.vote_account.pubkey(),
        &context.payer.pubkey(),
        &rent,
        minimum_pool_balance,
    );

    // guard against instruction moving in the builder function
    assert_eq!(&instructions[5].data, &[6]);
    let onramp_instruction = instructions.remove(5);

    let transaction = Transaction::new_signed_with_payer(
        &instructions,
        Some(&context.payer.pubkey()),
        &[&context.payer],
        context.last_blockhash,
    );

    context
        .banks_client
        .process_transaction(transaction)
        .await
        .unwrap();

    if activate {
        advance_epoch(&mut context).await;
    }

    // pool is now activating or active with no onramp account
    // replenish should fail because the onramp is required
    let replenish_instruction = instruction::replenish_pool(&id(), &accounts.vote_account.pubkey());
    let transaction = Transaction::new_signed_with_payer(
        core::slice::from_ref(&replenish_instruction),
        Some(&context.payer.pubkey()),
        &[&context.payer],
        context.last_blockhash,
    );

    let e = context
        .banks_client
        .process_transaction(transaction)
        .await
        .unwrap_err();
    check_error(e, SinglePoolError::OnRampDoesntExist);

    // creating onramp lets replenish succeed in the same epoch
    let transaction = Transaction::new_signed_with_payer(
        &[onramp_instruction],
        Some(&context.payer.pubkey()),
        &[&context.payer],
        context.last_blockhash,
    );

    context
        .banks_client
        .process_transaction(transaction)
        .await
        .unwrap();

    refresh_blockhash(&mut context).await;
    let transaction = Transaction::new_signed_with_payer(
        &[replenish_instruction],
        Some(&context.payer.pubkey()),
        &[&context.payer],
        context.last_blockhash,
    );

    context
        .banks_client
        .process_transaction(transaction)
        .await
        .unwrap();
}