From 27f882a333cc054623d23319440b8b791f3656ce Mon Sep 17 00:00:00 2001 From: Christian Kamm Date: Mon, 21 Mar 2022 10:45:55 +0100 Subject: [PATCH] Tests: More flexible test context construction --- programs/mango-v4/tests/program_test/mod.rs | 161 +++++++++++------- programs/mango-v4/tests/test_basic.rs | 2 +- .../tests/test_group_address_lookup_tables.rs | 2 +- .../mango-v4/tests/test_health_compute.rs | 4 +- programs/mango-v4/tests/test_margin_trade.rs | 27 +-- programs/mango-v4/tests/test_perp.rs | 2 +- .../mango-v4/tests/test_position_lifetime.rs | 2 +- programs/mango-v4/tests/test_serum.rs | 2 +- 8 files changed, 119 insertions(+), 83 deletions(-) diff --git a/programs/mango-v4/tests/program_test/mod.rs b/programs/mango-v4/tests/program_test/mod.rs index e392e3330..f7bed075b 100644 --- a/programs/mango-v4/tests/program_test/mod.rs +++ b/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 solana: Arc, - pub mints: Vec, - pub users: Vec, - pub quote_index: usize, - pub serum: Arc, +pub struct MarginTradeCookie { + pub program: Pubkey, + pub token_account: Keypair, + pub token_account_owner: Pubkey, + pub token_account_bump: u8, } -impl TestContext { - pub async fn new( - test_opt: Option, - margin_trade_program_id: Option<&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)) - }; +pub struct TestContextBuilder { + test: ProgramTest, + program_log_capture: Arc>>, + mint0: Pubkey, +} - let serum_program_id = anchor_spl::dex::id(); - test.add_program("serum_dex", serum_program_id, None); +impl TestContextBuilder { + 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 { let mut mints: Vec = 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 - 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(), - ); - } + mints + } - // Users + pub fn create_users(&mut self, mints: &[MintCookie]) -> Vec { 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; - let rent = context.banks_client.get_rent().await.unwrap(); + users + } - let solana = Arc::new(SolanaCookie { - context: RefCell::new(context), - rent, - program_log: program_log_capture.clone(), - }); + pub fn add_serum_program(&mut self) -> Pubkey { + let serum_program_id = anchor_spl::dex::id(); + self.test.add_program("serum_dex", serum_program_id, None); + 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 { + 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, + pub mints: Vec, + pub users: Vec, + pub serum: Arc, +} + +impl TestContext { + pub async fn new() -> Self { + TestContextBuilder::new().start_default().await + } } diff --git a/programs/mango-v4/tests/test_basic.rs b/programs/mango-v4/tests/test_basic.rs index af9ed5527..ae7fa6d4d 100644 --- a/programs/mango-v4/tests/test_basic.rs +++ b/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(); diff --git a/programs/mango-v4/tests/test_group_address_lookup_tables.rs b/programs/mango-v4/tests/test_group_address_lookup_tables.rs index 64bf68f2a..94affa9c2 100644 --- a/programs/mango-v4/tests/test_group_address_lookup_tables.rs +++ b/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(); diff --git a/programs/mango-v4/tests/test_health_compute.rs b/programs/mango-v4/tests/test_health_compute.rs index ec3c537ae..17ada3078 100644 --- a/programs/mango-v4/tests/test_health_compute.rs +++ b/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(); diff --git a/programs/mango-v4/tests/test_margin_trade.rs b/programs/mango-v4/tests/test_margin_trade.rs index 44b21add6..430e14d76 100644 --- a/programs/mango-v4/tests/test_margin_trade.rs +++ b/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 = - Pubkey::from_str("J83w4HKfqxwcq3BEMMkPFSppX3gqekLyLJBexebFVkix").unwrap(); - let margin_trade_token_account = Keypair::new(); - 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 mut builder = TestContextBuilder::new(); + let margin_trade = builder.add_margin_trade_program(); + let context = builder.start_default().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, - deposit_account: margin_trade_token_account.pubkey(), - deposit_account_owner: mtta_owner, + margin_trade_program_id: margin_trade.program, + deposit_account: margin_trade.token_account.pubkey(), + 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 ); diff --git a/programs/mango-v4/tests/test_perp.rs b/programs/mango-v4/tests/test_perp.rs index 7df05a14a..d33f0ee43 100644 --- a/programs/mango-v4/tests/test_perp.rs +++ b/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(); diff --git a/programs/mango-v4/tests/test_position_lifetime.rs b/programs/mango-v4/tests/test_position_lifetime.rs index 918f34444..d0e7c4aa2 100644 --- a/programs/mango-v4/tests/test_position_lifetime.rs +++ b/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(); diff --git a/programs/mango-v4/tests/test_serum.rs b/programs/mango-v4/tests/test_serum.rs index a55806515..3c8e1fba3 100644 --- a/programs/mango-v4/tests/test_serum.rs +++ b/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();