From 8eaaafed54d3fd4978e5d2fa0f2ee7f07dc9a0a7 Mon Sep 17 00:00:00 2001 From: Tyera Eulberg Date: Tue, 25 Jan 2022 13:45:53 -0700 Subject: [PATCH] Refactor TestContext to enable unwrap_err (#2805) --- .../program-2022/tests/initialize_account.rs | 19 ++--- token/program-2022/tests/initialize_mint.rs | 56 ++++++++------- .../tests/mint_close_authority.rs | 52 ++++++++------ token/program-2022/tests/program_test.rs | 35 +++++++--- token/program-2022/tests/transfer_fee.rs | 70 +++++++++++-------- 5 files changed, 137 insertions(+), 95 deletions(-) diff --git a/token/program-2022/tests/initialize_account.rs b/token/program-2022/tests/initialize_account.rs index 2c958f35..1e96a62f 100644 --- a/token/program-2022/tests/initialize_account.rs +++ b/token/program-2022/tests/initialize_account.rs @@ -1,7 +1,9 @@ #![cfg(feature = "test-bpf")] +mod program_test; use { - solana_program_test::{processor, tokio, ProgramTest}, + program_test::TestContext, + solana_program_test::tokio, solana_sdk::{ instruction::InstructionError, program_pack::Pack, @@ -17,16 +19,15 @@ use { transfer_fee::{self, TransferFeeAmount}, ExtensionType, StateWithExtensions, }, - id, instruction, - processor::Processor, + instruction, state::{Account, Mint}, }, }; #[tokio::test] async fn no_extensions() { - let program_test = ProgramTest::new("spl_token_2022", id(), processor!(Processor::process)); - let mut ctx = program_test.start_with_context().await; + let context = TestContext::new().await; + let mut ctx = context.context.lock().await; let rent = ctx.banks_client.get_rent().await.unwrap(); let mint_account = Keypair::new(); let mint_authority_pubkey = Pubkey::new_unique(); @@ -96,8 +97,8 @@ async fn no_extensions() { #[tokio::test] async fn fail_on_invalid_mint() { - let program_test = ProgramTest::new("spl_token_2022", id(), processor!(Processor::process)); - let mut ctx = program_test.start_with_context().await; + let context = TestContext::new().await; + let mut ctx = context.context.lock().await; let rent = ctx.banks_client.get_rent().await.unwrap(); let mint_account = Keypair::new(); @@ -161,8 +162,8 @@ async fn fail_on_invalid_mint() { #[tokio::test] async fn single_extension() { - let program_test = ProgramTest::new("spl_token_2022", id(), processor!(Processor::process)); - let mut ctx = program_test.start_with_context().await; + let context = TestContext::new().await; + let mut ctx = context.context.lock().await; let rent = ctx.banks_client.get_rent().await.unwrap(); let mint_account = Keypair::new(); let mint_authority_pubkey = Pubkey::new_unique(); diff --git a/token/program-2022/tests/initialize_mint.rs b/token/program-2022/tests/initialize_mint.rs index 344e3222..63fd66db 100644 --- a/token/program-2022/tests/initialize_mint.rs +++ b/token/program-2022/tests/initialize_mint.rs @@ -2,8 +2,8 @@ mod program_test; use { - program_test::TestContext, - solana_program_test::{processor, tokio, ProgramTest}, + program_test::{TestContext, TokenContext}, + solana_program_test::tokio, solana_sdk::{ instruction::InstructionError, program_option::COption, @@ -17,8 +17,7 @@ use { spl_token_2022::{ error::TokenError, extension::{mint_close_authority::MintCloseAuthority, transfer_fee, ExtensionType}, - id, instruction, - processor::Processor, + instruction, state::Mint, }, spl_token_client::token::ExtensionInitializationParams, @@ -27,12 +26,14 @@ use { #[tokio::test] async fn success_base() { - let TestContext { + let mut context = TestContext::new().await; + context.init_token_with_mint(vec![]).await.unwrap(); + let TokenContext { decimals, mint_authority, token, .. - } = TestContext::new(vec![]).await.unwrap(); + } = context.token_context.unwrap(); let mint = token.get_mint_info().await.unwrap(); assert_eq!(mint.base.decimals, decimals); @@ -47,8 +48,8 @@ async fn success_base() { #[tokio::test] async fn fail_extension_no_space() { - let program_test = ProgramTest::new("spl_token_2022", id(), processor!(Processor::process)); - let mut ctx = program_test.start_with_context().await; + let context = TestContext::new().await; + let mut ctx = context.context.lock().await; let rent = ctx.banks_client.get_rent().await.unwrap(); let mint_account = Keypair::new(); let mint_authority_pubkey = Pubkey::new_unique(); @@ -100,8 +101,8 @@ async fn fail_extension_no_space() { #[tokio::test] async fn fail_extension_after_mint_init() { - let program_test = ProgramTest::new("spl_token_2022", id(), processor!(Processor::process)); - let mut ctx = program_test.start_with_context().await; + let context = TestContext::new().await; + let mut ctx = context.context.lock().await; let rent = ctx.banks_client.get_rent().await.unwrap(); let mint_account = Keypair::new(); let mint_authority_pubkey = Pubkey::new_unique(); @@ -154,16 +155,19 @@ async fn fail_extension_after_mint_init() { #[tokio::test] async fn success_extension_and_base() { let close_authority = Some(Pubkey::new_unique()); - let TestContext { + let mut context = TestContext::new().await; + context + .init_token_with_mint(vec![ExtensionInitializationParams::MintCloseAuthority { + close_authority, + }]) + .await + .unwrap(); + let TokenContext { decimals, mint_authority, token, .. - } = TestContext::new(vec![ExtensionInitializationParams::MintCloseAuthority { - close_authority, - }]) - .await - .unwrap(); + } = context.token_context.unwrap(); let state = token.get_mint_info().await.unwrap(); assert_eq!(state.base.decimals, decimals); @@ -183,8 +187,8 @@ async fn success_extension_and_base() { #[tokio::test] async fn fail_init_overallocated_mint() { - let program_test = ProgramTest::new("spl_token_2022", id(), processor!(Processor::process)); - let mut ctx = program_test.start_with_context().await; + let context = TestContext::new().await; + let mut ctx = context.context.lock().await; let rent = ctx.banks_client.get_rent().await.unwrap(); let mint_account = Keypair::new(); let mint_authority_pubkey = Pubkey::new_unique(); @@ -230,8 +234,8 @@ async fn fail_init_overallocated_mint() { #[tokio::test] async fn fail_account_init_after_mint_extension() { - let program_test = ProgramTest::new("spl_token_2022", id(), processor!(Processor::process)); - let mut ctx = program_test.start_with_context().await; + let context = TestContext::new().await; + let mut ctx = context.context.lock().await; let rent = ctx.banks_client.get_rent().await.unwrap(); let mint_account = Keypair::new(); let mint_authority_pubkey = Pubkey::new_unique(); @@ -303,8 +307,8 @@ async fn fail_account_init_after_mint_extension() { #[tokio::test] async fn fail_account_init_after_mint_init() { - let program_test = ProgramTest::new("spl_token_2022", id(), processor!(Processor::process)); - let mut ctx = program_test.start_with_context().await; + let context = TestContext::new().await; + let mut ctx = context.context.lock().await; let rent = ctx.banks_client.get_rent().await.unwrap(); let mint_account = Keypair::new(); let mint_authority_pubkey = Pubkey::new_unique(); @@ -357,8 +361,8 @@ async fn fail_account_init_after_mint_init() { #[tokio::test] async fn fail_account_init_after_mint_init_with_extension() { - let program_test = ProgramTest::new("spl_token_2022", id(), processor!(Processor::process)); - let mut ctx = program_test.start_with_context().await; + let context = TestContext::new().await; + let mut ctx = context.context.lock().await; let rent = ctx.banks_client.get_rent().await.unwrap(); let mint_account = Keypair::new(); let mint_authority_pubkey = Pubkey::new_unique(); @@ -417,8 +421,8 @@ async fn fail_account_init_after_mint_init_with_extension() { #[tokio::test] async fn fail_fee_init_after_mint_init() { - let program_test = ProgramTest::new("spl_token_2022", id(), processor!(Processor::process)); - let mut ctx = program_test.start_with_context().await; + let context = TestContext::new().await; + let mut ctx = context.context.lock().await; let rent = ctx.banks_client.get_rent().await.unwrap(); let mint_account = Keypair::new(); let mint_authority_pubkey = Pubkey::new_unique(); diff --git a/token/program-2022/tests/mint_close_authority.rs b/token/program-2022/tests/mint_close_authority.rs index fb937a2a..08a986ab 100644 --- a/token/program-2022/tests/mint_close_authority.rs +++ b/token/program-2022/tests/mint_close_authority.rs @@ -2,7 +2,7 @@ mod program_test; use { - program_test::TestContext, + program_test::{TestContext, TokenContext}, solana_program_test::tokio, solana_sdk::{ instruction::InstructionError, program_option::COption, pubkey::Pubkey, signature::Signer, @@ -18,16 +18,19 @@ use { #[tokio::test] async fn success_init() { let close_authority = Some(Pubkey::new_unique()); - let TestContext { + let mut context = TestContext::new().await; + context + .init_token_with_mint(vec![ExtensionInitializationParams::MintCloseAuthority { + close_authority, + }]) + .await + .unwrap(); + let TokenContext { decimals, mint_authority, token, .. - } = TestContext::new(vec![ExtensionInitializationParams::MintCloseAuthority { - close_authority, - }]) - .await - .unwrap(); + } = context.token_context.unwrap(); let state = token.get_mint_info().await.unwrap(); assert_eq!(state.base.decimals, decimals); @@ -48,12 +51,14 @@ async fn success_init() { #[tokio::test] async fn set_authority() { let close_authority = Keypair::new(); - let TestContext { token, .. } = - TestContext::new(vec![ExtensionInitializationParams::MintCloseAuthority { + let mut context = TestContext::new().await; + context + .init_token_with_mint(vec![ExtensionInitializationParams::MintCloseAuthority { close_authority: Some(close_authority.pubkey()), }]) .await .unwrap(); + let token = context.token_context.unwrap().token; let new_authority = Keypair::new(); // fail, wrong signature @@ -148,12 +153,14 @@ async fn set_authority() { #[tokio::test] async fn success_close() { let close_authority = Keypair::new(); - let TestContext { token, .. } = - TestContext::new(vec![ExtensionInitializationParams::MintCloseAuthority { + let mut context = TestContext::new().await; + context + .init_token_with_mint(vec![ExtensionInitializationParams::MintCloseAuthority { close_authority: Some(close_authority.pubkey()), }]) .await .unwrap(); + let token = context.token_context.unwrap().token; let destination = Pubkey::new_unique(); token @@ -167,11 +174,13 @@ async fn success_close() { #[tokio::test] async fn fail_without_extension() { let close_authority = Pubkey::new_unique(); - let TestContext { + let mut context = TestContext::new().await; + context.init_token_with_mint(vec![]).await.unwrap(); + let TokenContext { mint_authority, token, .. - } = TestContext::new(vec![]).await.unwrap(); + } = context.token_context.unwrap(); // fail set let err = token @@ -207,15 +216,18 @@ async fn fail_without_extension() { #[tokio::test] async fn fail_close_with_supply() { let close_authority = Keypair::new(); - let TestContext { - token, + let mut context = TestContext::new().await; + context + .init_token_with_mint(vec![ExtensionInitializationParams::MintCloseAuthority { + close_authority: Some(close_authority.pubkey()), + }]) + .await + .unwrap(); + let TokenContext { mint_authority, + token, .. - } = TestContext::new(vec![ExtensionInitializationParams::MintCloseAuthority { - close_authority: Some(close_authority.pubkey()), - }]) - .await - .unwrap(); + } = context.token_context.unwrap(); // mint a token let owner = Pubkey::new_unique(); diff --git a/token/program-2022/tests/program_test.rs b/token/program-2022/tests/program_test.rs index dfec1e9c..529717c5 100644 --- a/token/program-2022/tests/program_test.rs +++ b/token/program-2022/tests/program_test.rs @@ -1,3 +1,5 @@ +#![allow(dead_code)] + use { solana_program_test::{processor, tokio::sync::Mutex, ProgramTest, ProgramTestContext}, solana_sdk::signer::{keypair::Keypair, Signer}, @@ -9,28 +11,39 @@ use { std::sync::Arc, }; -pub struct TestContext { +pub struct TokenContext { pub decimals: u8, pub mint_authority: Keypair, pub token: Token, pub alice: Keypair, pub bob: Keypair, - pub context: Arc>, // ProgramTestContext needs to #[derive(Debug)] +} + +pub struct TestContext { + pub context: Arc>, + pub token_context: Option, } impl TestContext { - pub async fn new( - extension_init_params: Vec, - ) -> TokenResult { + pub async fn new() -> Self { let program_test = ProgramTest::new("spl_token_2022", id(), processor!(Processor::process)); let context = program_test.start_with_context().await; let context = Arc::new(Mutex::new(context)); - let payer = keypair_clone(&context.lock().await.payer); + Self { + context, + token_context: None, + } + } + pub async fn init_token_with_mint( + &mut self, + extension_init_params: Vec, + ) -> TokenResult<()> { + let payer = keypair_clone(&self.context.lock().await.payer); let client: Arc> = Arc::new(ProgramBanksClient::new_from_context( - Arc::clone(&context), + Arc::clone(&self.context), ProgramBanksClientProcessTransaction, )); @@ -51,15 +64,15 @@ impl TestContext { extension_init_params, ) .await?; - - Ok(Self { + self.token_context = Some(TokenContext { decimals, mint_authority, token, alice: Keypair::new(), bob: Keypair::new(), - context, - }) + }); + + Ok(()) } } diff --git a/token/program-2022/tests/transfer_fee.rs b/token/program-2022/tests/transfer_fee.rs index faa60916..b69531c8 100644 --- a/token/program-2022/tests/transfer_fee.rs +++ b/token/program-2022/tests/transfer_fee.rs @@ -2,7 +2,7 @@ mod program_test; use { - program_test::TestContext, + program_test::{TestContext, TokenContext}, solana_program_test::tokio, solana_sdk::{ instruction::InstructionError, program_option::COption, pubkey::Pubkey, signature::Signer, @@ -72,19 +72,22 @@ async fn success_init() { newer_transfer_fee, .. } = test_transfer_fee_config(); - let TestContext { + let mut context = TestContext::new().await; + context + .init_token_with_mint(vec![ExtensionInitializationParams::TransferFeeConfig { + transfer_fee_config_authority: transfer_fee_config_authority.into(), + withdraw_withheld_authority: withdraw_withheld_authority.into(), + transfer_fee_basis_points: newer_transfer_fee.transfer_fee_basis_points.into(), + maximum_fee: newer_transfer_fee.maximum_fee.into(), + }]) + .await + .unwrap(); + let TokenContext { decimals, mint_authority, token, .. - } = TestContext::new(vec![ExtensionInitializationParams::TransferFeeConfig { - transfer_fee_config_authority: transfer_fee_config_authority.into(), - withdraw_withheld_authority: withdraw_withheld_authority.into(), - transfer_fee_basis_points: newer_transfer_fee.transfer_fee_basis_points.into(), - maximum_fee: newer_transfer_fee.maximum_fee.into(), - }]) - .await - .unwrap(); + } = context.token_context.unwrap(); let state = token.get_mint_info().await.unwrap(); assert_eq!(state.base.decimals, decimals); @@ -115,15 +118,16 @@ async fn fail_init_default_pubkey_as_authority() { newer_transfer_fee, .. } = test_transfer_fee_config(); - let err = TestContext::new(vec![ExtensionInitializationParams::TransferFeeConfig { - transfer_fee_config_authority: transfer_fee_config_authority.into(), - withdraw_withheld_authority: Some(Pubkey::default()), - transfer_fee_basis_points: newer_transfer_fee.transfer_fee_basis_points.into(), - maximum_fee: newer_transfer_fee.maximum_fee.into(), - }]) - .await - .err() - .unwrap(); + let mut context = TestContext::new().await; + let err = context + .init_token_with_mint(vec![ExtensionInitializationParams::TransferFeeConfig { + transfer_fee_config_authority: transfer_fee_config_authority.into(), + withdraw_withheld_authority: Some(Pubkey::default()), + transfer_fee_basis_points: newer_transfer_fee.transfer_fee_basis_points.into(), + maximum_fee: newer_transfer_fee.maximum_fee.into(), + }]) + .await + .unwrap_err(); assert_eq!( err, TokenClientError::Client(Box::new(TransportError::TransactionError( @@ -142,8 +146,9 @@ async fn set_fee() { }, .. } = test_transfer_fee_config_with_keypairs(); - let TestContext { context, token, .. } = - TestContext::new(vec![ExtensionInitializationParams::TransferFeeConfig { + let mut context = TestContext::new().await; + context + .init_token_with_mint(vec![ExtensionInitializationParams::TransferFeeConfig { transfer_fee_config_authority: transfer_fee_config_authority.pubkey().into(), withdraw_withheld_authority: withdraw_withheld_authority.pubkey().into(), transfer_fee_basis_points: newer_transfer_fee.transfer_fee_basis_points.into(), @@ -151,6 +156,7 @@ async fn set_fee() { }]) .await .unwrap(); + let token = context.token_context.unwrap().token; // set to something new, old fee not touched let new_transfer_fee_basis_points = u16::MAX; @@ -200,7 +206,7 @@ async fn set_fee() { // warp forward one epoch, new fee becomes old fee during set let newer_transfer_fee = extension.newer_transfer_fee; - context.lock().await.warp_to_slot(10_000).unwrap(); + context.context.lock().await.warp_to_slot(10_000).unwrap(); let new_transfer_fee_basis_points = u16::MAX; let new_maximum_fee = u64::MAX; token @@ -246,11 +252,13 @@ async fn set_fee() { #[tokio::test] async fn fail_set_fee_unsupported_mint() { - let TestContext { - token, + let mut context = TestContext::new().await; + context.init_token_with_mint(vec![]).await.unwrap(); + let TokenContext { mint_authority, + token, .. - } = TestContext::new(vec![]).await.unwrap(); + } = context.token_context.unwrap(); let transfer_fee_basis_points = u16::MAX; let maximum_fee = u64::MAX; let error = token @@ -276,8 +284,9 @@ async fn set_transfer_fee_config_authority() { }, .. } = test_transfer_fee_config_with_keypairs(); - let TestContext { token, .. } = - TestContext::new(vec![ExtensionInitializationParams::TransferFeeConfig { + let mut context = TestContext::new().await; + context + .init_token_with_mint(vec![ExtensionInitializationParams::TransferFeeConfig { transfer_fee_config_authority: transfer_fee_config_authority.pubkey().into(), withdraw_withheld_authority: withdraw_withheld_authority.pubkey().into(), transfer_fee_basis_points: newer_transfer_fee.transfer_fee_basis_points.into(), @@ -285,6 +294,7 @@ async fn set_transfer_fee_config_authority() { }]) .await .unwrap(); + let token = context.token_context.unwrap().token; let new_authority = Keypair::new(); let wrong = Keypair::new(); @@ -414,8 +424,9 @@ async fn set_withdraw_withheld_authority() { }, .. } = test_transfer_fee_config_with_keypairs(); - let TestContext { token, .. } = - TestContext::new(vec![ExtensionInitializationParams::TransferFeeConfig { + let mut context = TestContext::new().await; + context + .init_token_with_mint(vec![ExtensionInitializationParams::TransferFeeConfig { transfer_fee_config_authority: transfer_fee_config_authority.pubkey().into(), withdraw_withheld_authority: withdraw_withheld_authority.pubkey().into(), transfer_fee_basis_points: newer_transfer_fee.transfer_fee_basis_points.into(), @@ -423,6 +434,7 @@ async fn set_withdraw_withheld_authority() { }]) .await .unwrap(); + let token = context.token_context.unwrap().token; let new_authority = Keypair::new(); let wrong = Keypair::new();