Refactor TestContext to enable unwrap_err (#2805)

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,
+        instruction,
-        processor::Processor,
         state::{Account, Mint},
     },
 };
 
 #[tokio::test]
 async fn no_extensions() {
-    let program_test = ProgramTest::new("spl_token_2022", id(), processor!(Processor::process));
+    let context = TestContext::new().await;
-    let mut ctx = program_test.start_with_context().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 context = TestContext::new().await;
-    let mut ctx = program_test.start_with_context().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 context = TestContext::new().await;
-    let mut ctx = program_test.start_with_context().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();

token/program-2022/tests/initialize_mint.rs

@@ -2,8 +2,8 @@
 
 mod program_test;
 use {
-    program_test::TestContext,
+    program_test::{TestContext, TokenContext},
-    solana_program_test::{processor, tokio, ProgramTest},
+    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,
+        instruction,
-        processor::Processor,
         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 context = TestContext::new().await;
-    let mut ctx = program_test.start_with_context().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 context = TestContext::new().await;
-    let mut ctx = program_test.start_with_context().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 {
+    } = context.token_context.unwrap();
-        close_authority,
-    }])
-    .await
-    .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 context = TestContext::new().await;
-    let mut ctx = program_test.start_with_context().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 context = TestContext::new().await;
-    let mut ctx = program_test.start_with_context().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 context = TestContext::new().await;
-    let mut ctx = program_test.start_with_context().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 context = TestContext::new().await;
-    let mut ctx = program_test.start_with_context().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 context = TestContext::new().await;
-    let mut ctx = program_test.start_with_context().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();

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 {
+    } = context.token_context.unwrap();
-        close_authority,
-    }])
-    .await
-    .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, .. } =
+    let mut context = TestContext::new().await;
-        TestContext::new(vec![ExtensionInitializationParams::MintCloseAuthority {
+    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, .. } =
+    let mut context = TestContext::new().await;
-        TestContext::new(vec![ExtensionInitializationParams::MintCloseAuthority {
+    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 {
+    let mut context = TestContext::new().await;
-        token,
+    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 {
+    } = context.token_context.unwrap();
-        close_authority: Some(close_authority.pubkey()),
-    }])
-    .await
-    .unwrap();
 
     // mint a token
     let owner = Pubkey::new_unique();

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<ProgramBanksClientProcessTransaction, Keypair>,
     pub alice: Keypair,
     pub bob: Keypair,
-    pub context: Arc<Mutex<ProgramTestContext>>, // ProgramTestContext needs to #[derive(Debug)]
+}
+
+pub struct TestContext {
+    pub context: Arc<Mutex<ProgramTestContext>>,
+    pub token_context: Option<TokenContext>,
 }
 
 impl TestContext {
-    pub async fn new(
+    pub async fn new() -> Self {
-        extension_init_params: Vec<ExtensionInitializationParams>,
-    ) -> TokenResult<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<ExtensionInitializationParams>,
+    ) -> TokenResult<()> {
+        let payer = keypair_clone(&self.context.lock().await.payer);
         let client: Arc<dyn ProgramClient<ProgramBanksClientProcessTransaction>> =
             Arc::new(ProgramBanksClient::new_from_context(
-                Arc::clone(&context),
+                Arc::clone(&self.context),
                 ProgramBanksClientProcessTransaction,
             ));
 
@@ -51,15 +64,15 @@ impl TestContext {
             extension_init_params,
         )
         .await?;
-
+        self.token_context = Some(TokenContext {
-        Ok(Self {
             decimals,
             mint_authority,
             token,
             alice: Keypair::new(),
             bob: Keypair::new(),
-            context,
+        });
-        })
+
+        Ok(())
     }
 }
 

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 {
+    } = context.token_context.unwrap();
-        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 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 {
+    let mut context = TestContext::new().await;
-        transfer_fee_config_authority: transfer_fee_config_authority.into(),
+    let err = context
-        withdraw_withheld_authority: Some(Pubkey::default()),
+        .init_token_with_mint(vec![ExtensionInitializationParams::TransferFeeConfig {
-        transfer_fee_basis_points: newer_transfer_fee.transfer_fee_basis_points.into(),
+            transfer_fee_config_authority: transfer_fee_config_authority.into(),
-        maximum_fee: newer_transfer_fee.maximum_fee.into(),
+            withdraw_withheld_authority: Some(Pubkey::default()),
-    }])
+            transfer_fee_basis_points: newer_transfer_fee.transfer_fee_basis_points.into(),
-    .await
+            maximum_fee: newer_transfer_fee.maximum_fee.into(),
-    .err()
+        }])
-    .unwrap();
+        .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, .. } =
+    let mut context = TestContext::new().await;
-        TestContext::new(vec![ExtensionInitializationParams::TransferFeeConfig {
+    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 {
+    let mut context = TestContext::new().await;
-        token,
+    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, .. } =
+    let mut context = TestContext::new().await;
-        TestContext::new(vec![ExtensionInitializationParams::TransferFeeConfig {
+    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, .. } =
+    let mut context = TestContext::new().await;
-        TestContext::new(vec![ExtensionInitializationParams::TransferFeeConfig {
+    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();