From a1312bb67800f3d950e3d6d60b81f05e06782bf5 Mon Sep 17 00:00:00 2001 From: rwwwx Date: Tue, 16 Jul 2024 16:30:33 +0300 Subject: [PATCH] added trait extension for `Result` to check specific error-codes from on-chain calls + fmt --- program-states/src/state/deposit_entry.rs | 12 ++++-- .../tests/program_test/governance.rs | 6 +-- .../tests/program_test/mod.rs | 2 +- .../tests/program_test/utils.rs | 38 +++++++++++++++++-- .../tests/test_deposit_constant.rs | 13 +++++-- .../tests/test_deposit_no_locking.rs | 14 +++++-- .../tests/test_extend_deposit.rs | 7 +++- .../voter-stake-registry/tests/test_lockup.rs | 13 +++++-- 8 files changed, 79 insertions(+), 26 deletions(-) diff --git a/program-states/src/state/deposit_entry.rs b/program-states/src/state/deposit_entry.rs index 6dc650e7..8621b898 100644 --- a/program-states/src/state/deposit_entry.rs +++ b/program-states/src/state/deposit_entry.rs @@ -1,6 +1,7 @@ -use crate::state::lockup::{Lockup, LockupKind}; - -use crate::state::LockupPeriod; +use crate::state::{ + lockup::{Lockup, LockupKind}, + LockupPeriod, +}; use anchor_lang::prelude::*; /// Bookkeeping for a single deposit for a given mint and lockup schedule. @@ -84,7 +85,10 @@ impl DepositEntry { #[cfg(test)] mod tests { use super::*; - use crate::state::{LockupKind::Constant, LockupKind::None, LockupPeriod}; + use crate::state::{ + LockupKind::{Constant, None}, + LockupPeriod, + }; #[test] pub fn far_future_lockup_start_test() -> Result<()> { diff --git a/programs/voter-stake-registry/tests/program_test/governance.rs b/programs/voter-stake-registry/tests/program_test/governance.rs index c4321c75..f1bf64e9 100644 --- a/programs/voter-stake-registry/tests/program_test/governance.rs +++ b/programs/voter-stake-registry/tests/program_test/governance.rs @@ -4,9 +4,9 @@ use solana_sdk::{ pubkey::Pubkey, signature::{Keypair, Signer}, }; -use spl_governance::state::realm::GoverningTokenConfigAccountArgs; -use spl_governance::state::realm_config::GoverningTokenType; -use spl_governance::state::{proposal, vote_record}; +use spl_governance::state::{ + proposal, realm::GoverningTokenConfigAccountArgs, realm_config::GoverningTokenType, vote_record, +}; use std::rc::Rc; #[derive(Clone)] diff --git a/programs/voter-stake-registry/tests/program_test/mod.rs b/programs/voter-stake-registry/tests/program_test/mod.rs index 2d2c8c9c..7284937a 100644 --- a/programs/voter-stake-registry/tests/program_test/mod.rs +++ b/programs/voter-stake-registry/tests/program_test/mod.rs @@ -17,7 +17,7 @@ use std::{ str::FromStr, sync::{Arc, RwLock}, }; -pub use utils::*; +pub use utils::{assert_custom_on_chain_error::AssertCustomOnChainErr, *}; pub mod addin; pub mod cookies; diff --git a/programs/voter-stake-registry/tests/program_test/utils.rs b/programs/voter-stake-registry/tests/program_test/utils.rs index f08ecd9f..33f66bde 100644 --- a/programs/voter-stake-registry/tests/program_test/utils.rs +++ b/programs/voter-stake-registry/tests/program_test/utils.rs @@ -1,11 +1,16 @@ use bytemuck::{bytes_of, Contiguous}; -use solana_program::program_error::ProgramError; +use mplx_staking_states::error::VsrError; +use solana_program::{instruction::InstructionError, program_error::ProgramError}; use solana_program_test::{BanksClientError, ProgramTestContext}; use solana_sdk::{ - program_pack::Pack, pubkey::Pubkey, signature::Keypair, signer::Signer, system_instruction, - transaction::Transaction, + program_pack::Pack, + pubkey::Pubkey, + signature::Keypair, + signer::Signer, + system_instruction, + transaction::{Transaction, TransactionError}, }; -use std::borrow::BorrowMut; +use std::{borrow::BorrowMut, fmt::Debug}; #[allow(dead_code)] pub fn gen_signer_seeds<'a>(nonce: &'a u64, acc_pk: &'a Pubkey) -> [&'a [u8]; 2] { @@ -102,3 +107,28 @@ pub async fn advance_clock_by_ts(context: &mut ProgramTestContext, ts: i64) { new_clock.unix_timestamp += ts; context.borrow_mut().set_sysvar(&new_clock); } + +pub mod assert_custom_on_chain_error { + use super::*; + use mplx_staking_states::error::VsrError; + use std::fmt::Debug; + + pub trait AssertCustomOnChainErr { + fn assert_on_chain_err(self, expected_err: VsrError); + } + + impl AssertCustomOnChainErr for Result { + fn assert_on_chain_err(self, expected_err: VsrError) { + assert!(self.is_err()); + match self.unwrap_err() { + BanksClientError::TransactionError(TransactionError::InstructionError( + _, + InstructionError::Custom(code), + )) => { + debug_assert_eq!((expected_err as u32) + 6000, code); + } + _ => unreachable!("BanksClientError has no 'Custom' variant."), + } + } + } +} diff --git a/programs/voter-stake-registry/tests/test_deposit_constant.rs b/programs/voter-stake-registry/tests/test_deposit_constant.rs index 8865140d..76528bfb 100644 --- a/programs/voter-stake-registry/tests/test_deposit_constant.rs +++ b/programs/voter-stake-registry/tests/test_deposit_constant.rs @@ -1,5 +1,8 @@ use anchor_spl::token::TokenAccount; -use mplx_staking_states::state::{LockupKind, LockupPeriod}; +use mplx_staking_states::{ + error::VsrError, + state::{LockupKind, LockupPeriod}, +}; use program_test::*; use solana_program_test::*; use solana_sdk::{pubkey::Pubkey, signature::Keypair, signer::Signer, transport::TransportError}; @@ -194,7 +197,9 @@ async fn test_deposit_constant() -> Result<(), TransportError> { assert_eq!(after_deposit.voter_weight, after_deposit.vault); // unchanged assert_eq!(after_deposit.vault, 10_000); assert_eq!(after_deposit.deposit, 10_000); - withdraw(1, 1).await.expect_err("all locked up"); + withdraw(1, 1) + .await + .assert_on_chain_err(VsrError::UnlockMustBeCalledFirst); // advance to day 95. Just to be sure withdraw isn't possible without unlocking first // even at lockup period + cooldown period (90 + 5 respectively in that case) @@ -218,7 +223,7 @@ async fn test_deposit_constant() -> Result<(), TransportError> { .unwrap(); withdraw(10_000, 1) .await - .expect_err("Cooldown still not passed"); + .assert_on_chain_err(VsrError::InvalidTimestampArguments); context.solana.advance_clock_by_slots(2).await; // avoid caching of transactions // warp to day 100. (90 days of lockup + fake cooldown (5 days)) + 5 days of true cooldown @@ -370,7 +375,7 @@ async fn test_withdrawing_without_unlocking() -> Result<(), TransportError> { // withdraw withdraw(10_000) .await - .expect_err("impossible to withdraw without unlocking"); + .assert_on_chain_err(VsrError::InsufficientUnlockedTokens); Ok(()) } diff --git a/programs/voter-stake-registry/tests/test_deposit_no_locking.rs b/programs/voter-stake-registry/tests/test_deposit_no_locking.rs index 79791e8b..e1b383ca 100644 --- a/programs/voter-stake-registry/tests/test_deposit_no_locking.rs +++ b/programs/voter-stake-registry/tests/test_deposit_no_locking.rs @@ -1,5 +1,9 @@ +use crate::program_test::utils::assert_custom_on_chain_error::AssertCustomOnChainErr; use anchor_spl::token::TokenAccount; -use mplx_staking_states::state::{LockupKind, LockupPeriod}; +use mplx_staking_states::{ + error::VsrError, + state::{LockupKind, LockupPeriod}, +}; use program_test::*; use solana_program_test::*; use solana_sdk::{pubkey::Pubkey, signature::Keypair, signer::Signer, transport::TransportError}; @@ -222,7 +226,9 @@ async fn test_deposit_no_locking() -> Result<(), TransportError> { assert_eq!(after_withdraw1.vault, 12000); assert_eq!(after_withdraw1.deposit, 5000); - withdraw(5001).await.expect_err("withdrew too much"); + withdraw(5001) + .await + .assert_on_chain_err(VsrError::InsufficientUnlockedTokens); withdraw(5000).await.unwrap(); @@ -236,11 +242,11 @@ async fn test_deposit_no_locking() -> Result<(), TransportError> { addin .close_deposit_entry(&voter, voter_authority, 2) .await - .expect_err("deposit not in use"); + .assert_on_chain_err(VsrError::UnusedDepositEntryIndex); addin .close_deposit_entry(&voter, voter_authority, 1) .await - .expect_err("deposit not empty"); + .assert_on_chain_err(VsrError::VotingTokenNonZero); addin .close_deposit_entry(&voter, voter_authority, 0) .await diff --git a/programs/voter-stake-registry/tests/test_extend_deposit.rs b/programs/voter-stake-registry/tests/test_extend_deposit.rs index 50637ca1..e02ac3b8 100644 --- a/programs/voter-stake-registry/tests/test_extend_deposit.rs +++ b/programs/voter-stake-registry/tests/test_extend_deposit.rs @@ -1,5 +1,8 @@ use anchor_spl::token::TokenAccount; -use mplx_staking_states::state::{LockupKind, LockupPeriod}; +use mplx_staking_states::{ + error::VsrError, + state::{LockupKind, LockupPeriod}, +}; use program_test::*; use solana_program_test::*; use solana_sdk::{ @@ -808,7 +811,7 @@ async fn extend_from_three_month_to_one_year() -> Result<(), TransportError> { 50, ) .await - .expect_err("Impossible to extend stake from existing stake (not flex) to another period"); + .assert_on_chain_err(VsrError::ArithmeticOverflow); Ok(()) } diff --git a/programs/voter-stake-registry/tests/test_lockup.rs b/programs/voter-stake-registry/tests/test_lockup.rs index 9d10d7a2..ade538e4 100644 --- a/programs/voter-stake-registry/tests/test_lockup.rs +++ b/programs/voter-stake-registry/tests/test_lockup.rs @@ -1,8 +1,12 @@ use anchor_spl::token::TokenAccount; -use mplx_staking_states::state::{LockupKind, LockupPeriod}; +use mplx_staking_states::{ + error::VsrError, + state::{LockupKind, LockupPeriod}, +}; use program_test::*; use solana_program_test::*; use solana_sdk::{signature::Keypair, signer::Signer, transport::TransportError}; +use AssertCustomOnChainErr; mod program_test; @@ -153,7 +157,8 @@ async fn test_unlock_and_withdraw_before_end_ts() -> Result<(), TransportError> &context.rewards.program_id, ) .await - .expect_err("fails because it's too early to unlock is invalid"); + .assert_on_chain_err(VsrError::DepositStillLocked); + context .addin .withdraw( @@ -166,7 +171,7 @@ async fn test_unlock_and_withdraw_before_end_ts() -> Result<(), TransportError> 10000, ) .await - .expect_err("fails because it's impossible to withdraw without unlock"); + .assert_on_chain_err(VsrError::UnlockMustBeCalledFirst); Ok(()) } @@ -339,7 +344,7 @@ async fn test_unlock_after_end_ts() -> Result<(), TransportError> { 10000, ) .await - .expect_err("fails because cooldown is ongoing"); + .assert_on_chain_err(VsrError::InvalidTimestampArguments); Ok(()) }