Tests: More flexible test context construction

programs/mango-v4/tests/program_test/mod.rs

@@ -1,4 +1,7 @@
+#![allow(dead_code)]
+
 use std::cell::RefCell;
+use std::str::FromStr;
 use std::{sync::Arc, sync::RwLock};
 
 use log::*;
@@ -73,29 +76,22 @@ impl Log for LoggerWrapper {
     fn flush(&self) {}
 }
 
-pub struct TestContext {
+pub struct MarginTradeCookie {
-    pub solana: Arc<SolanaCookie>,
+    pub program: Pubkey,
-    pub mints: Vec<MintCookie>,
+    pub token_account: Keypair,
-    pub users: Vec<UserCookie>,
+    pub token_account_owner: Pubkey,
-    pub quote_index: usize,
+    pub token_account_bump: u8,
-    pub serum: Arc<SerumCookie>,
 }
 
-impl TestContext {
+pub struct TestContextBuilder {
-    pub async fn new(
+    test: ProgramTest,
-        test_opt: Option<ProgramTest>,
+    program_log_capture: Arc<RwLock<Vec<String>>>,
-        margin_trade_program_id: Option<&Pubkey>,
+    mint0: Pubkey,
-        margin_trade_token_account: Option<&Keypair>,
+}
-        mtta_owner: Option<&Pubkey>,
-    ) -> Self {
-        let mut test = if test_opt.is_some() {
-            test_opt.unwrap()
-        } else {
-            ProgramTest::new("mango_v4", mango_v4::id(), processor!(mango_v4::entry))
-        };
 
-        let serum_program_id = anchor_spl::dex::id();
+impl TestContextBuilder {
-        test.add_program("serum_dex", serum_program_id, None);
+    pub fn new() -> Self {
+        let mut test = ProgramTest::new("mango_v4", mango_v4::id(), processor!(mango_v4::entry));
 
         // We need to intercept logs to capture program log output
         let log_filter = "solana_rbpf=trace,\
@@ -115,9 +111,18 @@ impl TestContext {
         // intentionally set to half the limit, to catch potential problems early
         test.set_compute_max_units(100000);
 
-        // Setup the environment
+        Self {
+            test,
+            program_log_capture,
+            mint0: Pubkey::new_unique(),
+        }
+    }
 
-        // Mints
+    pub fn test(&mut self) -> &mut ProgramTest {
+        &mut self.test
+    }
+
+    pub fn create_mints(&mut self) -> Vec<MintCookie> {
         let mut mints: Vec<MintCookie> = vec![
             MintCookie {
                 index: 0,
@@ -125,7 +130,7 @@ impl TestContext {
                 unit: 10u64.pow(6) as f64,
                 base_lot: 100 as f64,
                 quote_lot: 10 as f64,
-                pubkey: Pubkey::default(),
+                pubkey: self.mint0,
                 authority: Keypair::new(),
             }, // symbol: "MNGO".to_string()
         ];
@@ -150,7 +155,7 @@ impl TestContext {
             }
             mints[mint_index].pubkey = mint_pk;
 
-            test.add_packable_account(
+            self.test.add_packable_account(
                 mint_pk,
                 u32::MAX as u64,
                 &Mint {
@@ -162,36 +167,16 @@ impl TestContext {
                 &spl_token::id(),
             );
         }
-        let quote_index = mints.len() - 1;
 
-        // margin trade
+        mints
-        if margin_trade_program_id.is_some() {
+    }
-            test.add_program(
-                "margin_trade",
-                *margin_trade_program_id.unwrap(),
-                std::option::Option::None,
-            );
-            test.add_packable_account(
-                margin_trade_token_account.unwrap().pubkey(),
-                u32::MAX as u64,
-                &Account {
-                    mint: mints[0].pubkey,
-                    owner: *mtta_owner.unwrap(),
-                    amount: 0,
-                    state: AccountState::Initialized,
-                    is_native: COption::None,
-                    ..Account::default()
-                },
-                &spl_token::id(),
-            );
-        }
 
-        // Users
+    pub fn create_users(&mut self, mints: &[MintCookie]) -> Vec<UserCookie> {
         let num_users = 4;
         let mut users = Vec::new();
         for _ in 0..num_users {
             let user_key = Keypair::new();
-            test.add_account(
+            self.test.add_account(
                 user_key.pubkey(),
                 solana_sdk::account::Account::new(
                     u32::MAX as u64,
@@ -205,7 +190,7 @@ impl TestContext {
             let mut token_accounts = Vec::new();
             for mint_index in 0..mints.len() {
                 let token_key = Pubkey::new_unique();
-                test.add_packable_account(
+                self.test.add_packable_account(
                     token_key,
                     u32::MAX as u64,
                     &spl_token::state::Account {
@@ -226,14 +211,51 @@ impl TestContext {
             });
         }
 
-        let mut context = test.start_with_context().await;
+        users
-        let rent = context.banks_client.get_rent().await.unwrap();
+    }
 
-        let solana = Arc::new(SolanaCookie {
+    pub fn add_serum_program(&mut self) -> Pubkey {
-            context: RefCell::new(context),
+        let serum_program_id = anchor_spl::dex::id();
-            rent,
+        self.test.add_program("serum_dex", serum_program_id, None);
-            program_log: program_log_capture.clone(),
+        serum_program_id
-        });
+    }
+
+    pub fn add_margin_trade_program(&mut self) -> MarginTradeCookie {
+        let program = Pubkey::from_str("J83w4HKfqxwcq3BEMMkPFSppX3gqekLyLJBexebFVkix").unwrap();
+        let token_account = Keypair::new();
+        let (token_account_owner, token_account_bump) =
+            Pubkey::find_program_address(&[b"MarginTrade"], &program);
+
+        self.test
+            .add_program("margin_trade", program, std::option::Option::None);
+        self.test.add_packable_account(
+            token_account.pubkey(),
+            u32::MAX as u64,
+            &Account {
+                mint: self.mint0,
+                owner: token_account_owner,
+                amount: 0,
+                state: AccountState::Initialized,
+                is_native: COption::None,
+                ..Account::default()
+            },
+            &spl_token::id(),
+        );
+
+        MarginTradeCookie {
+            program,
+            token_account,
+            token_account_owner,
+            token_account_bump,
+        }
+    }
+
+    pub async fn start_default(mut self) -> TestContext {
+        let mints = self.create_mints();
+        let users = self.create_users(&mints);
+        let serum_program_id = self.add_serum_program();
+
+        let solana = self.start().await;
 
         let serum = Arc::new(SerumCookie {
             solana: solana.clone(),
@@ -244,8 +266,33 @@ impl TestContext {
             solana: solana.clone(),
             mints,
             users,
-            quote_index,
             serum,
         }
     }
+
+    pub async fn start(self) -> Arc<SolanaCookie> {
+        let mut context = self.test.start_with_context().await;
+        let rent = context.banks_client.get_rent().await.unwrap();
+
+        let solana = Arc::new(SolanaCookie {
+            context: RefCell::new(context),
+            rent,
+            program_log: self.program_log_capture.clone(),
+        });
+
+        solana
+    }
+}
+
+pub struct TestContext {
+    pub solana: Arc<SolanaCookie>,
+    pub mints: Vec<MintCookie>,
+    pub users: Vec<UserCookie>,
+    pub serum: Arc<SerumCookie>,
+}
+
+impl TestContext {
+    pub async fn new() -> Self {
+        TestContextBuilder::new().start_default().await
+    }
 }

programs/mango-v4/tests/test_basic.rs

@@ -13,7 +13,7 @@ mod program_test;
 // that they work in principle. It should be split up / renamed.
 #[tokio::test]
 async fn test_basic() -> Result<(), TransportError> {
-    let context = TestContext::new(Option::None, Option::None, Option::None, Option::None).await;
+    let context = TestContext::new().await;
     let solana = &context.solana.clone();
 
     let admin = &Keypair::new();

programs/mango-v4/tests/test_group_address_lookup_tables.rs

@@ -14,7 +14,7 @@ mod program_test;
 // that they work in principle. It should be split up / renamed.
 #[tokio::test]
 async fn test_group_address_lookup_tables() -> Result<()> {
-    let context = TestContext::new(None, None, None, None).await;
+    let context = TestContext::new().await;
     let solana = &context.solana.clone();
 
     let admin = &Keypair::new();

programs/mango-v4/tests/test_health_compute.rs

@@ -10,7 +10,7 @@ mod program_test;
 // Try to reach compute limits in health checks by having many different tokens in an account
 #[tokio::test]
 async fn test_health_compute_tokens() -> Result<(), TransportError> {
-    let context = TestContext::new(Option::None, Option::None, Option::None, Option::None).await;
+    let context = TestContext::new().await;
     let solana = &context.solana.clone();
 
     let admin = &Keypair::new();
@@ -73,7 +73,7 @@ async fn test_health_compute_tokens() -> Result<(), TransportError> {
 // Try to reach compute limits in health checks by having many serum markets in an account
 #[tokio::test]
 async fn test_health_compute_serum() -> Result<(), TransportError> {
-    let context = TestContext::new(Option::None, Option::None, Option::None, Option::None).await;
+    let context = TestContext::new().await;
     let solana = &context.solana.clone();
 
     let admin = &Keypair::new();

programs/mango-v4/tests/test_margin_trade.rs

@@ -2,11 +2,9 @@
 
 use anchor_lang::InstructionData;
 use fixed::types::I80F48;
-use solana_program::pubkey::Pubkey;
 use solana_program_test::*;
 use solana_sdk::signature::Signer;
 use solana_sdk::{signature::Keypair, transport::TransportError};
-use std::str::FromStr;
 
 use mango_v4::state::*;
 use program_test::*;
@@ -17,18 +15,9 @@ mod program_test;
 // that they work in principle. It should be split up / renamed.
 #[tokio::test]
 async fn test_margin_trade() -> Result<(), TransportError> {
-    let margin_trade_program_id =
+    let mut builder = TestContextBuilder::new();
-        Pubkey::from_str("J83w4HKfqxwcq3BEMMkPFSppX3gqekLyLJBexebFVkix").unwrap();
+    let margin_trade = builder.add_margin_trade_program();
-    let margin_trade_token_account = Keypair::new();
+    let context = builder.start_default().await;
-    let (mtta_owner, mtta_bump_seeds) =
-        Pubkey::find_program_address(&[b"MarginTrade"], &margin_trade_program_id);
-    let context = TestContext::new(
-        Option::None,
-        Some(&margin_trade_program_id),
-        Some(&margin_trade_token_account),
-        Some(&mtta_owner),
-    )
-    .await;
     let solana = &context.solana.clone();
 
     let admin = &Keypair::new();
@@ -117,14 +106,14 @@ async fn test_margin_trade() -> Result<(), TransportError> {
                 account,
                 owner,
                 mango_token_vault: vault,
-                margin_trade_program_id,
+                margin_trade_program_id: margin_trade.program,
-                deposit_account: margin_trade_token_account.pubkey(),
+                deposit_account: margin_trade.token_account.pubkey(),
-                deposit_account_owner: mtta_owner,
+                deposit_account_owner: margin_trade.token_account_owner,
                 margin_trade_program_ix_cpi_data: {
                     let ix = margin_trade::instruction::MarginTrade {
                         amount_from: 2,
                         amount_to: 1,
-                        deposit_account_owner_bump_seeds: mtta_bump_seeds,
+                        deposit_account_owner_bump_seeds: margin_trade.token_account_bump,
                     };
                     ix.data()
                 },
@@ -139,7 +128,7 @@ async fn test_margin_trade() -> Result<(), TransportError> {
     );
     assert_eq!(
         solana
-            .token_account_balance(margin_trade_token_account.pubkey())
+            .token_account_balance(margin_trade.token_account.pubkey())
             .await,
         withdraw_amount - deposit_amount
     );

programs/mango-v4/tests/test_perp.rs

@@ -9,7 +9,7 @@ mod program_test;
 
 #[tokio::test]
 async fn test_perp() -> Result<(), TransportError> {
-    let context = TestContext::new(Option::None, Option::None, Option::None, Option::None).await;
+    let context = TestContext::new().await;
     let solana = &context.solana.clone();
 
     let admin = &Keypair::new();

programs/mango-v4/tests/test_position_lifetime.rs

@@ -12,7 +12,7 @@ mod program_test;
 // Check opening and closing positions
 #[tokio::test]
 async fn test_position_lifetime() -> Result<()> {
-    let context = TestContext::new(None, None, None, None).await;
+    let context = TestContext::new().await;
     let solana = &context.solana.clone();
 
     let admin = &Keypair::new();

programs/mango-v4/tests/test_serum.rs

@@ -11,7 +11,7 @@ mod program_test;
 
 #[tokio::test]
 async fn test_serum() -> Result<(), TransportError> {
-    let context = TestContext::new(Option::None, Option::None, Option::None, Option::None).await;
+    let context = TestContext::new().await;
     let solana = &context.solana.clone();
 
     let admin = &Keypair::new();