wp-solana-test-core 0.1.1

//! SPL-Token helpers for [`TestContext`].
//!
//! Adapted from `solana-test-lib/src/common/svm_setup.rs` to work with the
//! new [`TestContext`] wrapper. All functions take the SVM `std::sync::Mutex`
//! synchronously via `lock()` — critical sections never span `.await`.

use anyhow::Result;
use solana_sdk::{
    account::Account, message::Message, program_option::COption, program_pack::Pack,
    pubkey::Pubkey, signature::Signer, transaction::Transaction,
};
use spl_associated_token_account_interface::{
    address::get_associated_token_address_with_program_id,
    instruction::create_associated_token_account,
};
use spl_token::{
    instruction::{mint_to, sync_native},
    state::Mint,
    ID as TOKEN_PROGRAM_ID,
};

use crate::context::TestContext;

/// Override a SPL-Token mint's `mint_authority` to `new_authority`.
///
/// This is a no-op if the mint account does not exist in the SVM.
pub fn override_mint_authority(
    ctx: &TestContext,
    mint: Pubkey,
    new_authority: Pubkey,
) -> Result<()> {
    let mut svm = ctx.lock_svm();

    let Some(mint_account) = svm.get_account(&mint) else {
        tracing::warn!(%mint, "override_mint_authority: mint account not found in SVM");
        return Ok(());
    };

    let mut mint_state = Mint::unpack(&mint_account.data)
        .map_err(|e| anyhow::anyhow!("Failed to unpack mint: {:?}", e))?;
    mint_state.mint_authority = COption::Some(new_authority);

    let mut mint_data = vec![0u8; mint_account.data.len()];
    mint_state.pack_into_slice(&mut mint_data);

    svm.set_account(
        mint,
        Account {
            lamports: mint_account.lamports,
            data: mint_data,
            owner: mint_account.owner,
            executable: mint_account.executable,
            rent_epoch: mint_account.rent_epoch,
        },
    )?;

    Ok(())
}

/// Create an Associated Token Account for `owner` on `mint`, then mint
/// `amount` tokens into it.
///
/// Returns the ATA address.
pub fn fund_token(ctx: &TestContext, mint: Pubkey, owner: Pubkey, amount: u64) -> Result<Pubkey> {
    let mut svm = ctx.lock_svm();
    let payer = &ctx.payer;

    let ata = get_associated_token_address_with_program_id(&owner, &mint, &TOKEN_PROGRAM_ID);

    if svm.get_account(&ata).is_none() {
        let ix = create_associated_token_account(&payer.pubkey(), &owner, &mint, &TOKEN_PROGRAM_ID);
        let blockhash = svm.latest_blockhash();
        let msg = Message::new(&[ix], Some(&payer.pubkey()));
        let tx = Transaction::new(&[payer.as_ref()], msg, blockhash);
        svm.send_transaction(tx).map_err(|e| anyhow::anyhow!("Failed to create ATA: {:?}", e))?;
    }

    let mint_ix = mint_to(&TOKEN_PROGRAM_ID, &mint, &ata, &payer.pubkey(), &[], amount)?;
    let blockhash = svm.latest_blockhash();
    let msg = Message::new(&[mint_ix], Some(&payer.pubkey()));
    let tx = Transaction::new(&[payer.as_ref()], msg, blockhash);
    svm.send_transaction(tx).map_err(|e| anyhow::anyhow!("Failed to mint tokens: {:?}", e))?;

    Ok(ata)
}

/// Create an Associated Token Account for WSOL and wrap `amount` lamports
/// of SOL into it via `sync_native`.
///
/// Returns the ATA address.
pub fn fund_wsol(ctx: &TestContext, owner: Pubkey, amount: u64) -> Result<Pubkey> {
    let mut svm = ctx.lock_svm();
    let payer = &ctx.payer;
    let wsol_mint = spl_token::native_mint::id();

    let ata = get_associated_token_address_with_program_id(&owner, &wsol_mint, &TOKEN_PROGRAM_ID);

    if svm.get_account(&ata).is_none() {
        let ix =
            create_associated_token_account(&payer.pubkey(), &owner, &wsol_mint, &TOKEN_PROGRAM_ID);
        let blockhash = svm.latest_blockhash();
        let msg = Message::new(&[ix], Some(&payer.pubkey()));
        let tx = Transaction::new(&[payer.as_ref()], msg, blockhash);
        svm.send_transaction(tx)
            .map_err(|e| anyhow::anyhow!("Failed to create WSOL ATA: {:?}", e))?;
    }

    // Add lamports directly to the token account (test-only shortcut).
    let token_account =
        svm.get_account(&ata).ok_or_else(|| anyhow::anyhow!("Failed to get WSOL token account"))?;

    svm.set_account(ata, Account { lamports: token_account.lamports + amount, ..token_account })?;

    let sync_ix = sync_native(&TOKEN_PROGRAM_ID, &ata)?;
    let blockhash = svm.latest_blockhash();
    let mut tx = Transaction::new_with_payer(&[sync_ix], Some(&payer.pubkey()));
    tx.sign(&[payer.as_ref()], blockhash);
    svm.send_transaction(tx).map_err(|e| anyhow::anyhow!("Failed to sync WSOL: {:?}", e))?;

    Ok(ata)
}