diff --git a/integration_tests/tests/fixtures/mod.rs b/integration_tests/tests/fixtures/mod.rs index c15b19b..cfc6717 100644 --- a/integration_tests/tests/fixtures/mod.rs +++ b/integration_tests/tests/fixtures/mod.rs @@ -6,6 +6,7 @@ use thiserror::Error; pub mod restaking_client; pub mod test_builder; pub mod tip_router_client; +pub mod vault_client; pub type TestResult = Result; diff --git a/integration_tests/tests/fixtures/restaking_client.rs b/integration_tests/tests/fixtures/restaking_client.rs index 9128bf8..ab0d6b1 100644 --- a/integration_tests/tests/fixtures/restaking_client.rs +++ b/integration_tests/tests/fixtures/restaking_client.rs @@ -1,12 +1,34 @@ -use jito_restaking_core::{config::Config, ncn::Ncn}; -use jito_restaking_sdk::sdk::{initialize_config, initialize_ncn}; +use jito_bytemuck::AccountDeserialize; +use jito_restaking_core::{ + config::Config, ncn::Ncn, ncn_operator_state::NcnOperatorState, + ncn_vault_slasher_ticket::NcnVaultSlasherTicket, ncn_vault_ticket::NcnVaultTicket, + operator::Operator, operator_vault_ticket::OperatorVaultTicket, +}; +use jito_restaking_sdk::{ + error::RestakingError, + instruction::OperatorAdminRole, + sdk::{ + cooldown_ncn_vault_ticket, initialize_config, initialize_ncn, + initialize_ncn_operator_state, initialize_ncn_vault_slasher_ticket, + initialize_ncn_vault_ticket, initialize_operator, initialize_operator_vault_ticket, + ncn_cooldown_operator, ncn_set_admin, ncn_warmup_operator, operator_cooldown_ncn, + operator_set_admin, operator_set_fee, operator_set_secondary_admin, operator_warmup_ncn, + set_config_admin, warmup_ncn_vault_slasher_ticket, warmup_ncn_vault_ticket, + warmup_operator_vault_ticket, + }, +}; +use solana_program::{ + instruction::InstructionError, native_token::sol_to_lamports, pubkey::Pubkey, + system_instruction::transfer, +}; use solana_program_test::BanksClient; use solana_sdk::{ - commitment_config::CommitmentLevel, native_token::sol_to_lamports, pubkey::Pubkey, - signature::Keypair, signer::Signer, system_instruction::transfer, transaction::Transaction, + commitment_config::CommitmentLevel, + signature::{Keypair, Signer}, + transaction::{Transaction, TransactionError}, }; -use super::TestResult; +use crate::fixtures::{TestError, TestResult}; #[derive(Debug)] pub struct NcnRoot { @@ -14,6 +36,12 @@ pub struct NcnRoot { pub ncn_admin: Keypair, } +#[derive(Debug)] +pub struct OperatorRoot { + pub operator_pubkey: Pubkey, + pub operator_admin: Keypair, +} + pub struct RestakingProgramClient { banks_client: BanksClient, payer: Keypair, @@ -27,43 +55,189 @@ impl RestakingProgramClient { } } - pub async fn process_transaction(&mut self, tx: &Transaction) -> TestResult<()> { - self.banks_client - .process_transaction_with_preflight_and_commitment( - tx.clone(), - CommitmentLevel::Processed, - ) - .await?; - Ok(()) + pub async fn get_ncn(&mut self, ncn: &Pubkey) -> TestResult { + let account = self + .banks_client + .get_account_with_commitment(*ncn, CommitmentLevel::Processed) + .await? + .unwrap(); + + Ok(*Ncn::try_from_slice_unchecked(account.data.as_slice())?) } - pub async fn airdrop(&mut self, to: &Pubkey, sol: f64) -> TestResult<()> { - let blockhash = self.banks_client.get_latest_blockhash().await?; - self.banks_client - .process_transaction_with_preflight_and_commitment( - Transaction::new_signed_with_payer( - &[transfer(&self.payer.pubkey(), to, sol_to_lamports(sol))], - Some(&self.payer.pubkey()), - &[&self.payer], - blockhash, - ), - CommitmentLevel::Processed, - ) - .await?; - Ok(()) + pub async fn get_config(&mut self, account: &Pubkey) -> TestResult { + let account = self.banks_client.get_account(*account).await?.unwrap(); + Ok(*Config::try_from_slice_unchecked(account.data.as_slice())?) + } + + pub async fn get_ncn_vault_ticket( + &mut self, + ncn: &Pubkey, + vault: &Pubkey, + ) -> TestResult { + let account = + NcnVaultTicket::find_program_address(&jito_restaking_program::id(), ncn, vault).0; + let account = self.banks_client.get_account(account).await?.unwrap(); + Ok(*NcnVaultTicket::try_from_slice_unchecked( + account.data.as_slice(), + )?) + } + + pub async fn get_ncn_operator_state( + &mut self, + ncn: &Pubkey, + operator: &Pubkey, + ) -> TestResult { + let account = + NcnOperatorState::find_program_address(&jito_restaking_program::id(), ncn, operator).0; + let account = self.banks_client.get_account(account).await?.unwrap(); + Ok(*NcnOperatorState::try_from_slice_unchecked( + account.data.as_slice(), + )?) + } + + pub async fn get_operator(&mut self, account: &Pubkey) -> TestResult { + let account = self.banks_client.get_account(*account).await?.unwrap(); + Ok(*Operator::try_from_slice_unchecked( + account.data.as_slice(), + )?) + } + + pub async fn get_operator_vault_ticket( + &mut self, + operator: &Pubkey, + vault: &Pubkey, + ) -> TestResult { + let account = OperatorVaultTicket::find_program_address( + &jito_restaking_program::id(), + operator, + vault, + ) + .0; + let account = self.banks_client.get_account(account).await?.unwrap(); + Ok(*OperatorVaultTicket::try_from_slice_unchecked( + account.data.as_slice(), + )?) + } + + #[allow(dead_code)] + pub async fn get_operator_ncn_ticket( + &mut self, + operator: &Pubkey, + ncn: &Pubkey, + ) -> TestResult { + let account = + NcnOperatorState::find_program_address(&jito_restaking_program::id(), operator, ncn).0; + let account = self.banks_client.get_account(account).await?.unwrap(); + Ok(*NcnOperatorState::try_from_slice_unchecked( + account.data.as_slice(), + )?) } pub async fn do_initialize_config(&mut self) -> TestResult { let restaking_config_pubkey = Config::find_program_address(&jito_restaking_program::id()).0; let restaking_config_admin = Keypair::new(); - self.airdrop(&restaking_config_admin.pubkey(), 10.0).await?; + self.airdrop(&restaking_config_admin.pubkey(), 1.0).await?; self.initialize_config(&restaking_config_pubkey, &restaking_config_admin) .await?; Ok(restaking_config_admin) } + pub async fn do_initialize_operator(&mut self) -> TestResult { + // create operator + add operator vault + let operator_base = Keypair::new(); + let operator_pubkey = + Operator::find_program_address(&jito_restaking_program::id(), &operator_base.pubkey()) + .0; + let operator_admin = Keypair::new(); + self.airdrop(&operator_admin.pubkey(), 1.0).await?; + self.initialize_operator( + &Config::find_program_address(&jito_restaking_program::id()).0, + &operator_pubkey, + &operator_admin, + &operator_base, + 0, + ) + .await?; + Ok(OperatorRoot { + operator_pubkey, + operator_admin, + }) + } + + pub async fn do_initialize_operator_vault_ticket( + &mut self, + operator_root: &OperatorRoot, + vault_pubkey: &Pubkey, + ) -> TestResult<()> { + let operator_vault_ticket = OperatorVaultTicket::find_program_address( + &jito_restaking_program::id(), + &operator_root.operator_pubkey, + vault_pubkey, + ) + .0; + self.initialize_operator_vault_ticket( + &Config::find_program_address(&jito_restaking_program::id()).0, + &operator_root.operator_pubkey, + vault_pubkey, + &operator_vault_ticket, + &operator_root.operator_admin, + &operator_root.operator_admin, + ) + .await?; + + Ok(()) + } + + pub async fn do_warmup_operator_vault_ticket( + &mut self, + operator_root: &OperatorRoot, + vault_pubkey: &Pubkey, + ) -> TestResult<()> { + let operator_vault_ticket = OperatorVaultTicket::find_program_address( + &jito_restaking_program::id(), + &operator_root.operator_pubkey, + vault_pubkey, + ) + .0; + self.warmup_operator_vault_ticket( + &Config::find_program_address(&jito_restaking_program::id()).0, + &operator_root.operator_pubkey, + vault_pubkey, + &operator_vault_ticket, + &operator_root.operator_admin, + ) + .await + } + + pub async fn warmup_operator_vault_ticket( + &mut self, + config: &Pubkey, + operator: &Pubkey, + vault: &Pubkey, + operator_vault_ticket: &Pubkey, + admin: &Keypair, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self.process_transaction(&Transaction::new_signed_with_payer( + &[warmup_operator_vault_ticket( + &jito_restaking_program::id(), + config, + operator, + vault, + operator_vault_ticket, + &admin.pubkey(), + )], + Some(&self.payer.pubkey()), + &[admin, &self.payer], + blockhash, + )) + .await + } + pub async fn initialize_config( &mut self, config: &Pubkey, @@ -106,6 +280,429 @@ impl RestakingProgramClient { }) } + pub async fn do_initialize_ncn_vault_ticket( + &mut self, + ncn_root: &NcnRoot, + vault: &Pubkey, + ) -> TestResult<()> { + let ncn_vault_ticket = NcnVaultTicket::find_program_address( + &jito_restaking_program::id(), + &ncn_root.ncn_pubkey, + vault, + ) + .0; + + self.initialize_ncn_vault_ticket( + &Config::find_program_address(&jito_restaking_program::id()).0, + &ncn_root.ncn_pubkey, + vault, + &ncn_vault_ticket, + &ncn_root.ncn_admin, + &self.payer.insecure_clone(), + ) + .await + } + + pub async fn do_warmup_ncn_vault_ticket( + &mut self, + ncn_root: &NcnRoot, + vault: &Pubkey, + ) -> TestResult<()> { + let ncn_vault_ticket = NcnVaultTicket::find_program_address( + &jito_restaking_program::id(), + &ncn_root.ncn_pubkey, + vault, + ) + .0; + self.warmup_ncn_vault_ticket( + &Config::find_program_address(&jito_restaking_program::id()).0, + &ncn_root.ncn_pubkey, + vault, + &ncn_vault_ticket, + &ncn_root.ncn_admin, + ) + .await + } + + pub async fn warmup_ncn_vault_ticket( + &mut self, + config: &Pubkey, + ncn: &Pubkey, + vault: &Pubkey, + ncn_vault_ticket: &Pubkey, + admin: &Keypair, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self.process_transaction(&Transaction::new_signed_with_payer( + &[warmup_ncn_vault_ticket( + &jito_restaking_program::id(), + config, + ncn, + vault, + ncn_vault_ticket, + &admin.pubkey(), + )], + Some(&self.payer.pubkey()), + &[admin, &self.payer], + blockhash, + )) + .await + } + + #[allow(dead_code)] + pub async fn do_cooldown_ncn_vault_ticket( + &mut self, + ncn_root: &NcnRoot, + vault: &Pubkey, + ) -> TestResult<()> { + let ncn_vault_ticket = NcnVaultTicket::find_program_address( + &jito_restaking_program::id(), + &ncn_root.ncn_pubkey, + vault, + ) + .0; + self.cooldown_ncn_vault_ticket( + &Config::find_program_address(&jito_restaking_program::id()).0, + &ncn_root.ncn_pubkey, + vault, + &ncn_vault_ticket, + &ncn_root.ncn_admin, + ) + .await + } + + #[allow(dead_code)] + pub async fn cooldown_ncn_vault_ticket( + &mut self, + config: &Pubkey, + ncn: &Pubkey, + vault: &Pubkey, + ncn_vault_ticket: &Pubkey, + admin: &Keypair, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self.process_transaction(&Transaction::new_signed_with_payer( + &[cooldown_ncn_vault_ticket( + &jito_restaking_program::id(), + config, + ncn, + vault, + ncn_vault_ticket, + &admin.pubkey(), + )], + Some(&self.payer.pubkey()), + &[admin, &self.payer], + blockhash, + )) + .await + } + + pub async fn do_ncn_warmup_operator( + &mut self, + ncn_root: &NcnRoot, + operator_pubkey: &Pubkey, + ) -> TestResult<()> { + self.ncn_warmup_operator( + &Config::find_program_address(&jito_restaking_program::id()).0, + &ncn_root.ncn_pubkey, + operator_pubkey, + &NcnOperatorState::find_program_address( + &jito_restaking_program::id(), + &ncn_root.ncn_pubkey, + operator_pubkey, + ) + .0, + &ncn_root.ncn_admin, + ) + .await + } + + pub async fn do_ncn_cooldown_operator( + &mut self, + ncn_root: &NcnRoot, + operator_pubkey: &Pubkey, + ) -> TestResult<()> { + self.ncn_cooldown_operator( + &Config::find_program_address(&jito_restaking_program::id()).0, + &ncn_root.ncn_pubkey, + operator_pubkey, + &NcnOperatorState::find_program_address( + &jito_restaking_program::id(), + &ncn_root.ncn_pubkey, + operator_pubkey, + ) + .0, + &ncn_root.ncn_admin, + ) + .await + } + + pub async fn ncn_cooldown_operator( + &mut self, + config: &Pubkey, + ncn: &Pubkey, + operator_pubkey: &Pubkey, + ncn_operator_state: &Pubkey, + admin: &Keypair, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self.process_transaction(&Transaction::new_signed_with_payer( + &[ncn_cooldown_operator( + &jito_restaking_program::id(), + config, + ncn, + operator_pubkey, + ncn_operator_state, + &admin.pubkey(), + )], + Some(&self.payer.pubkey()), + &[admin, &self.payer], + blockhash, + )) + .await + } + + pub async fn ncn_warmup_operator( + &mut self, + config: &Pubkey, + ncn: &Pubkey, + operator_pubkey: &Pubkey, + ncn_operator_state: &Pubkey, + admin: &Keypair, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self.process_transaction(&Transaction::new_signed_with_payer( + &[ncn_warmup_operator( + &jito_restaking_program::id(), + config, + ncn, + operator_pubkey, + ncn_operator_state, + &admin.pubkey(), + )], + Some(&self.payer.pubkey()), + &[admin, &self.payer], + blockhash, + )) + .await + } + + pub async fn do_operator_warmup_ncn( + &mut self, + operator_root: &OperatorRoot, + ncn_pubkey: &Pubkey, + ) -> TestResult<()> { + self.operator_warmup_ncn( + &Config::find_program_address(&jito_restaking_program::id()).0, + ncn_pubkey, + &operator_root.operator_pubkey, + &NcnOperatorState::find_program_address( + &jito_restaking_program::id(), + ncn_pubkey, + &operator_root.operator_pubkey, + ) + .0, + &operator_root.operator_admin, + ) + .await + } + + pub async fn operator_warmup_ncn( + &mut self, + config: &Pubkey, + ncn_pubkey: &Pubkey, + operator_pubkey: &Pubkey, + ncn_operator_state: &Pubkey, + admin: &Keypair, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self.process_transaction(&Transaction::new_signed_with_payer( + &[operator_warmup_ncn( + &jito_restaking_program::id(), + config, + ncn_pubkey, + operator_pubkey, + ncn_operator_state, + &admin.pubkey(), + )], + Some(&self.payer.pubkey()), + &[admin, &self.payer], + blockhash, + )) + .await + } + + pub async fn do_operator_cooldown_ncn( + &mut self, + operator_root: &OperatorRoot, + ncn_pubkey: &Pubkey, + ) -> TestResult<()> { + self.operator_cooldown_ncn( + &Config::find_program_address(&jito_restaking_program::id()).0, + ncn_pubkey, + &operator_root.operator_pubkey, + &NcnOperatorState::find_program_address( + &jito_restaking_program::id(), + ncn_pubkey, + &operator_root.operator_pubkey, + ) + .0, + &operator_root.operator_admin, + ) + .await + } + + pub async fn operator_cooldown_ncn( + &mut self, + config: &Pubkey, + ncn_pubkey: &Pubkey, + operator_pubkey: &Pubkey, + ncn_operator_state: &Pubkey, + admin: &Keypair, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self.process_transaction(&Transaction::new_signed_with_payer( + &[operator_cooldown_ncn( + &jito_restaking_program::id(), + config, + ncn_pubkey, + operator_pubkey, + ncn_operator_state, + &admin.pubkey(), + )], + Some(&self.payer.pubkey()), + &[admin, &self.payer], + blockhash, + )) + .await + } + + pub async fn do_initialize_ncn_operator_state( + &mut self, + ncn_root: &NcnRoot, + operator: &Pubkey, + ) -> TestResult<()> { + let ncn_operator_state = NcnOperatorState::find_program_address( + &jito_restaking_program::id(), + &ncn_root.ncn_pubkey, + operator, + ) + .0; + + self.initialize_ncn_operator_state( + &Config::find_program_address(&jito_restaking_program::id()).0, + &ncn_root.ncn_pubkey, + operator, + &ncn_operator_state, + &ncn_root.ncn_admin, + &self.payer.insecure_clone(), + ) + .await + } + + pub async fn do_initialize_ncn_vault_slasher_ticket( + &mut self, + ncn_root: &NcnRoot, + vault: &Pubkey, + slasher: &Pubkey, + max_slash_amount: u64, + ) -> TestResult<()> { + let ncn_vault_ticket = NcnVaultTicket::find_program_address( + &jito_restaking_program::id(), + &ncn_root.ncn_pubkey, + vault, + ) + .0; + let ncn_slasher_ticket = NcnVaultSlasherTicket::find_program_address( + &jito_restaking_program::id(), + &ncn_root.ncn_pubkey, + vault, + slasher, + ) + .0; + + self.initialize_ncn_vault_slasher_ticket( + &Config::find_program_address(&jito_restaking_program::id()).0, + &ncn_root.ncn_pubkey, + vault, + slasher, + &ncn_vault_ticket, + &ncn_slasher_ticket, + &ncn_root.ncn_admin, + &self.payer.insecure_clone(), + max_slash_amount, + ) + .await + } + + pub async fn do_warmup_ncn_vault_slasher_ticket( + &mut self, + ncn_root: &NcnRoot, + vault: &Pubkey, + slasher: &Pubkey, + ) -> TestResult<()> { + let ncn_vault_ticket = NcnVaultTicket::find_program_address( + &jito_restaking_program::id(), + &ncn_root.ncn_pubkey, + vault, + ) + .0; + let ncn_slasher_ticket = NcnVaultSlasherTicket::find_program_address( + &jito_restaking_program::id(), + &ncn_root.ncn_pubkey, + vault, + slasher, + ) + .0; + + self.warmup_ncn_vault_slasher_ticket( + &Config::find_program_address(&jito_restaking_program::id()).0, + &ncn_root.ncn_pubkey, + vault, + slasher, + &ncn_vault_ticket, + &ncn_slasher_ticket, + &ncn_root.ncn_admin, + ) + .await + } + + pub async fn warmup_ncn_vault_slasher_ticket( + &mut self, + config: &Pubkey, + ncn: &Pubkey, + vault: &Pubkey, + slasher: &Pubkey, + ncn_vault_ticket: &Pubkey, + ncn_slasher_ticket: &Pubkey, + admin: &Keypair, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self.process_transaction(&Transaction::new_signed_with_payer( + &[warmup_ncn_vault_slasher_ticket( + &jito_restaking_program::id(), + config, + ncn, + vault, + slasher, + ncn_vault_ticket, + ncn_slasher_ticket, + &admin.pubkey(), + )], + Some(&self.payer.pubkey()), + &[admin, &self.payer], + blockhash, + )) + .await + } + pub async fn initialize_ncn( &mut self, config: &Pubkey, @@ -129,4 +726,354 @@ impl RestakingProgramClient { )) .await } + + pub async fn initialize_ncn_vault_ticket( + &mut self, + config: &Pubkey, + ncn: &Pubkey, + vault: &Pubkey, + ncn_vault_ticket: &Pubkey, + ncn_admin: &Keypair, + payer: &Keypair, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self.process_transaction(&Transaction::new_signed_with_payer( + &[initialize_ncn_vault_ticket( + &jito_restaking_program::id(), + config, + ncn, + vault, + ncn_vault_ticket, + &ncn_admin.pubkey(), + &payer.pubkey(), + )], + Some(&payer.pubkey()), + &[ncn_admin, payer], + blockhash, + )) + .await + } + + pub async fn initialize_ncn_operator_state( + &mut self, + config: &Pubkey, + ncn: &Pubkey, + operator: &Pubkey, + ncn_operator_state: &Pubkey, + ncn_admin: &Keypair, + payer: &Keypair, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self.process_transaction(&Transaction::new_signed_with_payer( + &[initialize_ncn_operator_state( + &jito_restaking_program::id(), + config, + ncn, + operator, + ncn_operator_state, + &ncn_admin.pubkey(), + &payer.pubkey(), + )], + Some(&payer.pubkey()), + &[ncn_admin, payer], + blockhash, + )) + .await + } + + pub async fn initialize_ncn_vault_slasher_ticket( + &mut self, + config: &Pubkey, + ncn: &Pubkey, + vault: &Pubkey, + slasher: &Pubkey, + ncn_vault_ticket: &Pubkey, + ncn_slasher_ticket: &Pubkey, + ncn_admin: &Keypair, + payer: &Keypair, + max_slash_amount: u64, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self.process_transaction(&Transaction::new_signed_with_payer( + &[initialize_ncn_vault_slasher_ticket( + &jito_restaking_program::id(), + config, + ncn, + vault, + slasher, + ncn_vault_ticket, + ncn_slasher_ticket, + &ncn_admin.pubkey(), + &payer.pubkey(), + max_slash_amount, + )], + Some(&payer.pubkey()), + &[ncn_admin, payer], + blockhash, + )) + .await + } + + pub async fn ncn_set_admin( + &mut self, + ncn: &Pubkey, + old_admin: &Keypair, + new_admin: &Keypair, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self.process_transaction(&Transaction::new_signed_with_payer( + &[ncn_set_admin( + &jito_restaking_program::id(), + ncn, + &old_admin.pubkey(), + &new_admin.pubkey(), + )], + Some(&old_admin.pubkey()), + &[old_admin, new_admin], + blockhash, + )) + .await + } + + pub async fn operator_set_admin( + &mut self, + operator: &Pubkey, + old_admin: &Keypair, + new_admin: &Keypair, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self.process_transaction(&Transaction::new_signed_with_payer( + &[operator_set_admin( + &jito_restaking_program::id(), + operator, + &old_admin.pubkey(), + &new_admin.pubkey(), + )], + Some(&old_admin.pubkey()), + &[old_admin, new_admin], + blockhash, + )) + .await + } + + pub async fn operator_set_secondary_admin( + &mut self, + operator: &Pubkey, + old_admin: &Keypair, + new_admin: &Keypair, + operator_admin_role: OperatorAdminRole, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self.process_transaction(&Transaction::new_signed_with_payer( + &[operator_set_secondary_admin( + &jito_restaking_program::id(), + operator, + &old_admin.pubkey(), + &new_admin.pubkey(), + operator_admin_role, + )], + Some(&old_admin.pubkey()), + &[old_admin], + blockhash, + )) + .await + } + + pub async fn initialize_operator( + &mut self, + config: &Pubkey, + operator: &Pubkey, + admin: &Keypair, + base: &Keypair, + operator_fee_bps: u16, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self.process_transaction(&Transaction::new_signed_with_payer( + &[initialize_operator( + &jito_restaking_program::id(), + config, + operator, + &admin.pubkey(), + &base.pubkey(), + operator_fee_bps, + )], + Some(&admin.pubkey()), + &[admin, base], + blockhash, + )) + .await + } + + pub async fn initialize_operator_vault_ticket( + &mut self, + config: &Pubkey, + operator: &Pubkey, + vault: &Pubkey, + operator_vault_ticket: &Pubkey, + admin: &Keypair, + payer: &Keypair, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self.process_transaction(&Transaction::new_signed_with_payer( + &[initialize_operator_vault_ticket( + &jito_restaking_program::id(), + config, + operator, + vault, + operator_vault_ticket, + &admin.pubkey(), + &payer.pubkey(), + )], + Some(&payer.pubkey()), + &[admin, payer], + blockhash, + )) + .await + } + + pub async fn operator_set_fee( + &mut self, + config: &Pubkey, + operator: &Pubkey, + admin: &Keypair, + new_fee_bps: u16, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self.process_transaction(&Transaction::new_signed_with_payer( + &[operator_set_fee( + &jito_restaking_program::id(), + config, + operator, + &admin.pubkey(), + new_fee_bps, + )], + Some(&self.payer.pubkey()), + &[admin, &self.payer], + blockhash, + )) + .await + } + + pub async fn ncn_delegate_token_account( + &mut self, + ncn_pubkey: &Pubkey, + delegate_admin: &Keypair, + token_mint: &Pubkey, + token_account: &Pubkey, + delegate: &Pubkey, + token_program_id: &Pubkey, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self.process_transaction(&Transaction::new_signed_with_payer( + &[jito_restaking_sdk::sdk::ncn_delegate_token_account( + &jito_restaking_program::id(), + ncn_pubkey, + &delegate_admin.pubkey(), + token_mint, + token_account, + delegate, + token_program_id, + )], + Some(&self.payer.pubkey()), + &[&self.payer, delegate_admin], + blockhash, + )) + .await + } + + pub async fn operator_delegate_token_account( + &mut self, + operator_pubkey: &Pubkey, + delegate_admin: &Keypair, + token_mint: &Pubkey, + token_account: &Pubkey, + delegate: &Pubkey, + token_program_id: &Pubkey, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self.process_transaction(&Transaction::new_signed_with_payer( + &[jito_restaking_sdk::sdk::operator_delegate_token_account( + &jito_restaking_program::id(), + operator_pubkey, + &delegate_admin.pubkey(), + token_mint, + token_account, + delegate, + token_program_id, + )], + Some(&self.payer.pubkey()), + &[&self.payer, delegate_admin], + blockhash, + )) + .await + } + + pub async fn process_transaction(&mut self, tx: &Transaction) -> TestResult<()> { + self.banks_client + .process_transaction_with_preflight_and_commitment( + tx.clone(), + CommitmentLevel::Processed, + ) + .await?; + Ok(()) + } + + pub async fn airdrop(&mut self, to: &Pubkey, sol: f64) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self.banks_client + .process_transaction_with_preflight_and_commitment( + Transaction::new_signed_with_payer( + &[transfer(&self.payer.pubkey(), to, sol_to_lamports(sol))], + Some(&self.payer.pubkey()), + &[&self.payer], + blockhash, + ), + CommitmentLevel::Processed, + ) + .await?; + Ok(()) + } + + pub async fn set_config_admin( + &mut self, + config: &Pubkey, + old_admin: &Keypair, + new_admin: &Keypair, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self.process_transaction(&Transaction::new_signed_with_payer( + &[set_config_admin( + &jito_restaking_program::id(), + config, + &old_admin.pubkey(), + &new_admin.pubkey(), + )], + Some(&old_admin.pubkey()), + &[old_admin], + blockhash, + )) + .await + } +} + +#[track_caller] +#[inline(always)] +pub fn assert_restaking_error( + test_error: Result, + restaking_error: RestakingError, +) { + assert!(test_error.is_err()); + assert_eq!( + test_error.err().unwrap().to_transaction_error().unwrap(), + TransactionError::InstructionError(0, InstructionError::Custom(restaking_error as u32)) + ); } diff --git a/integration_tests/tests/fixtures/test_builder.rs b/integration_tests/tests/fixtures/test_builder.rs index 74545ef..4df80d4 100644 --- a/integration_tests/tests/fixtures/test_builder.rs +++ b/integration_tests/tests/fixtures/test_builder.rs @@ -6,6 +6,7 @@ use solana_program_test::{processor, BanksClientError, ProgramTest, ProgramTestC use super::{ restaking_client::{NcnRoot, RestakingProgramClient}, tip_router_client::TipRouterClient, + vault_client::VaultProgramClient, TestResult, }; @@ -72,6 +73,13 @@ impl TestBuilder { ) } + pub fn vault_client(&self) -> VaultProgramClient { + VaultProgramClient::new( + self.context.banks_client.clone(), + self.context.payer.insecure_clone(), + ) + } + pub async fn setup_ncn(&mut self) -> TestResult { let mut restaking_program_client = self.restaking_program_client(); diff --git a/integration_tests/tests/fixtures/vault_client.rs b/integration_tests/tests/fixtures/vault_client.rs index 06a9639..7f86a88 100644 --- a/integration_tests/tests/fixtures/vault_client.rs +++ b/integration_tests/tests/fixtures/vault_client.rs @@ -1,8 +1,48 @@ use std::{fmt, fmt::Debug}; -use solana_program::pubkey::Pubkey; -use solana_program_test::BanksClient; -use solana_sdk::signature::Keypair; +use borsh::BorshDeserialize; +use jito_bytemuck::AccountDeserialize; +use jito_restaking_core::{ + ncn_vault_slasher_ticket::NcnVaultSlasherTicket, ncn_vault_ticket::NcnVaultTicket, + operator_vault_ticket::OperatorVaultTicket, +}; +use jito_vault_core::{ + config::Config, vault::Vault, vault_ncn_slasher_operator_ticket::VaultNcnSlasherOperatorTicket, + vault_ncn_slasher_ticket::VaultNcnSlasherTicket, vault_ncn_ticket::VaultNcnTicket, + vault_operator_delegation::VaultOperatorDelegation, + vault_staker_withdrawal_ticket::VaultStakerWithdrawalTicket, + vault_update_state_tracker::VaultUpdateStateTracker, +}; +use jito_vault_sdk::{ + error::VaultError, + instruction::{VaultAdminRole, WithdrawalAllocationMethod}, + sdk::{ + add_delegation, cooldown_delegation, initialize_config, initialize_vault, + set_deposit_capacity, warmup_vault_ncn_slasher_ticket, warmup_vault_ncn_ticket, + }, +}; +use log::info; +use solana_program::{ + clock::Clock, + native_token::sol_to_lamports, + program_pack::Pack, + pubkey::Pubkey, + rent::Rent, + system_instruction::{create_account, transfer}, +}; +use solana_program_test::{BanksClient, BanksClientError}; +use solana_sdk::{ + commitment_config::CommitmentLevel, + instruction::InstructionError, + signature::{Keypair, Signer}, + transaction::{Transaction, TransactionError}, +}; +use spl_associated_token_account::{ + get_associated_token_address, instruction::create_associated_token_account_idempotent, +}; +use spl_token::state::Account as SPLTokenAccount; + +use crate::fixtures::{TestError, TestResult}; pub struct VaultRoot { pub vault_pubkey: Pubkey, @@ -19,7 +59,11 @@ impl Debug for VaultRoot { } } -#[allow(dead_code)] +#[derive(Debug)] +pub struct VaultStakerWithdrawalTicketRoot { + pub base: Pubkey, +} + pub struct VaultProgramClient { banks_client: BanksClient, payer: Keypair, @@ -32,4 +76,1602 @@ impl VaultProgramClient { payer, } } + + pub async fn configure_depositor( + &mut self, + vault_root: &VaultRoot, + depositor: &Pubkey, + amount_to_mint: u64, + ) -> TestResult<()> { + self.airdrop(depositor, 100.0).await?; + let vault = self.get_vault(&vault_root.vault_pubkey).await?; + self.create_ata(&vault.supported_mint, depositor).await?; + self.create_ata(&vault.vrt_mint, depositor).await?; + self.mint_spl_to(&vault.supported_mint, depositor, amount_to_mint) + .await?; + + Ok(()) + } + + pub async fn get_config(&mut self, account: &Pubkey) -> Result { + let account = self.banks_client.get_account(*account).await?.unwrap(); + Ok(*Config::try_from_slice_unchecked(account.data.as_slice())?) + } + + pub async fn get_vault(&mut self, account: &Pubkey) -> Result { + let account = self.banks_client.get_account(*account).await?.unwrap(); + Ok(*Vault::try_from_slice_unchecked(account.data.as_slice())?) + } + + pub async fn get_vault_ncn_ticket( + &mut self, + vault: &Pubkey, + ncn: &Pubkey, + ) -> Result { + let account = VaultNcnTicket::find_program_address(&jito_vault_program::id(), vault, ncn).0; + let account = self.banks_client.get_account(account).await?.unwrap(); + Ok(*VaultNcnTicket::try_from_slice_unchecked( + account.data.as_slice(), + )?) + } + + pub async fn get_vault_operator_delegation( + &mut self, + vault: &Pubkey, + operator: &Pubkey, + ) -> Result { + let account = VaultOperatorDelegation::find_program_address( + &jito_vault_program::id(), + vault, + operator, + ) + .0; + let account = self.banks_client.get_account(account).await?.unwrap(); + Ok(*VaultOperatorDelegation::try_from_slice_unchecked( + account.data.as_slice(), + )?) + } + + pub async fn get_vault_staker_withdrawal_ticket( + &mut self, + vault: &Pubkey, + staker: &Pubkey, + base: &Pubkey, + ) -> Result { + let account = VaultStakerWithdrawalTicket::find_program_address( + &jito_vault_program::id(), + vault, + base, + ) + .0; + let account = self.banks_client.get_account(account).await?.unwrap(); + let withdrawal_ticket = + *VaultStakerWithdrawalTicket::try_from_slice_unchecked(account.data.as_slice())?; + assert_eq!(withdrawal_ticket.staker, *staker); + Ok(withdrawal_ticket) + } + + pub async fn get_vault_ncn_slasher_ticket( + &mut self, + vault: &Pubkey, + ncn: &Pubkey, + slasher: &Pubkey, + ) -> Result { + let account = VaultNcnSlasherTicket::find_program_address( + &jito_vault_program::id(), + vault, + ncn, + slasher, + ) + .0; + let account = self.banks_client.get_account(account).await?.unwrap(); + Ok(*VaultNcnSlasherTicket::try_from_slice_unchecked( + account.data.as_slice(), + )?) + } + + #[allow(dead_code)] + pub async fn get_vault_ncn_slasher_operator_ticket( + &mut self, + vault: &Pubkey, + ncn: &Pubkey, + slasher: &Pubkey, + operator: &Pubkey, + epoch: u64, + ) -> Result { + let account = VaultNcnSlasherOperatorTicket::find_program_address( + &jito_vault_program::id(), + vault, + ncn, + slasher, + operator, + epoch, + ) + .0; + let account = self.banks_client.get_account(account).await?.unwrap(); + Ok(*VaultNcnSlasherOperatorTicket::try_from_slice_unchecked( + account.data.as_slice(), + )?) + } + + pub async fn get_vault_update_state_tracker( + &mut self, + vault: &Pubkey, + epoch: u64, + ) -> Result { + let account = + VaultUpdateStateTracker::find_program_address(&jito_vault_program::id(), vault, epoch) + .0; + let account = self.banks_client.get_account(account).await?.unwrap(); + Ok(*VaultUpdateStateTracker::try_from_slice_unchecked( + account.data.as_slice(), + )?) + } + + pub async fn do_initialize_config(&mut self) -> Result { + let config_admin = Keypair::new(); + + self.airdrop(&config_admin.pubkey(), 1.0).await?; + + let config_pubkey = Config::find_program_address(&jito_vault_program::id()).0; + self.initialize_config(&config_pubkey, &config_admin, &config_admin.pubkey(), 0) + .await?; + + Ok(config_admin) + } + + pub async fn initialize_config( + &mut self, + config: &Pubkey, + config_admin: &Keypair, + program_fee_wallet: &Pubkey, + program_fee_bps: u16, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self._process_transaction(&Transaction::new_signed_with_payer( + &[initialize_config( + &jito_vault_program::id(), + config, + &config_admin.pubkey(), + &jito_restaking_program::id(), + program_fee_wallet, + program_fee_bps, + )], + Some(&config_admin.pubkey()), + &[config_admin], + blockhash, + )) + .await + } + + pub async fn setup_config_and_vault( + &mut self, + deposit_fee_bps: u16, + withdrawal_fee_bps: u16, + reward_fee_bps: u16, + ) -> Result<(Keypair, VaultRoot), TestError> { + let config_admin = self.do_initialize_config().await?; + let vault_root = self + .do_initialize_vault( + deposit_fee_bps, + withdrawal_fee_bps, + reward_fee_bps, + 9, + &config_admin.pubkey(), + ) + .await?; + + Ok((config_admin, vault_root)) + } + + pub async fn do_initialize_vault( + &mut self, + deposit_fee_bps: u16, + withdrawal_fee_bps: u16, + reward_fee_bps: u16, + decimals: u8, + program_fee_wallet: &Pubkey, + ) -> Result { + let vault_base = Keypair::new(); + + let vault_pubkey = + Vault::find_program_address(&jito_vault_program::id(), &vault_base.pubkey()).0; + + let vrt_mint = Keypair::new(); + let vault_admin = Keypair::new(); + let token_mint = Keypair::new(); + + self.airdrop(&vault_admin.pubkey(), 100.0).await?; + self.create_token_mint(&token_mint, &spl_token::id()) + .await?; + + self.initialize_vault( + &Config::find_program_address(&jito_vault_program::id()).0, + &vault_pubkey, + &vrt_mint, + &token_mint, + &vault_admin, + &vault_base, + deposit_fee_bps, + withdrawal_fee_bps, + reward_fee_bps, + decimals, + ) + .await?; + + // for holding the backed asset in the vault + self.create_ata(&token_mint.pubkey(), &vault_pubkey).await?; + // for holding fees + self.create_ata(&vrt_mint.pubkey(), &vault_admin.pubkey()) + .await?; + // for holding program fee + self.create_ata(&vrt_mint.pubkey(), program_fee_wallet) + .await?; + + // for holding program fee + Ok(VaultRoot { + vault_admin, + vault_pubkey, + }) + } + + pub async fn do_initialize_vault_ncn_ticket( + &mut self, + vault_root: &VaultRoot, + ncn: &Pubkey, + ) -> Result<(), TestError> { + let vault_ncn_ticket = VaultNcnTicket::find_program_address( + &jito_vault_program::id(), + &vault_root.vault_pubkey, + ncn, + ) + .0; + let ncn_vault_ticket = NcnVaultTicket::find_program_address( + &jito_restaking_program::id(), + ncn, + &vault_root.vault_pubkey, + ) + .0; + self.initialize_vault_ncn_ticket( + &Config::find_program_address(&jito_vault_program::id()).0, + &vault_root.vault_pubkey, + ncn, + &ncn_vault_ticket, + &vault_ncn_ticket, + &vault_root.vault_admin, + &self.payer.insecure_clone(), + ) + .await?; + + Ok(()) + } + + pub async fn set_capacity( + &mut self, + config: &Pubkey, + vault: &Pubkey, + admin: &Keypair, + capacity: u64, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self._process_transaction(&Transaction::new_signed_with_payer( + &[set_deposit_capacity( + &jito_vault_program::id(), + config, + vault, + &admin.pubkey(), + capacity, + )], + Some(&admin.pubkey()), + &[&admin], + blockhash, + )) + .await + } + + pub async fn do_warmup_vault_ncn_ticket( + &mut self, + vault_root: &VaultRoot, + ncn: &Pubkey, + ) -> Result<(), TestError> { + let vault_ncn_ticket = VaultNcnTicket::find_program_address( + &jito_vault_program::id(), + &vault_root.vault_pubkey, + ncn, + ) + .0; + + self.warmup_vault_ncn_ticket( + &Config::find_program_address(&jito_vault_program::id()).0, + &vault_root.vault_pubkey, + ncn, + &vault_ncn_ticket, + &vault_root.vault_admin, + ) + .await?; + + Ok(()) + } + + pub async fn warmup_vault_ncn_ticket( + &mut self, + config: &Pubkey, + vault: &Pubkey, + ncn: &Pubkey, + vault_ncn_ticket: &Pubkey, + ncn_vault_admin: &Keypair, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self._process_transaction(&Transaction::new_signed_with_payer( + &[warmup_vault_ncn_ticket( + &jito_vault_program::id(), + config, + vault, + ncn, + vault_ncn_ticket, + &ncn_vault_admin.pubkey(), + )], + Some(&ncn_vault_admin.pubkey()), + &[&ncn_vault_admin], + blockhash, + )) + .await + } + + #[allow(dead_code)] + pub async fn setup_vault_ncn_slasher_operator_ticket( + &mut self, + vault_root: &VaultRoot, + ncn_pubkey: &Pubkey, + slasher: &Pubkey, + operator_pubkey: &Pubkey, + ) -> Result<(), TestError> { + let config = self + .get_config(&Config::find_program_address(&jito_vault_program::id()).0) + .await + .unwrap(); + let clock: Clock = self.banks_client.get_sysvar().await?; + + let vault_ncn_slasher_ticket = VaultNcnSlasherTicket::find_program_address( + &jito_vault_program::id(), + &vault_root.vault_pubkey, + ncn_pubkey, + slasher, + ) + .0; + let vault_ncn_slasher_operator_ticket = + VaultNcnSlasherOperatorTicket::find_program_address( + &jito_vault_program::id(), + &vault_root.vault_pubkey, + ncn_pubkey, + slasher, + operator_pubkey, + clock.slot / config.epoch_length(), + ) + .0; + self.initialize_vault_ncn_slasher_operator_ticket( + &Config::find_program_address(&jito_vault_program::id()).0, + &vault_root.vault_pubkey, + ncn_pubkey, + slasher, + operator_pubkey, + &vault_ncn_slasher_ticket, + &vault_ncn_slasher_operator_ticket, + &self.payer.insecure_clone(), + ) + .await + .unwrap(); + + Ok(()) + } + + pub async fn do_initialize_vault_operator_delegation( + &mut self, + vault_root: &VaultRoot, + operator_pubkey: &Pubkey, + ) -> Result<(), TestError> { + let vault_operator_delegation = VaultOperatorDelegation::find_program_address( + &jito_vault_program::id(), + &vault_root.vault_pubkey, + operator_pubkey, + ) + .0; + let operator_vault_ticket = OperatorVaultTicket::find_program_address( + &jito_restaking_program::id(), + operator_pubkey, + &vault_root.vault_pubkey, + ) + .0; + self.initialize_vault_operator_delegation( + &Config::find_program_address(&jito_vault_program::id()).0, + &vault_root.vault_pubkey, + operator_pubkey, + &operator_vault_ticket, + &vault_operator_delegation, + &vault_root.vault_admin, + &vault_root.vault_admin, + ) + .await?; + + Ok(()) + } + + pub async fn do_initialize_vault_ncn_slasher_ticket( + &mut self, + vault_root: &VaultRoot, + ncn_pubkey: &Pubkey, + slasher: &Pubkey, + ) -> Result<(), TestError> { + let vault_slasher_ticket_pubkey = VaultNcnSlasherTicket::find_program_address( + &jito_vault_program::id(), + &vault_root.vault_pubkey, + ncn_pubkey, + slasher, + ) + .0; + let ncn_slasher_ticket_pubkey = NcnVaultSlasherTicket::find_program_address( + &jito_restaking_program::id(), + ncn_pubkey, + &vault_root.vault_pubkey, + slasher, + ) + .0; + + self.initialize_vault_ncn_slasher_ticket( + &Config::find_program_address(&jito_vault_program::id()).0, + &vault_root.vault_pubkey, + ncn_pubkey, + slasher, + &ncn_slasher_ticket_pubkey, + &vault_slasher_ticket_pubkey, + &vault_root.vault_admin, + &vault_root.vault_admin, + ) + .await?; + + Ok(()) + } + + pub async fn do_warmup_vault_ncn_slasher_ticket( + &mut self, + vault_root: &VaultRoot, + ncn_pubkey: &Pubkey, + slasher: &Pubkey, + ) -> Result<(), TestError> { + let vault_slasher_ticket_pubkey = VaultNcnSlasherTicket::find_program_address( + &jito_vault_program::id(), + &vault_root.vault_pubkey, + ncn_pubkey, + slasher, + ) + .0; + + self.warmup_vault_ncn_slasher_ticket( + &Config::find_program_address(&jito_vault_program::id()).0, + &vault_root.vault_pubkey, + ncn_pubkey, + slasher, + &vault_slasher_ticket_pubkey, + &vault_root.vault_admin, + ) + .await?; + + Ok(()) + } + + pub async fn warmup_vault_ncn_slasher_ticket( + &mut self, + config: &Pubkey, + vault: &Pubkey, + ncn: &Pubkey, + slasher: &Pubkey, + vault_ncn_slasher_ticket: &Pubkey, + admin: &Keypair, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self._process_transaction(&Transaction::new_signed_with_payer( + &[warmup_vault_ncn_slasher_ticket( + &jito_vault_program::id(), + config, + vault, + ncn, + slasher, + vault_ncn_slasher_ticket, + &admin.pubkey(), + )], + Some(&admin.pubkey()), + &[admin], + blockhash, + )) + .await + } + + pub async fn do_add_delegation( + &mut self, + vault_root: &VaultRoot, + operator: &Pubkey, + amount: u64, + ) -> Result<(), TestError> { + self.add_delegation( + &Config::find_program_address(&jito_vault_program::id()).0, + &vault_root.vault_pubkey, + operator, + &VaultOperatorDelegation::find_program_address( + &jito_vault_program::id(), + &vault_root.vault_pubkey, + operator, + ) + .0, + &vault_root.vault_admin, + amount, + ) + .await?; + + Ok(()) + } + + pub async fn initialize_vault( + &mut self, + config: &Pubkey, + vault: &Pubkey, + vrt_mint: &Keypair, + token_mint: &Keypair, + vault_admin: &Keypair, + vault_base: &Keypair, + deposit_fee_bps: u16, + withdrawal_fee_bps: u16, + reward_fee_bps: u16, + decimals: u8, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self._process_transaction(&Transaction::new_signed_with_payer( + &[initialize_vault( + &jito_vault_program::id(), + config, + vault, + &vrt_mint.pubkey(), + &token_mint.pubkey(), + &vault_admin.pubkey(), + &vault_base.pubkey(), + deposit_fee_bps, + withdrawal_fee_bps, + reward_fee_bps, + decimals, + )], + Some(&vault_admin.pubkey()), + &[&vault_admin, &vrt_mint, &vault_base], + blockhash, + )) + .await + } + + pub async fn initialize_vault_ncn_ticket( + &mut self, + config: &Pubkey, + vault: &Pubkey, + ncn: &Pubkey, + ncn_vault_ticket: &Pubkey, + vault_ncn_ticket: &Pubkey, + admin: &Keypair, + payer: &Keypair, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::initialize_vault_ncn_ticket( + &jito_vault_program::id(), + config, + vault, + ncn, + ncn_vault_ticket, + vault_ncn_ticket, + &admin.pubkey(), + &payer.pubkey(), + )], + Some(&payer.pubkey()), + &[admin, payer], + blockhash, + )) + .await + } + + pub async fn initialize_vault_operator_delegation( + &mut self, + config: &Pubkey, + vault: &Pubkey, + operator: &Pubkey, + operator_vault_ticket: &Pubkey, + vault_operator_delegation: &Pubkey, + admin: &Keypair, + payer: &Keypair, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::initialize_vault_operator_delegation( + &jito_vault_program::id(), + config, + vault, + operator, + operator_vault_ticket, + vault_operator_delegation, + &admin.pubkey(), + &payer.pubkey(), + )], + Some(&payer.pubkey()), + &[admin, payer], + blockhash, + )) + .await + } + + pub async fn delegate_token_account( + &mut self, + config: &Pubkey, + vault: &Pubkey, + delegate_asset_admin: &Keypair, + token_mint: &Pubkey, + token_account: &Pubkey, + delegate: &Pubkey, + token_program_id: &Pubkey, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::delegate_token_account( + &jito_vault_program::id(), + config, + vault, + &delegate_asset_admin.pubkey(), + token_mint, + token_account, + delegate, + token_program_id, + )], + Some(&self.payer.pubkey()), + &[&self.payer, delegate_asset_admin], + blockhash, + )) + .await + } + + pub async fn set_admin( + &mut self, + config: &Pubkey, + vault: &Pubkey, + old_admin: &Keypair, + new_admin: &Keypair, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::set_admin( + &jito_vault_program::id(), + config, + vault, + &old_admin.pubkey(), + &new_admin.pubkey(), + )], + Some(&old_admin.pubkey()), + &[old_admin, new_admin], + blockhash, + )) + .await + } + + pub async fn set_secondary_admin( + &mut self, + config: &Pubkey, + vault: &Pubkey, + admin: &Keypair, + new_admin: &Pubkey, + role: VaultAdminRole, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::set_secondary_admin( + &jito_vault_program::id(), + config, + vault, + &admin.pubkey(), + new_admin, + role, + )], + Some(&admin.pubkey()), + &[admin], + blockhash, + )) + .await + } + + pub async fn set_fees( + &mut self, + config: &Pubkey, + vault: &Pubkey, + fee_admin: &Keypair, + deposit_fee_bps: Option, + withdrawal_fee_bps: Option, + reward_fee_bps: Option, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::set_fees( + &jito_vault_program::id(), + config, + vault, + &fee_admin.pubkey(), + deposit_fee_bps, + withdrawal_fee_bps, + reward_fee_bps, + )], + Some(&fee_admin.pubkey()), + &[fee_admin], + blockhash, + )) + .await + } + + pub async fn set_program_fee( + &mut self, + config_admin: &Keypair, + new_fee_bps: u16, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::set_program_fee( + &jito_vault_program::id(), + &Config::find_program_address(&jito_vault_program::id()).0, + &config_admin.pubkey(), + new_fee_bps, + )], + Some(&config_admin.pubkey()), + &[config_admin], + blockhash, + )) + .await + } + + pub async fn do_enqueue_withdrawal( + &mut self, + vault_root: &VaultRoot, + depositor: &Keypair, + amount: u64, + ) -> Result { + let vault = self.get_vault(&vault_root.vault_pubkey).await.unwrap(); + let depositor_vrt_token_account = + get_associated_token_address(&depositor.pubkey(), &vault.vrt_mint); + + let base = Keypair::new(); + let vault_staker_withdrawal_ticket = VaultStakerWithdrawalTicket::find_program_address( + &jito_vault_program::id(), + &vault_root.vault_pubkey, + &base.pubkey(), + ) + .0; + info!( + "vault_staker_withdrawal_ticket: {:?}", + vault_staker_withdrawal_ticket + ); + let vault_staker_withdrawal_ticket_token_account = + get_associated_token_address(&vault_staker_withdrawal_ticket, &vault.vrt_mint); + + self.create_ata(&vault.vrt_mint, &vault_staker_withdrawal_ticket) + .await?; + + self.enqueue_withdrawal( + &Config::find_program_address(&jito_vault_program::id()).0, + &vault_root.vault_pubkey, + &vault_staker_withdrawal_ticket, + &vault_staker_withdrawal_ticket_token_account, + depositor, + &depositor_vrt_token_account, + &base, + amount, + ) + .await?; + + Ok(VaultStakerWithdrawalTicketRoot { + base: base.pubkey(), + }) + } + + pub async fn do_cooldown_delegation( + &mut self, + vault_root: &VaultRoot, + operator: &Pubkey, + amount: u64, + ) -> TestResult<()> { + self.cooldown_delegation( + &Config::find_program_address(&jito_vault_program::id()).0, + &vault_root.vault_pubkey, + operator, + &VaultOperatorDelegation::find_program_address( + &jito_vault_program::id(), + &vault_root.vault_pubkey, + operator, + ) + .0, + &vault_root.vault_admin, + amount, + ) + .await + } + + pub async fn cooldown_delegation( + &mut self, + config: &Pubkey, + vault: &Pubkey, + operator: &Pubkey, + vault_operator_delegation: &Pubkey, + admin: &Keypair, + amount: u64, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self._process_transaction(&Transaction::new_signed_with_payer( + &[cooldown_delegation( + &jito_vault_program::id(), + config, + vault, + operator, + vault_operator_delegation, + &admin.pubkey(), + amount, + )], + Some(&self.payer.pubkey()), + &[&self.payer, admin], + blockhash, + )) + .await + } + + pub async fn do_full_vault_update( + &mut self, + vault_pubkey: &Pubkey, + operators: &[Pubkey], + ) -> Result<(), TestError> { + let slot = self.banks_client.get_sysvar::().await?.slot; + + let config = self + .get_config(&Config::find_program_address(&jito_vault_program::id()).0) + .await?; + + let ncn_epoch = slot / config.epoch_length(); + + let vault_update_state_tracker = VaultUpdateStateTracker::find_program_address( + &jito_vault_program::id(), + vault_pubkey, + ncn_epoch, + ) + .0; + self.initialize_vault_update_state_tracker(vault_pubkey, &vault_update_state_tracker) + .await?; + + for i in 0..operators.len() { + let operator_index = (i + ncn_epoch as usize) % operators.len(); + let operator = &operators[operator_index]; + self.crank_vault_update_state_tracker( + vault_pubkey, + operator, + &VaultOperatorDelegation::find_program_address( + &jito_vault_program::id(), + vault_pubkey, + operator, + ) + .0, + &vault_update_state_tracker, + ) + .await?; + } + + self.close_vault_update_state_tracker( + vault_pubkey, + &vault_update_state_tracker, + slot / config.epoch_length(), + ) + .await?; + + self.update_vault_balance(vault_pubkey).await?; + + Ok(()) + } + + pub async fn do_crank_vault_update_state_tracker( + &mut self, + vault: &Pubkey, + operator: &Pubkey, + ) -> TestResult<()> { + let slot = self.banks_client.get_sysvar::().await?.slot; + let config = self + .get_config(&Config::find_program_address(&jito_vault_program::id()).0) + .await?; + let ncn_epoch = slot / config.epoch_length(); + self.crank_vault_update_state_tracker( + vault, + operator, + &VaultOperatorDelegation::find_program_address( + &jito_vault_program::id(), + vault, + operator, + ) + .0, + &VaultUpdateStateTracker::find_program_address( + &jito_vault_program::id(), + vault, + ncn_epoch, + ) + .0, + ) + .await + } + + pub async fn crank_vault_update_state_tracker( + &mut self, + vault: &Pubkey, + operator: &Pubkey, + vault_operator_delegation: &Pubkey, + vault_update_state_tracker: &Pubkey, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::crank_vault_update_state_tracker( + &jito_vault_program::id(), + &Config::find_program_address(&jito_vault_program::id()).0, + vault, + operator, + vault_operator_delegation, + vault_update_state_tracker, + )], + Some(&self.payer.pubkey()), + &[&self.payer], + blockhash, + )) + .await?; + Ok(()) + } + + pub async fn update_vault_balance(&mut self, vault_pubkey: &Pubkey) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + let vault = self.get_vault(vault_pubkey).await?; + + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::update_vault_balance( + &jito_vault_program::id(), + &Config::find_program_address(&jito_vault_program::id()).0, + vault_pubkey, + &get_associated_token_address(vault_pubkey, &vault.supported_mint), + &vault.vrt_mint, + &get_associated_token_address(&vault.fee_wallet, &vault.vrt_mint), + &spl_token::ID, + )], + Some(&self.payer.pubkey()), + &[&self.payer], + blockhash, + )) + .await?; + + Ok(()) + } + + pub async fn initialize_vault_update_state_tracker( + &mut self, + vault_pubkey: &Pubkey, + vault_update_state_tracker: &Pubkey, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::initialize_vault_update_state_tracker( + &jito_vault_program::id(), + &Config::find_program_address(&jito_vault_program::id()).0, + vault_pubkey, + vault_update_state_tracker, + &self.payer.pubkey(), + WithdrawalAllocationMethod::Greedy, + )], + Some(&self.payer.pubkey()), + &[&self.payer], + blockhash, + )) + .await?; + Ok(()) + } + + pub async fn close_vault_update_state_tracker( + &mut self, + vault_pubkey: &Pubkey, + vault_update_state_tracker: &Pubkey, + ncn_epoch: u64, + ) -> TestResult<()> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::close_vault_update_state_tracker( + &jito_vault_program::id(), + &Config::find_program_address(&jito_vault_program::id()).0, + vault_pubkey, + vault_update_state_tracker, + &self.payer.pubkey(), + ncn_epoch, + )], + Some(&self.payer.pubkey()), + &[&self.payer], + blockhash, + )) + .await + } + + pub async fn enqueue_withdrawal( + &mut self, + config: &Pubkey, + vault: &Pubkey, + vault_staker_withdrawal_ticket: &Pubkey, + vault_staker_withdrawal_ticket_token_account: &Pubkey, + staker: &Keypair, + staker_vrt_token_account: &Pubkey, + base: &Keypair, + amount: u64, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::enqueue_withdrawal( + &jito_vault_program::id(), + config, + vault, + vault_staker_withdrawal_ticket, + vault_staker_withdrawal_ticket_token_account, + &staker.pubkey(), + staker_vrt_token_account, + &base.pubkey(), + amount, + )], + Some(&staker.pubkey()), + &[staker, base], + blockhash, + )) + .await + } + + pub async fn do_burn_withdrawal_ticket( + &mut self, + vault_root: &VaultRoot, + staker: &Keypair, + vault_staker_withdrawal_ticket_base: &Pubkey, + program_fee_wallet: &Pubkey, + ) -> Result<(), TestError> { + let vault = self.get_vault(&vault_root.vault_pubkey).await.unwrap(); + let vault_staker_withdrawal_ticket = VaultStakerWithdrawalTicket::find_program_address( + &jito_vault_program::id(), + &vault_root.vault_pubkey, + vault_staker_withdrawal_ticket_base, + ) + .0; + + self.burn_withdrawal_ticket( + &Config::find_program_address(&jito_vault_program::id()).0, + &vault_root.vault_pubkey, + &get_associated_token_address(&vault_root.vault_pubkey, &vault.supported_mint), + &vault.vrt_mint, + &staker.pubkey(), + &get_associated_token_address(&staker.pubkey(), &vault.supported_mint), + &vault_staker_withdrawal_ticket, + &get_associated_token_address(&vault_staker_withdrawal_ticket, &vault.vrt_mint), + &get_associated_token_address(&vault.fee_wallet, &vault.vrt_mint), + &get_associated_token_address(program_fee_wallet, &vault.vrt_mint), + ) + .await?; + + Ok(()) + } + + pub async fn burn_withdrawal_ticket( + &mut self, + config: &Pubkey, + vault: &Pubkey, + vault_token_account: &Pubkey, + vrt_mint: &Pubkey, + staker: &Pubkey, + staker_token_account: &Pubkey, + vault_staker_withdrawal_ticket: &Pubkey, + vault_staker_withdrawal_ticket_token_account: &Pubkey, + vault_fee_token_account: &Pubkey, + program_fee_vrt_token_account: &Pubkey, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::burn_withdrawal_ticket( + &jito_vault_program::id(), + config, + vault, + vault_token_account, + vrt_mint, + staker, + staker_token_account, + vault_staker_withdrawal_ticket, + vault_staker_withdrawal_ticket_token_account, + vault_fee_token_account, + program_fee_vrt_token_account, + )], + Some(&self.payer.pubkey()), + &[&self.payer], + blockhash, + )) + .await + } + + pub async fn add_delegation( + &mut self, + config: &Pubkey, + vault: &Pubkey, + operator: &Pubkey, + vault_operator_delegation: &Pubkey, + admin: &Keypair, + amount: u64, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self._process_transaction(&Transaction::new_signed_with_payer( + &[add_delegation( + &jito_vault_program::id(), + config, + vault, + operator, + vault_operator_delegation, + &admin.pubkey(), + amount, + )], + Some(&admin.pubkey()), + &[admin], + blockhash, + )) + .await + } + + pub async fn do_mint_to( + &mut self, + vault_root: &VaultRoot, + depositor: &Keypair, + amount_in: u64, + min_amount_out: u64, + ) -> TestResult<()> { + let vault = self.get_vault(&vault_root.vault_pubkey).await.unwrap(); + self.mint_to( + &vault_root.vault_pubkey, + &vault.vrt_mint, + depositor, + &get_associated_token_address(&depositor.pubkey(), &vault.supported_mint), + &get_associated_token_address(&vault_root.vault_pubkey, &vault.supported_mint), + &get_associated_token_address(&depositor.pubkey(), &vault.vrt_mint), + &get_associated_token_address(&vault.fee_wallet, &vault.vrt_mint), + None, + amount_in, + min_amount_out, + ) + .await + } + + pub async fn mint_to( + &mut self, + vault: &Pubkey, + vrt_mint: &Pubkey, + depositor: &Keypair, + depositor_token_account: &Pubkey, + vault_token_account: &Pubkey, + depositor_vrt_token_account: &Pubkey, + vault_fee_token_account: &Pubkey, + mint_signer: Option<&Keypair>, + amount_in: u64, + min_amount_out: u64, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + let mut signers = vec![depositor]; + if let Some(signer) = mint_signer { + signers.push(signer); + } + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::mint_to( + &jito_vault_program::id(), + &Config::find_program_address(&jito_vault_program::id()).0, + vault, + vrt_mint, + &depositor.pubkey(), + depositor_token_account, + vault_token_account, + depositor_vrt_token_account, + vault_fee_token_account, + mint_signer.map(|s| s.pubkey()).as_ref(), + amount_in, + min_amount_out, + )], + Some(&depositor.pubkey()), + &signers, + blockhash, + )) + .await + } + + pub async fn initialize_vault_ncn_slasher_ticket( + &mut self, + config: &Pubkey, + vault: &Pubkey, + ncn: &Pubkey, + slasher: &Pubkey, + ncn_slasher_ticket: &Pubkey, + vault_slasher_ticket: &Pubkey, + admin: &Keypair, + payer: &Keypair, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::initialize_vault_ncn_slasher_ticket( + &jito_vault_program::id(), + config, + vault, + ncn, + slasher, + ncn_slasher_ticket, + vault_slasher_ticket, + &admin.pubkey(), + &payer.pubkey(), + )], + Some(&payer.pubkey()), + &[admin, payer], + blockhash, + )) + .await + } + + #[allow(dead_code)] + pub async fn initialize_vault_ncn_slasher_operator_ticket( + &mut self, + config: &Pubkey, + vault: &Pubkey, + ncn: &Pubkey, + slasher: &Pubkey, + operator: &Pubkey, + vault_ncn_slasher_ticket: &Pubkey, + vault_ncn_slasher_operator_ticket: &Pubkey, + payer: &Keypair, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self._process_transaction(&Transaction::new_signed_with_payer( + &[ + jito_vault_sdk::sdk::initialize_vault_ncn_slasher_operator_ticket( + &jito_vault_program::id(), + config, + vault, + ncn, + slasher, + operator, + vault_ncn_slasher_ticket, + vault_ncn_slasher_operator_ticket, + &payer.pubkey(), + ), + ], + Some(&payer.pubkey()), + &[payer], + blockhash, + )) + .await + } + + pub async fn create_token_metadata( + &mut self, + vault: &Pubkey, + admin: &Keypair, + vrt_mint: &Pubkey, + payer: &Keypair, + metadata: &Pubkey, + name: String, + symbol: String, + uri: String, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::create_token_metadata( + &jito_vault_program::id(), + vault, + &admin.pubkey(), + vrt_mint, + &payer.pubkey(), + metadata, + name, + symbol, + uri, + )], + Some(&payer.pubkey()), + &[admin, payer], + blockhash, + )) + .await + } + + pub async fn update_token_metadata( + &mut self, + vault: &Pubkey, + admin: &Keypair, + vrt_mint: &Pubkey, + metadata: &Pubkey, + name: String, + symbol: String, + uri: String, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::update_token_metadata( + &jito_vault_program::id(), + vault, + &admin.pubkey(), + vrt_mint, + metadata, + name, + symbol, + uri, + )], + Some(&self.payer.pubkey()), + &[&self.payer, admin], + blockhash, + )) + .await + } + + async fn _process_transaction(&mut self, tx: &Transaction) -> Result<(), TestError> { + self.banks_client + .process_transaction_with_preflight_and_commitment( + tx.clone(), + CommitmentLevel::Processed, + ) + .await?; + Ok(()) + } + + pub async fn airdrop(&mut self, to: &Pubkey, sol: f64) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self.banks_client + .process_transaction_with_preflight_and_commitment( + Transaction::new_signed_with_payer( + &[transfer(&self.payer.pubkey(), to, sol_to_lamports(sol))], + Some(&self.payer.pubkey()), + &[&self.payer], + blockhash, + ), + CommitmentLevel::Processed, + ) + .await?; + Ok(()) + } + + pub async fn create_token_mint( + &mut self, + mint: &Keypair, + token_program_id: &Pubkey, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + let rent: Rent = self.banks_client.get_sysvar().await?; + let ixs = vec![ + create_account( + &self.payer.pubkey(), + &mint.pubkey(), + rent.minimum_balance(spl_token::state::Mint::LEN), + spl_token::state::Mint::LEN as u64, + token_program_id, + ), + spl_token::instruction::initialize_mint2( + token_program_id, + &mint.pubkey(), + &self.payer.pubkey(), + None, + 9, + ) + .unwrap(), + ]; + self.banks_client + .process_transaction_with_preflight_and_commitment( + Transaction::new_signed_with_payer( + &ixs, + Some(&self.payer.pubkey()), + &[&self.payer, mint], + blockhash, + ), + CommitmentLevel::Processed, + ) + .await?; + Ok(()) + } + + pub async fn create_ata(&mut self, mint: &Pubkey, owner: &Pubkey) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self.banks_client + .process_transaction_with_preflight_and_commitment( + Transaction::new_signed_with_payer( + &[create_associated_token_account_idempotent( + &self.payer.pubkey(), + owner, + mint, + &spl_token::id(), + )], + Some(&self.payer.pubkey()), + &[&self.payer], + blockhash, + ), + CommitmentLevel::Processed, + ) + .await?; + Ok(()) + } + + /// Mints tokens to an ATA owned by the `to` address + pub async fn mint_spl_to( + &mut self, + mint: &Pubkey, + to: &Pubkey, + amount: u64, + ) -> Result<(), BanksClientError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self.banks_client + .process_transaction_with_preflight_and_commitment( + Transaction::new_signed_with_payer( + &[ + create_associated_token_account_idempotent( + &self.payer.pubkey(), + to, + mint, + &spl_token::id(), + ), + spl_token::instruction::mint_to( + &spl_token::id(), + mint, + &get_associated_token_address(to, mint), + &self.payer.pubkey(), + &[], + amount, + ) + .unwrap(), + ], + Some(&self.payer.pubkey()), + &[&self.payer], + blockhash, + ), + CommitmentLevel::Processed, + ) + .await + } + + pub async fn get_reward_fee_token_account( + &mut self, + vault: &Pubkey, + ) -> Result { + let vault = self.get_vault(vault).await.unwrap(); + + let vault_fee_token_account = + get_associated_token_address(&vault.fee_wallet, &vault.vrt_mint); + + let account = self + .banks_client + .get_account(vault_fee_token_account) + .await + .unwrap() + .unwrap(); + + Ok(SPLTokenAccount::unpack(&account.data).unwrap()) + } + + pub async fn create_and_fund_reward_vault( + &mut self, + vault: &Pubkey, + rewarder: &Keypair, + amount: u64, + ) -> Result<(), BanksClientError> { + let vault_account = self.get_vault(vault).await.unwrap(); + + let rewarder_token_account = + get_associated_token_address(&rewarder.pubkey(), &vault_account.supported_mint); + + let vault_token_account = + get_associated_token_address(vault, &vault_account.supported_mint); + + let blockhash = self.banks_client.get_latest_blockhash().await?; + self.banks_client + .process_transaction_with_preflight_and_commitment( + Transaction::new_signed_with_payer( + &[ + create_associated_token_account_idempotent( + &rewarder.pubkey(), + &vault_token_account, + &vault_account.supported_mint, + &spl_token::id(), + ), + spl_token::instruction::transfer( + &spl_token::id(), + &rewarder_token_account, + &vault_token_account, + &rewarder.pubkey(), + &[], + amount, + ) + .unwrap(), + ], + Some(&rewarder.pubkey()), + &[&rewarder], + blockhash, + ), + CommitmentLevel::Processed, + ) + .await + } + + pub async fn set_program_fee_wallet( + &mut self, + program_fee_admin: &Keypair, + new_fee_wallet: &Pubkey, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::set_program_fee_wallet( + &jito_vault_program::id(), + &Config::find_program_address(&jito_vault_program::id()).0, + &program_fee_admin.pubkey(), + new_fee_wallet, + )], + Some(&program_fee_admin.pubkey()), + &[program_fee_admin], + blockhash, + )) + .await + } + + pub async fn set_is_paused( + &mut self, + vault: &Pubkey, + admin: &Keypair, + is_paused: bool, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::set_is_paused( + &jito_vault_program::id(), + &Config::find_program_address(&jito_vault_program::id()).0, + vault, + &admin.pubkey(), + is_paused, + )], + Some(&admin.pubkey()), + &[admin], + blockhash, + )) + .await + } + + pub async fn set_config_admin( + &mut self, + config: &Pubkey, + old_admin: &Keypair, + new_admin: &Keypair, + ) -> Result<(), TestError> { + let blockhash = self.banks_client.get_latest_blockhash().await?; + self._process_transaction(&Transaction::new_signed_with_payer( + &[jito_vault_sdk::sdk::set_config_admin( + &jito_vault_program::id(), + config, + &old_admin.pubkey(), + &new_admin.pubkey(), + )], + Some(&old_admin.pubkey()), + &[old_admin], + blockhash, + )) + .await + } +} + +#[inline(always)] +#[track_caller] +pub fn assert_vault_error(test_error: Result, vault_error: VaultError) { + assert!(test_error.is_err()); + assert_eq!( + test_error.err().unwrap().to_transaction_error().unwrap(), + TransactionError::InstructionError(0, InstructionError::Custom(vault_error as u32)) + ); } diff --git a/integration_tests/tests/tip_router/register_mint.rs b/integration_tests/tests/tip_router/register_mint.rs index 27c6c20..72d7381 100644 --- a/integration_tests/tests/tip_router/register_mint.rs +++ b/integration_tests/tests/tip_router/register_mint.rs @@ -1,5 +1,7 @@ #[cfg(test)] mod tests { + use jito_restaking_core::{config::Config, ncn_vault_ticket::NcnVaultTicket}; + use jito_vault_core::vault_ncn_ticket::VaultNcnTicket; use solana_sdk::{signature::Keypair, signer::Signer}; use crate::fixtures::{test_builder::TestBuilder, TestResult}; @@ -8,23 +10,61 @@ mod tests { async fn test_register_mint_success() -> TestResult<()> { let mut fixture = TestBuilder::new().await; let mut tip_router_client = fixture.tip_router_client(); + let mut vault_client = fixture.vault_client(); + let mut restaking_client = fixture.restaking_program_client(); let ncn_root = fixture.setup_ncn().await?; // Setup initial state tip_router_client.setup_tip_router(&ncn_root).await?; // Setup vault and tickets - let vault = Keypair::new(); - let vault_ncn_ticket = Keypair::new(); - let ncn_vault_ticket = Keypair::new(); + let _ = vault_client.do_initialize_config().await?; + let vault_root = vault_client + .do_initialize_vault(0, 0, 0, 9, &ncn_root.ncn_pubkey) + .await?; + restaking_client + .do_initialize_ncn_vault_ticket(&ncn_root, &vault_root.vault_pubkey) + .await?; + vault_client + .do_initialize_vault_ncn_ticket(&vault_root, &ncn_root.ncn_pubkey) + .await?; + + let vault = vault_root.vault_pubkey; + let vault_ncn_ticket = VaultNcnTicket::find_program_address( + &jito_vault_program::id(), + &vault_root.vault_pubkey, + &ncn_root.ncn_pubkey, + ) + .0; + let ncn_vault_ticket = NcnVaultTicket::find_program_address( + &jito_restaking_program::id(), + &ncn_root.ncn_pubkey, + &vault_root.vault_pubkey, + ) + .0; + + fixture.warp_slot_incremental(2).await?; + + vault_client + .do_warmup_vault_ncn_ticket(&vault_root, &ncn_root.ncn_pubkey) + .await?; + restaking_client + .do_warmup_ncn_vault_ticket(&ncn_root, &vault_root.vault_pubkey) + .await?; + let restaking_config_pubkey = Config::find_program_address(&jito_restaking_program::id()).0; + let epoch_length = restaking_client + .get_config(&restaking_config_pubkey) + .await? + .epoch_length(); + fixture.warp_slot_incremental(2 * epoch_length).await?; // Register mint tip_router_client .do_register_mint( ncn_root.ncn_pubkey, - vault.pubkey(), - vault_ncn_ticket.pubkey(), - ncn_vault_ticket.pubkey(), + vault, + vault_ncn_ticket, + ncn_vault_ticket, ) .await?; @@ -66,33 +106,73 @@ mod tests { async fn test_register_mint_duplicate() -> TestResult<()> { let mut fixture = TestBuilder::new().await; let mut tip_router_client = fixture.tip_router_client(); + let mut vault_client = fixture.vault_client(); + let mut restaking_client = fixture.restaking_program_client(); let ncn_root = fixture.setup_ncn().await?; // Setup initial state tip_router_client.setup_tip_router(&ncn_root).await?; // Setup vault and tickets - let vault = Keypair::new(); - let vault_ncn_ticket = Keypair::new(); - let ncn_vault_ticket = Keypair::new(); + let _ = vault_client.do_initialize_config().await?; + let vault_root = vault_client + .do_initialize_vault(0, 0, 0, 9, &ncn_root.ncn_pubkey) + .await?; + restaking_client + .do_initialize_ncn_vault_ticket(&ncn_root, &vault_root.vault_pubkey) + .await?; + vault_client + .do_initialize_vault_ncn_ticket(&vault_root, &ncn_root.ncn_pubkey) + .await?; + + let vault = vault_root.vault_pubkey; + let vault_ncn_ticket = VaultNcnTicket::find_program_address( + &jito_vault_program::id(), + &vault_root.vault_pubkey, + &ncn_root.ncn_pubkey, + ) + .0; + let ncn_vault_ticket = NcnVaultTicket::find_program_address( + &jito_restaking_program::id(), + &ncn_root.ncn_pubkey, + &vault_root.vault_pubkey, + ) + .0; + + fixture.warp_slot_incremental(2).await?; + + vault_client + .do_warmup_vault_ncn_ticket(&vault_root, &ncn_root.ncn_pubkey) + .await?; + restaking_client + .do_warmup_ncn_vault_ticket(&ncn_root, &vault_root.vault_pubkey) + .await?; + let restaking_config_pubkey = Config::find_program_address(&jito_restaking_program::id()).0; + let epoch_length = restaking_client + .get_config(&restaking_config_pubkey) + .await? + .epoch_length(); + fixture.warp_slot_incremental(2 * epoch_length).await?; // Register mint first time tip_router_client .do_register_mint( ncn_root.ncn_pubkey, - vault.pubkey(), - vault_ncn_ticket.pubkey(), - ncn_vault_ticket.pubkey(), + vault, + vault_ncn_ticket, + ncn_vault_ticket, ) .await?; + fixture.warp_slot_incremental(1).await?; + // Register same mint again tip_router_client .do_register_mint( ncn_root.ncn_pubkey, - vault.pubkey(), - vault_ncn_ticket.pubkey(), - ncn_vault_ticket.pubkey(), + vault, + vault_ncn_ticket, + ncn_vault_ticket, ) .await?; @@ -109,12 +189,75 @@ mod tests { async fn test_register_mint_fails_with_weight_table() -> TestResult<()> { let mut fixture = TestBuilder::new().await; let mut tip_router_client = fixture.tip_router_client(); + let mut vault_client = fixture.vault_client(); + let mut restaking_client = fixture.restaking_program_client(); let ncn_root = fixture.setup_ncn().await?; tip_router_client.setup_tip_router(&ncn_root).await?; - // TODO create ncn and vault with 1 mint, register mint, initialize weight table - // TODO verify weight table locks register_mint + let _ = vault_client.do_initialize_config().await?; + let vault_root = vault_client + .do_initialize_vault(0, 0, 0, 9, &ncn_root.ncn_pubkey) + .await?; + restaking_client + .do_initialize_ncn_vault_ticket(&ncn_root, &vault_root.vault_pubkey) + .await?; + vault_client + .do_initialize_vault_ncn_ticket(&vault_root, &ncn_root.ncn_pubkey) + .await?; + + let vault = vault_root.vault_pubkey; + let vault_ncn_ticket = VaultNcnTicket::find_program_address( + &jito_vault_program::id(), + &vault_root.vault_pubkey, + &ncn_root.ncn_pubkey, + ) + .0; + let ncn_vault_ticket = NcnVaultTicket::find_program_address( + &jito_restaking_program::id(), + &ncn_root.ncn_pubkey, + &vault_root.vault_pubkey, + ) + .0; + + fixture.warp_slot_incremental(2).await?; + + vault_client + .do_warmup_vault_ncn_ticket(&vault_root, &ncn_root.ncn_pubkey) + .await?; + restaking_client + .do_warmup_ncn_vault_ticket(&ncn_root, &vault_root.vault_pubkey) + .await?; + let restaking_config_pubkey = Config::find_program_address(&jito_restaking_program::id()).0; + let epoch_length = restaking_client + .get_config(&restaking_config_pubkey) + .await? + .epoch_length(); + fixture.warp_slot_incremental(2 * epoch_length).await?; + + tip_router_client + .do_register_mint( + ncn_root.ncn_pubkey, + vault, + vault_ncn_ticket, + ncn_vault_ticket, + ) + .await?; + + tip_router_client + .initialize_weight_table(ncn_root.ncn_pubkey, fixture.clock().await.slot) + .await?; + + let result = tip_router_client + .do_register_mint( + ncn_root.ncn_pubkey, + vault, + vault_ncn_ticket, + ncn_vault_ticket, + ) + .await; + + assert!(result.is_err()); Ok(()) } diff --git a/program/src/register_mint.rs b/program/src/register_mint.rs index 8fe2b0e..ae8fcd8 100644 --- a/program/src/register_mint.rs +++ b/program/src/register_mint.rs @@ -45,6 +45,7 @@ pub fn process_register_mint(program_id: &Pubkey, accounts: &[AccountInfo]) -> P .checked_div(epoch_length) .ok_or(TipRouterError::DenominatorIsZero)?; + // TODO: is there a way to DOS this by changing program owner to tiprouter or something before weight table is initialized? // Once tracked_mints.mint_count() == ncn.vault_count, the weight table can be initialized // Once the weight table is initialized, you can't add any more mints if weight_table.owner.eq(&system_program::ID) {