wp-solana-test-utils 0.1.1

//! Whirlpool online fixture setup (requires RPC)

use std::str::FromStr;

use anyhow::Result;
use solana_client::nonblocking::rpc_client::RpcClient;
use solana_sdk::{account::Account, pubkey::Pubkey};
use wp_solana_orca_whirlpool_client::{
    constants::TICK_ARRAY_SIZE,
    generated::accounts::Whirlpool,
    pda::tick_array::{get_tick_array_address, get_tick_array_start_tick_index},
};
use wp_solana_test_core::{new_test_context, set_account, TestContext};

use super::{offline::WhirlpoolAccounts, types::whirlpool_program_id};
use crate::internal::{fetch_and_add_program, setup_mint_and_fund};

/// Setup Whirlpool test fixture with LiteSVM, pool, USDC mint.
///
/// Fetches live state from RPC.
pub async fn setup_whirlpool_fixture_online(
    pool_address: Option<Pubkey>,
) -> Result<(TestContext, WhirlpoolAccounts)> {
    let ctx = new_test_context()?;

    let program_id = whirlpool_program_id();

    let rpc_url = std::env::var("SOLANA_RPC_URL")
        .unwrap_or_else(|_| "https://api.mainnet-beta.solana.com".to_string());
    let client = RpcClient::new(rpc_url);

    tracing::info!("Fetching Whirlpool program from RPC");
    fetch_and_add_program(&ctx, &client, program_id, "target/deploy/whirlpool.so").await?;

    let pool_address = pool_address.unwrap_or_else(|| {
        Pubkey::from_str("Czfq3xZZDmsdGdUyrNLtRhGc47cXcZtLG4crryfu44zE")
            .unwrap_or_else(|_| Pubkey::new_unique())
    });

    const USDC_MINT: &str = "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v";
    let usdc_mint = Pubkey::from_str(USDC_MINT)
        .map_err(|e| anyhow::anyhow!("Invalid USDC mint address: {}", e))?;

    const WSOL_MINT: &str = "So11111111111111111111111111111111111111112";
    let wsol_mint = Pubkey::from_str(WSOL_MINT)
        .map_err(|e| anyhow::anyhow!("Invalid WSOL mint address: {}", e))?;

    tracing::info!(
        %pool_address,
        "Fetching pool, USDC mint, and WSOL mint accounts in parallel"
    );
    let (pool_account_result, usdc_mint_account_result, wsol_mint_account_result) = tokio::join!(
        client.get_account(&pool_address),
        client.get_account(&usdc_mint),
        client.get_account(&wsol_mint),
    );

    let pool_account =
        pool_account_result.map_err(|e| anyhow::anyhow!("Failed to fetch pool account: {}", e))?;
    let usdc_mint_account = usdc_mint_account_result
        .map_err(|e| anyhow::anyhow!("Failed to fetch USDC mint account: {}", e))?;
    let wsol_mint_account = wsol_mint_account_result
        .map_err(|e| anyhow::anyhow!("Failed to fetch WSOL mint account: {}", e))?;

    let pool_account_data = pool_account.data.clone();

    let whirlpool_parsed = Whirlpool::from_bytes(pool_account_data.as_slice())
        .map_err(|e| anyhow::anyhow!("Failed to parse Whirlpool pool account: {:?}", e))?;

    set_account(&ctx, pool_address, pool_account)?;
    tracing::info!("Pool account added successfully");

    tracing::info!(
        program_id = %program_id,
        %pool_address,
        "Program is ready for execution"
    );

    // Setup USDC mint, override authority, and fund payer
    tracing::info!("Setting up USDC mint and token account");
    let usdc_amount: u64 = 1_000_000_000_000_000;
    let payer_usdc_ata =
        setup_mint_and_fund(&ctx, usdc_mint, usdc_mint_account, false, usdc_amount)?;

    // Setup WSOL mint and fund payer
    tracing::info!("Setting up WSOL mint and token account");
    let wsol_amount: u64 = 10_000_000_000;
    let payer_wsol_ata =
        setup_mint_and_fund(&ctx, wsol_mint, wsol_mint_account, true, wsol_amount)?;

    tracing::info!(
        usdc_mint = %usdc_mint,
        payer_usdc_ata = %payer_usdc_ata,
        wsol_mint = %wsol_mint,
        payer_wsol_ata = %payer_wsol_ata,
        "Setup complete"
    );

    // Fetch reward token mints (up to 3, only non-default).
    //
    // The harvest fetch path calls `get_multiple_accounts(addresses=[ticks, mints,
    // reward_mints])` and surfaces `AccountNotFound` if any reward mint is
    // absent from the SVM. Vault and authority pubkeys live inside on-chain
    // instructions only and do not need to be present in the SVM for the fetch
    // layer.
    let reward_mints: Vec<Pubkey> = whirlpool_parsed
        .reward_infos
        .iter()
        .filter(|ri| ri.mint != Pubkey::default())
        .map(|ri| ri.mint)
        .collect();

    // Reward vaults needed by close_position (CollectRewardV2 instruction).
    let reward_vaults: Vec<Pubkey> = whirlpool_parsed
        .reward_infos
        .iter()
        .filter(|ri| ri.mint != Pubkey::default())
        .map(|ri| ri.vault)
        .collect();

    if !reward_mints.is_empty() {
        tracing::info!(count = reward_mints.len(), "Fetching reward mint accounts");
        let reward_futures: Vec<_> = reward_mints.iter().map(|m| client.get_account(m)).collect();
        let reward_results = futures::future::join_all(reward_futures).await;
        for (mint, result) in reward_mints.iter().zip(reward_results) {
            match result {
                Ok(account) => {
                    set_account(&ctx, *mint, account)?;
                    tracing::info!(%mint, "Reward mint added");
                }
                Err(e) => {
                    // Soft-fail: harvest will surface AccountNotFound downstream
                    // if the missing mint is actually needed.
                    tracing::warn!(
                        %mint,
                        error = %e,
                        "Failed to fetch reward mint; harvest may fail"
                    );
                }
            }
        }
    }

    if !reward_vaults.is_empty() {
        tracing::info!(count = reward_vaults.len(), "Fetching reward vault accounts");
        let vault_futures: Vec<_> = reward_vaults.iter().map(|v| client.get_account(v)).collect();
        let vault_results = futures::future::join_all(vault_futures).await;
        for (vault, result) in reward_vaults.iter().zip(vault_results) {
            match result {
                Ok(account) => {
                    set_account(&ctx, *vault, account)?;
                    tracing::info!(%vault, "Reward vault added");
                }
                Err(e) => {
                    tracing::warn!(
                        %vault,
                        error = %e,
                        "Failed to fetch reward vault; close_position may fail"
                    );
                }
            }
        }
    }

    // Extract pool state
    let current_tick = whirlpool_parsed.tick_current_index;
    let tick_spacing = whirlpool_parsed.tick_spacing;
    let token_vault_a = whirlpool_parsed.token_vault_a;
    let token_vault_b = whirlpool_parsed.token_vault_b;

    // Fetch token vaults and extra accounts
    tracing::info!("Fetching pool token vaults and extra accounts");
    const METADATA_UPDATE_AUTHORITY: &str = "3axbTs2z5GBy6usVbNVoqEgZMng3vZvMnAoX29BFfwhr";
    let metadata_update_authority = Pubkey::from_str(METADATA_UPDATE_AUTHORITY)
        .map_err(|e| anyhow::anyhow!("Invalid metadata update authority address: {}", e))?;
    let spl_token_2022_program_id =
        Pubkey::from_str("TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb")?;

    let (vault_a_result, vault_b_result, metadata_result, spl_token_2022_result) = tokio::join!(
        client.get_account(&token_vault_a),
        client.get_account(&token_vault_b),
        client.get_account(&metadata_update_authority),
        client.get_account(&spl_token_2022_program_id),
    );

    match vault_a_result {
        Ok(vault_account) => {
            set_account(&ctx, token_vault_a, vault_account)?;
            tracing::info!(
                vault = %token_vault_a,
                "Token vault A added"
            );
        }
        Err(e) => {
            tracing::warn!(
                error = %e,
                "Failed to fetch token vault A"
            );
        }
    }

    match vault_b_result {
        Ok(vault_account) => {
            set_account(&ctx, token_vault_b, vault_account)?;
            tracing::info!(
                vault = %token_vault_b,
                "Token vault B added"
            );
        }
        Err(e) => {
            tracing::warn!(
                error = %e,
                "Failed to fetch token vault B"
            );
        }
    }

    match metadata_result {
        Ok(metadata_account) => {
            set_account(&ctx, metadata_update_authority, metadata_account)?;
            tracing::info!(
                address = %metadata_update_authority,
                "Metadata update authority added"
            );
        }
        Err(e) => {
            tracing::warn!(
                error = %e,
                "Failed to fetch metadata update authority"
            );
            const SYSTEM_PROGRAM_ID: &str = "11111111111111111111111111111111";
            let system_program_id = Pubkey::from_str(SYSTEM_PROGRAM_ID)?;
            let min_rent = {
                let svm = ctx.lock_svm();
                svm.minimum_balance_for_rent_exemption(0)
            };
            set_account(
                &ctx,
                metadata_update_authority,
                Account {
                    lamports: min_rent,
                    data: vec![],
                    owner: system_program_id,
                    executable: false,
                    rent_epoch: 0,
                },
            )?;
            tracing::info!(
                address = %metadata_update_authority,
                "Created empty metadata update authority account"
            );
        }
    }

    // Get tick arrays
    tracing::info!("Fetching tick array data");

    let _spl_token_2022_account_info = spl_token_2022_result?;

    let offset = tick_spacing as i32 * TICK_ARRAY_SIZE;
    let tick_array_start_index = get_tick_array_start_tick_index(current_tick, tick_spacing);

    let tick_array_indices = [
        tick_array_start_index - offset * 2,
        tick_array_start_index - offset,
        tick_array_start_index,
        tick_array_start_index + offset,
        tick_array_start_index + offset * 2,
    ];

    let tick_array_addresses: Vec<(i32, Pubkey)> = tick_array_indices
        .iter()
        .map(|&idx| {
            let addr = get_tick_array_address(&pool_address, idx);
            (idx, addr)
        })
        .collect();

    let tick_array_futures: Vec<_> =
        tick_array_addresses.iter().map(|(_, addr)| client.get_account(addr)).collect();

    tracing::info!("Fetching all tick arrays in parallel");
    let tick_array_results = futures::future::join_all(tick_array_futures).await;

    for ((tick_array_index, tick_array_address), result) in
        tick_array_addresses.iter().zip(tick_array_results)
    {
        tracing::debug!(
            address = %tick_array_address,
            "Tick array address"
        );

        match result {
            Ok(tick_array_account) => {
                tracing::debug!(tick_array_index, "Fetched tick array from mainnet");

                set_account(&ctx, *tick_array_address, tick_array_account)?;
                tracing::info!(
                    address = %tick_array_address,
                    tick_array_index,
                    "Tick array added"
                );
            }
            Err(e) => {
                const TICK_ARRAY_ACCOUNT_SIZE: usize = 8 + 16 + (88 * 20);

                tracing::warn!(
                    tick_array_index,
                    error = %e,
                    "Tick array not found on mainnet"
                );
                tracing::debug!(tick_array_index, "Creating empty tick array account");

                let min_rent = {
                    let svm = ctx.lock_svm();
                    svm.minimum_balance_for_rent_exemption(TICK_ARRAY_ACCOUNT_SIZE)
                };
                let empty_data = vec![0u8; TICK_ARRAY_ACCOUNT_SIZE];

                set_account(
                    &ctx,
                    *tick_array_address,
                    Account {
                        lamports: min_rent,
                        data: empty_data,
                        owner: program_id,
                        executable: false,
                        rent_epoch: 0,
                    },
                )?;

                tracing::info!(
                    address = %tick_array_address,
                    tick_array_index,
                    size = TICK_ARRAY_ACCOUNT_SIZE,
                    rent = min_rent,
                    "Empty tick array account created"
                );
            }
        }
    }

    tracing::info!("Tick array fetching complete");

    Ok((
        ctx,
        WhirlpoolAccounts {
            program_id,
            pool_address,
            usdc_mint,
            wsol_mint,
            payer_usdc_ata,
            payer_wsol_ata,
            tick_spacing,
        },
    ))
}