diff --git a/programs/registrar/src/deposit_governance_token.rs b/programs/registrar/src/deposit_governance_token.rs new file mode 100644 index 00000000..9be94a2d --- /dev/null +++ b/programs/registrar/src/deposit_governance_token.rs @@ -0,0 +1,91 @@ +use crate::{ + error::RegistrarError, + registrar_config::RegistrarConfig, + voter_weight_record::VoterWeightRecord, +}; +use solana_program::{ + account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey, sysvar::clock::Clock, +}; +use spl_token::instruction::transfer; + +#[derive(Accounts)] +pub struct DepositGovernanceToken<'info> { + #[account(mut)] + pub voter_weight_record: Account<'info, VoterWeightRecord>, + #[account(mut)] + pub voter_token_account: AccountInfo<'info>, + #[account(mut)] + pub governance_token_mint: AccountInfo<'info>, + pub registrar_config: AccountInfo<'info>, + pub token_program: AccountInfo<'info>, + pub authority: AccountInfo<'info>, +} + +pub fn deposit_governance_token( + ctx: Context, + amount: u64, +) -> ProgramResult { + let registrar_config = RegistrarConfig::unpack_from_slice(&ctx.accounts.registrar_config.data.borrow())?; + + // Check if the token is an accepted governance token + let token_mint = ctx.accounts.governance_token_mint.key(); + if !registrar_config.accepted_tokens.contains(&token_mint) { + return Err(RegistrarError::InvalidArgument.into()); + } + + // Transfer tokens from the voter's account to the governance program's account + transfer_tokens( + &ctx.accounts.voter_token_account, + &ctx.accounts.governance_token_mint, // Replace with the governance program's token account + amount, + &ctx.accounts.token_program, + &ctx.accounts.authority, + )?; + + // Update the VoterWeightRecord account + let voter_weight_record = &mut ctx.accounts.voter_weight_record; + let clock = Clock::get()?; + + if voter_weight_record.last_deposit_or_withdrawal_slot == clock.slot { + return Err(RegistrarError::InvalidOperation.into()); + } + + let weight_increase = amount * registrar_config.weights[registrar_config.accepted_tokens.iter().position(|&token| token == token_mint).ok_or(RegistrarError::InvalidArgument)?]; + voter_weight_record.weight = voter_weight_record.weight.checked_add(weight_increase).ok_or(RegistrarError::Overflow)?; + voter_weight_record.last_deposit_or_withdrawal_slot = clock.slot; + voter_weight_record.serialize(&mut ctx.accounts.voter_weight_record.data.borrow_mut()[..])?; + + // Update the MaxVoterWeightRecord account + // ... + + Ok(()) +} + +fn transfer_tokens( + source_account: &AccountInfo, + destination_account: &AccountInfo, + amount: u64, + token_program: &AccountInfo, + authority: &AccountInfo, +) -> ProgramResult { + let transfer_instruction = transfer( + token_program.key, + source_account.key, + destination_account.key, + authority.key, + &[], + amount, + )?; + + invoke( + &transfer_instruction, + &[ + source_account.clone(), + destination_account.clone(), + authority.clone(), + token_program.clone(), + ], + )?; + + Ok(()) +} diff --git a/programs/registrar/src/error.rs b/programs/registrar/src/error.rs new file mode 100644 index 00000000..a1358e67 --- /dev/null +++ b/programs/registrar/src/error.rs @@ -0,0 +1,15 @@ +#[derive(Debug, PartialEq, Eq)] +pub enum RegistrarError { + InvalidArgument, + InvalidAccountData, + InvalidOperation, + Overflow, + InsufficientFunds, + // Add more error variants as needed +} + +impl From for ProgramError { + fn from(error: RegistrarError) -> Self { + ProgramError::Custom(error as u32) + } +} diff --git a/programs/registrar/src/initialize_registrar.rs b/programs/registrar/src/initialize_registrar.rs new file mode 100644 index 00000000..03926592 --- /dev/null +++ b/programs/registrar/src/initialize_registrar.rs @@ -0,0 +1,35 @@ +use crate::{error::RegistrarError, registrar_config::RegistrarConfig}; +use solana_program::{ + account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey, system_program, +}; + +#[derive(Accounts)] +pub struct InitializeRegistrar<'info> { + #[account(init, payer = payer, space = RegistrarConfig::LEN)] + pub registrar_config: Account<'info, RegistrarConfig>, + #[account(mut)] + pub payer: AccountInfo<'info>, + pub system_program: Program<'info, System>, +} + +pub fn initialize_registrar( + ctx: Context, + accepted_tokens: Vec, + weights: Vec, +) -> ProgramResult { + if accepted_tokens.len() != weights.len() { + return Err(RegistrarError::InvalidArgument.into()); + } + + if accepted_tokens.len() > MAX_ACCEPTED_TOKENS { + return Err(RegistrarError::InvalidArgument.into()); + } + + let registrar_info = &mut ctx.accounts.registrar_config; + let config = RegistrarConfig { + accepted_tokens, + weights, + }; + config.pack_into_slice(&mut registrar_info.data.borrow_mut()); + Ok(()) +} diff --git a/programs/registrar/src/lib.rs b/programs/registrar/src/lib.rs new file mode 100644 index 00000000..9abc1841 --- /dev/null +++ b/programs/registrar/src/lib.rs @@ -0,0 +1,52 @@ +pub mod error; +pub mod initialize_registrar; +pub mod deposit_governance_token; +pub mod withdraw_governance_token; +pub mod registrar_config; +pub mod voter_weight_record; +pub mod max_voter_weight_record; + +use solana_program::{ + account_info::AccountInfo, entrypoint, entrypoint::ProgramResult, pubkey::Pubkey, +}; + +entrypoint!(process_instruction); + +pub fn process_instruction( + program_id: &Pubkey, + accounts: &[AccountInfo], + instruction_data: &[u8], +) -> ProgramResult { + let instruction = RegistrarInstruction::try_from_slice(instruction_data)?; + + match instruction { + RegistrarInstruction::InitializeRegistrar { + accepted_tokens, + weights, + } => { + let accounts = initialize_registrar::InitializeRegistrar::try_accounts( + program_id, + accounts, + &accepted_tokens, + &weights, + )?; + initialize_registrar::initialize_registrar(accounts, accepted_tokens, weights) + } + RegistrarInstruction::DepositGovernanceToken { amount } => { + let accounts = deposit_governance_token::DepositGovernanceToken::try_accounts( + program_id, + accounts, + &amount, + )?; + deposit_governance_token::deposit_governance_token(accounts, amount) + } + RegistrarInstruction::WithdrawGovernanceToken { amount } => { + let accounts = withdraw_governance_token::WithdrawGovernanceToken::try_accounts( + program_id, + accounts, + &amount, + )?; + withdraw_governance_token::withdraw_governance_token(accounts, amount) + } + } +} diff --git a/programs/registrar/src/max_voter_weight_record.rs b/programs/registrar/src/max_voter_weight_record.rs new file mode 100644 index 00000000..b797864b --- /dev/null +++ b/programs/registrar/src/max_voter_weight_record.rs @@ -0,0 +1,25 @@ +use borsh::{BorshDeserialize, BorshSerialize}; +use solana_program::program_error::ProgramError; + +#[derive(BorshSerialize, BorshDeserialize)] +pub struct MaxVoterWeightRecord { + pub max_weight: u64, +} + +impl Sealed for MaxVoterWeightRecord {} + +impl Pack for MaxVoterWeightRecord { + const LEN: usize = 8; + + fn unpack_from_slice(src: &[u8]) -> Result { + let max_weight = src.get(..8).ok_or(ProgramError::InvalidAccountData)?.get_u64(); + Ok(MaxVoterWeightRecord { max_weight }) + } + + fn pack_into_slice(&self, dst: &mut [u8]) { + dst.get_mut(..8) + .ok_or(ProgramError::InvalidAccountData) + .unwrap() + .copy_from_slice(&self.max_weight.to_le_bytes()); + } +} diff --git a/programs/registrar/src/registrar_config.rs b/programs/registrar/src/registrar_config.rs new file mode 100644 index 00000000..12017644 --- /dev/null +++ b/programs/registrar/src/registrar_config.rs @@ -0,0 +1,34 @@ +use borsh::{BorshDeserialize, BorshSerialize}; +use solana_program::{program_error::ProgramError, pubkey::Pubkey}; + +const MAX_ACCEPTED_TOKENS: usize = 10; + +#[derive(BorshSerialize, BorshDeserialize)] +pub struct RegistrarConfig { + pub accepted_tokens: Vec, + pub weights: Vec, +} + +impl Sealed for RegistrarConfig {} + +impl Pack for RegistrarConfig { + const LEN: usize = 8 + (4 + 32) * MAX_ACCEPTED_TOKENS; + + fn unpack_from_slice(src: &[u8]) -> Result { + let (accepted_tokens, weights) = array_refs![src, MAX_ACCEPTED_TOKENS; Pubkey, u64]; + let accepted_tokens = accepted_tokens.to_vec(); + let weights = weights.to_vec(); + Ok(RegistrarConfig { + accepted_tokens, + weights, + }) + } + + fn pack_into_slice(&self, dst: &mut [u8]) { + let (accepted_tokens, weights) = array_mut_refs![dst, MAX_ACCEPTED_TOKENS; Pubkey, u64]; + for (i, token) in self.accepted_tokens.iter().enumerate() { + accepted_tokens[i] = *token; + weights[i] = self.weights[i]; + } + } +} diff --git a/programs/registrar/src/voter_weight_record.rs b/programs/registrar/src/voter_weight_record.rs new file mode 100644 index 00000000..ff5cc2cc --- /dev/null +++ b/programs/registrar/src/voter_weight_record.rs @@ -0,0 +1,33 @@ +use borsh::{BorshDeserialize, BorshSerialize}; +use solana_program::{program_error::ProgramError, pubkey::Pubkey}; + +#[derive(BorshSerialize, BorshDeserialize)] +pub struct VoterWeightRecord { + pub voter: Pubkey, + pub weight: u64, + pub last_deposit_or_withdrawal_slot: u64, +} + +impl Sealed for VoterWeightRecord {} + +impl Pack for VoterWeightRecord { + const LEN: usize = 8 + 32 + 8; + + fn unpack_from_slice(src: &[u8]) -> Result { + let (voter, weight, last_deposit_or_withdrawal_slot) = + array_refs![src, 0, 32, 40; Pubkey, u64, u64]; + Ok(VoterWeightRecord { + voter: *voter, + weight: *weight, + last_deposit_or_withdrawal_slot: *last_deposit_or_withdrawal_slot, + }) + } + + fn pack_into_slice(&self, dst: &mut [u8]) { + let (voter, weight, last_deposit_or_withdrawal_slot) = + array_mut_refs![dst, 0, 32, 40; Pubkey, u64, u64]; + *voter = self.voter; + *weight = self.weight; + *last_deposit_or_withdrawal_slot = self.last_deposit_or_withdrawal_slot; + } +} diff --git a/programs/registrar/src/withdraw_governance_token.rs b/programs/registrar/src/withdraw_governance_token.rs new file mode 100644 index 00000000..b2ca68a9 --- /dev/null +++ b/programs/registrar/src/withdraw_governance_token.rs @@ -0,0 +1,99 @@ +use crate::{ + error::RegistrarError, + registrar_config::RegistrarConfig, + voter_weight_record::VoterWeightRecord, +}; +use solana_program::{ + account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey, sysvar::clock::Clock, +}; +use spl_token::instruction::transfer; + +#[derive(Accounts)] +pub struct WithdrawGovernanceToken<'info> { + #[account(mut)] + pub voter_weight_record: Account<'info, VoterWeightRecord>, + #[account(mut)] + pub voter_token_account: AccountInfo<'info>, + pub governance_token_mint: AccountInfo<'info>, + pub registrar_config: AccountInfo<'info>, + pub token_program: AccountInfo<'info>, + pub spl_governance_program: AccountInfo<'info>, + pub authority: AccountInfo<'info>, +} + +pub fn withdraw_governance_token( + ctx: Context, + amount: u64, +) -> ProgramResult { + let voter_weight_record = &mut ctx.accounts.voter_weight_record; + let clock = Clock::get()?; + + // Check if the withdrawal is allowed by the spl-governance program + // ... + + // Check if the current slot is the same as the last deposit or withdrawal slot + if voter_weight_record.last_deposit_or_withdrawal_slot == clock.slot { + return Err(RegistrarError::InvalidOperation.into()); + } + + let registrar_config = RegistrarConfig::unpack_from_slice(&ctx.accounts.registrar_config.data.borrow())?; + + // Check if the token is an accepted governance token + let token_mint = ctx.accounts.governance_token_mint.key(); + if !registrar_config.accepted_tokens.contains(&token_mint) { + return Err(RegistrarError::InvalidArgument.into()); + } + + let weight_decrease = amount * registrar_config.weights[registrar_config.accepted_tokens.iter().position(|&token| token == token_mint).ok_or(RegistrarError::InvalidArgument)?]; + if voter_weight_record.weight < weight_decrease { + return Err(RegistrarError::InsufficientFunds.into()); + } + + // Transfer tokens from the governance program's account to the voter's account + transfer_tokens( + &ctx.accounts.governance_token_mint, // Replace with the governance program's token account + &ctx.accounts.voter_token_account, + amount, + &ctx.accounts.token_program, + &ctx.accounts.authority, + )?; + + // Update the VoterWeightRecord account + voter_weight_record.weight = voter_weight_record.weight.checked_sub(weight_decrease).ok_or(RegistrarError::Overflow)?; + voter_weight_record.last_deposit_or_withdrawal_slot = clock.slot; + voter_weight_record.serialize(&mut ctx.accounts.voter_weight_record.data.borrow_mut()[..])?; + + // Update the MaxVoterWeightRecord account + // ... + + Ok(()) +} + +fn transfer_tokens( + source_account: &AccountInfo, + destination_account: &AccountInfo, + amount: u64, + token_program: &AccountInfo, + authority: &AccountInfo, +) -> ProgramResult { + let transfer_instruction = transfer( + token_program.key, + source_account.key, + destination_account.key, + authority.key, + &[], + amount, + )?; + + invoke( + &transfer_instruction, + &[ + source_account.clone(), + destination_account.clone(), + authority.clone(), + token_program.clone(), + ], + )?; + + Ok(()) +} diff --git a/tests/deposit_withdraw_governance_token.rs b/tests/deposit_withdraw_governance_token.rs new file mode 100644 index 00000000..8afe0ef6 --- /dev/null +++ b/tests/deposit_withdraw_governance_token.rs @@ -0,0 +1,84 @@ +use solana_program_test::*; +use solana_sdk::{pubkey::Pubkey, signature::Signer}; + +#[tokio::test] +async fn test_deposit_governance_token() { + let program_id = Pubkey::new_unique(); + let registrar_config_account = Pubkey::new_unique(); + let voter_weight_record_account = Pubkey::new_unique(); + let voter_token_account = Pubkey::new_unique(); + let governance_token_mint = Pubkey::new_unique(); + let token_program = Pubkey::new_unique(); + let authority = Pubkey::new_unique(); + let payer_account = Pubkey::new_unique(); + + let mut test_context = ProgramTest::new( + "registrar", + program_id, + processor!(process_instruction), + ); + + test_context + .add_account_with_data( + registrar_config_account, + &RegistrarConfig { + accepted_tokens: vec![governance_token_mint], + weights: vec![10], + }, + RegistrarConfig::LEN, + &program_id, + ) + .await + .unwrap(); + + test_context + .add_account_with_data( + voter_weight_record_account, + &VoterWeightRecord { + voter: payer_account, + weight: 0, + last_deposit_or_withdrawal_slot: 0, + }, + VoterWeightRecord::LEN, + &program_id, + ) + .await + .unwrap(); + + let (mut banks_client, payer, recent_blockhash) = test_context + .start_with_new_bank() + .await + .unwrap(); + + let deposit_amount = 100; + + let deposit_governance_token_instruction = deposit_governance_token( + program_id, + voter_weight_record_account, + voter_token_account, + governance_token_mint, + registrar_config_account, + token_program, + authority, + deposit_amount, + ); + + banks_client + .process_transaction_with_preflight(deposit_governance_token_instruction, &[&payer]) + .await + .unwrap(); + + let voter_weight_record_account_data = banks_client + .get_account_data_with_borsh(voter_weight_record_account) + .await + .unwrap(); + + assert_eq!( + voter_weight_record_account_data, + VoterWeightRecord { + voter: payer_account, + weight: deposit_amount * 10, + last_deposit_or_withdrawal_slot: banks_client.get_latest_slot().await.unwrap(), + } + ); +} \ No newline at end of file diff --git a/tests/registrar_config.rs b/tests/registrar_config.rs new file mode 100644 index 00000000..debc85c2 --- /dev/null +++ b/tests/registrar_config.rs @@ -0,0 +1,59 @@ +use solana_program_test::*; +use solana_sdk::{pubkey::Pubkey, signature::Signer}; + +#[tokio::test] +async fn test_initialize_registrar() { + let program_id = Pubkey::new_unique(); + let registrar_config_account = Pubkey::new_unique(); + let payer_account = Pubkey::new_unique(); + + let mut test_context = ProgramTest::new( + "registrar", + program_id, + processor!(process_instruction), + ); + + test_context + .add_account_with_data( + registrar_config_account, + &[], + RegistrarConfig::LEN, + &program_id, + ) + .await + .unwrap(); + + let (mut banks_client, payer, recent_blockhash) = test_context + .start_with_new_bank() + .await + .unwrap(); + + let accepted_tokens = vec![Pubkey::new_unique(), Pubkey::new_unique()]; + let weights = vec![10, 20]; + + let initialize_registrar_instruction = initialize_registrar( + program_id, + registrar_config_account, + payer_account, + accepted_tokens.clone(), + weights.clone(), + ); + + banks_client + .process_transaction_with_preflight(initialize_registrar_instruction, &[&payer]) + .await + .unwrap(); + + let registrar_config_account_data = banks_client + .get_account_data_with_borsh(registrar_config_account) + .await + .unwrap(); + + assert_eq!( + registrar_config_account_data, + RegistrarConfig { + accepted_tokens, + weights + } + ); +} diff --git a/tests/withdraw_governance_token.rs b/tests/withdraw_governance_token.rs new file mode 100644 index 00000000..d7e11ee9 --- /dev/null +++ b/tests/withdraw_governance_token.rs @@ -0,0 +1,393 @@ +#[tokio::test] +async fn test_withdraw_governance_token() { + let program_id = Pubkey::new_unique(); + let registrar_config_account = Pubkey::new_unique(); + let voter_weight_record_account = Pubkey::new_unique(); + let voter_token_account = Pubkey::new_unique(); + let governance_token_mint = Pubkey::new_unique(); + let token_program = Pubkey::new_unique(); + let spl_governance_program = Pubkey::new_unique(); + let authority = Pubkey::new_unique(); + let payer_account = Pubkey::new_unique(); + + let mut test_context = ProgramTest::new( + "registrar", + program_id, + processor!(process_instruction), + ); + + test_context + .add_account_with_data( + registrar_config_account, + &RegistrarConfig { + accepted_tokens: vec![governance_token_mint], + weights: vec![10], + }, + RegistrarConfig::LEN, + &program_id, + ) + .await + .unwrap(); + + test_context + .add_account_with_data( + voter_weight_record_account, + &VoterWeightRecord { + voter: payer_account, + weight: 1000, + last_deposit_or_withdrawal_slot: 0, + }, + VoterWeightRecord::LEN, + &program_id, + ) + .await + .unwrap(); + + let (mut banks_client, payer, recent_blockhash) = test_context + .start_with_new_bank() + .await + .unwrap(); + + let withdraw_amount = 100; + + let withdraw_governance_token_instruction = withdraw_governance_token( + program_id, + voter_weight_record_account, + voter_token_account, + governance_token_mint, + registrar_config_account, + token_program, + spl_governance_program, + authority, + withdraw_amount, + ); + + banks_client + .process_transaction_with_preflight(withdraw_governance_token_instruction, &[&payer]) + .await + .unwrap(); + + let voter_weight_record_account_data = banks_client + .get_account_data_with_borsh(voter_weight_record_account) + .await + .unwrap(); + + assert_eq!( + voter_weight_record_account_data, + VoterWeightRecord { + voter: payer_account, + weight: 900, + last_deposit_or_withdrawal_slot: banks_client.get_latest_slot().await.unwrap(), + } + ); +} + +#[tokio::test] +async fn test_deposit_withdraw_different_governance_tokens() { + let program_id = Pubkey::new_unique(); + let registrar_config_account = Pubkey::new_unique(); + let voter_weight_record_account = Pubkey::new_unique(); + let voter_token_account_1 = Pubkey::new_unique(); + let voter_token_account_2 = Pubkey::new_unique(); + let governance_token_mint_1 = Pubkey::new_unique(); + let governance_token_mint_2 = Pubkey::new_unique(); + let token_program = Pubkey::new_unique(); + let authority = Pubkey::new_unique(); + let payer_account = Pubkey::new_unique(); + + let mut test_context = ProgramTest::new( + "registrar", + program_id, + processor!(process_instruction), + ); + + test_context + .add_account_with_data( + registrar_config_account, + &RegistrarConfig { + accepted_tokens: vec![governance_token_mint_1, governance_token_mint_2], + weights: vec![10, 20], + }, + RegistrarConfig::LEN, + &program_id, + ) + .await + .unwrap(); + + test_context + .add_account_with_data( + voter_weight_record_account, + &VoterWeightRecord { + voter: payer_account, + weight: 0, + last_deposit_or_withdrawal_slot: 0, + }, + VoterWeightRecord::LEN, + &program_id, + ) + .await + .unwrap(); + + let (mut banks_client, payer, recent_blockhash) = test_context + .start_with_new_bank() + .await + .unwrap(); + + let deposit_amount_1 = 100; + let deposit_amount_2 = 200; + + let deposit_governance_token_instruction_1 = deposit_governance_token( + program_id, + voter_weight_record_account, + voter_token_account_1, + governance_token_mint_1, + registrar_config_account, + token_program, + authority, + deposit_amount_1, + ); + + banks_client + .process_transaction_with_preflight(deposit_governance_token_instruction_1, &[&payer]) + .await + .unwrap(); + + let deposit_governance_token_instruction_2 = deposit_governance_token( + program_id, + voter_weight_record_account, + voter_token_account_2, + governance_token_mint_2, + registrar_config_account, + token_program, + authority, + deposit_amount_2, + ); + + banks_client + .process_transaction_with_preflight(deposit_governance_token_instruction_2, &[&payer]) + .await + .unwrap(); + + let voter_weight_record_account_data = banks_client + .get_account_data_with_borsh(voter_weight_record_account) + .await + .unwrap(); + + assert_eq!( + voter_weight_record_account_data, + VoterWeightRecord { + voter: payer_account, + weight: deposit_amount_1 * 10 + deposit_amount_2 * 20, + last_deposit_or_withdrawal_slot: banks_client.get_latest_slot().await.unwrap(), + } + ); + + let withdraw_amount_1 = 50; + let withdraw_amount_2 = 100; + + let withdraw_governance_token_instruction_1 = withdraw_governance_token( + program_id, + voter_weight_record_account, + voter_token_account_1, + governance_token_mint_1, + registrar_config_account, + token_program, + spl_governance_program, + authority, + withdraw_amount_1, + ); + + banks_client + .process_transaction_with_preflight(withdraw_governance_token_instruction_1, &[&payer]) + .await + .unwrap(); + + let withdraw_governance_token_instruction_2 = withdraw_governance_token( + program_id, + voter_weight_record_account, + voter_token_account_2, + governance_token_mint_2, + registrar_config_account, + token_program, + spl_governance_program, + authority, + withdraw_amount_2, + ); + + banks_client + .process_transaction_with_preflight(withdraw_governance_token_instruction_2, &[&payer]) + .await + .unwrap(); + + let voter_weight_record_account_data = banks_client + .get_account_data_with_borsh(voter_weight_record_account) + .await + .unwrap(); + + assert_eq!( + voter_weight_record_account_data, + VoterWeightRecord { + voter: payer_account, + weight: (deposit_amount_1 - withdraw_amount_1) * 10 + + (deposit_amount_2 - withdraw_amount_2) * 20, + last_deposit_or_withdrawal_slot: banks_client.get_latest_slot().await.unwrap(), + } + ); +} + +#[tokio::test] +async fn test_invalid_deposit_withdraw_amount() { + let program_id = Pubkey::new_unique(); + let registrar_config_account = Pubkey::new_unique(); + let voter_weight_record_account = Pubkey::new_unique(); + let voter_token_account = Pubkey::new_unique(); + let governance_token_mint = Pubkey::new_unique(); + let token_program = Pubkey::new_unique(); + let authority = Pubkey::new_unique(); + let payer_account = Pubkey::new_unique(); + + let mut test_context = ProgramTest::new( + "registrar", + program_id, + processor!(process_instruction), + ); + + test_context + .add_account_with_data( + registrar_config_account, + &RegistrarConfig { + accepted_tokens: vec![governance_token_mint], + weights: vec![10], + }, + RegistrarConfig::LEN, + &program_id, + ) + .await + .unwrap(); + + test_context + .add_account_with_data( + voter_weight_record_account, + &VoterWeightRecord { + voter: payer_account, + weight: 0, + last_deposit_or_withdrawal_slot: 0, + }, + VoterWeightRecord::LEN, + &program_id, + ) + .await + .unwrap(); + + let (mut banks_client, payer, recent_blockhash) = test_context + .start_with_new_bank() + .await + .unwrap(); + + let invalid_deposit_amount = 0; + + let deposit_governance_token_instruction = deposit_governance_token( + program_id, + voter_weight_record_account, + voter_token_account, + governance_token_mint, + registrar_config_account, + token_program, + authority, + invalid_deposit_amount, + ); + + assert!(banks_client + .process_transaction_with_preflight(deposit_governance_token_instruction, &[&payer]) + .await + .is_err()); + + let invalid_withdraw_amount = 1000; + + let withdraw_governance_token_instruction = withdraw_governance_token( + program_id, + voter_weight_record_account, + voter_token_account, + governance_token_mint, + registrar_config_account, + token_program, + spl_governance_program, + authority, + invalid_withdraw_amount, + ); + + assert!(banks_client + .process_transaction_with_preflight(withdraw_governance_token_instruction, &[&payer]) + .await + .is_err()); +} + +#[tokio::test] +async fn test_insufficient_balance_for_withdraw() { + let program_id = Pubkey::new_unique(); + let registrar_config_account = Pubkey::new_unique(); + let voter_weight_record_account = Pubkey::new_unique(); + let voter_token_account = Pubkey::new_unique(); + let governance_token_mint = Pubkey::new_unique(); + let token_program = Pubkey::new_unique(); + let spl_governance_program = Pubkey::new_unique(); + let authority = Pubkey::new_unique(); + let payer_account = Pubkey::new_unique(); + + let mut test_context = ProgramTest::new( + "registrar", + program_id, + processor!(process_instruction), + ); + + test_context + .add_account_with_data( + registrar_config_account, + &RegistrarConfig { + accepted_tokens: vec![governance_token_mint], + weights: vec![10], + }, + RegistrarConfig::LEN, + &program_id, + ) + .await + .unwrap(); + + test_context + .add_account_with_data( + voter_weight_record_account, + &VoterWeightRecord { + voter: payer_account, + weight: 100, + last_deposit_or_withdrawal_slot: 0, + }, + VoterWeightRecord::LEN, + &program_id, + ) + .await + .unwrap(); + + let (mut banks_client, payer, recent_blockhash) = test_context + .start_with_new_bank() + .await + .unwrap(); + + let withdraw_amount = 1000; + + let withdraw_governance_token_instruction = withdraw_governance_token( + program_id, + voter_weight_record_account, + voter_token_account, + governance_token_mint, + registrar_config_account, + token_program, + spl_governance_program, + authority, + withdraw_amount, + ); + + assert!(banks_client + .process_transaction_with_preflight(withdraw_governance_token_instruction, &[&payer]) + .await + .is_err()); +}