#![cfg(not(tarpaulin_include))] #![allow(dead_code)] use std::collections::HashMap; use anyhow::Result as AnyResult; use astroport_test::coins::TestCoin; use astroport_test::cw_multi_test::{AppBuilder, AppResponse, Contract, ContractWrapper, Executor}; use astroport_test::modules::stargate::{MockStargate, StargateApp as TestApp}; use cosmwasm_std::{coin, to_json_binary, Addr, Coin, Decimal, Empty, StdResult, Uint128}; use cw20::{BalanceResponse, Cw20Coin, Cw20ExecuteMsg, Cw20QueryMsg}; use derivative::Derivative; use itertools::Itertools; use astroport::asset::{native_asset_info, token_asset_info, Asset, AssetInfo, PairInfo}; use astroport::factory::{PairConfig, PairType}; use astroport::observation::OracleObservation; use astroport::pair::{ CumulativePricesResponse, Cw20HookMsg, ExecuteMsg, QueryMsg, ReverseSimulationResponse, SimulationResponse, StablePoolParams, }; pub const NATIVE_TOKEN_PRECISION: u8 = 6; use astroport_pair_stable::contract::{execute, instantiate, query, reply}; const INIT_BALANCE: u128 = 1_000_000_000000; pub fn init_native_coins(test_coins: &[TestCoin]) -> Vec { test_coins .iter() .filter_map(|test_coin| match test_coin { TestCoin::Native(name) => { let init_balance = INIT_BALANCE * 10u128.pow(NATIVE_TOKEN_PRECISION as u32); Some(coin(init_balance, name)) } _ => None, }) .collect() } fn token_contract() -> Box> { Box::new(ContractWrapper::new_with_empty( cw20_base::contract::execute, cw20_base::contract::instantiate, cw20_base::contract::query, )) } fn pair_contract() -> Box> { Box::new(ContractWrapper::new_with_empty(execute, instantiate, query).with_reply_empty(reply)) } fn factory_contract() -> Box> { Box::new( ContractWrapper::new_with_empty( astroport_factory::contract::execute, astroport_factory::contract::instantiate, astroport_factory::contract::query, ) .with_reply_empty(astroport_factory::contract::reply), ) } fn store_coin_registry_code() -> Box> { Box::new(ContractWrapper::new_with_empty( astroport_native_coin_registry::contract::execute, astroport_native_coin_registry::contract::instantiate, astroport_native_coin_registry::contract::query, )) } #[derive(Derivative)] #[derivative(Debug)] pub struct Helper { #[derivative(Debug = "ignore")] pub app: TestApp, pub owner: Addr, pub assets: HashMap, pub factory: Addr, pub pair_addr: Addr, pub lp_token: String, pub amp: u64, } impl Helper { pub fn new( owner: &Addr, test_coins: Vec, amp: u64, swap_fee: Option, ) -> AnyResult { let mut app = AppBuilder::new_custom() .with_stargate(MockStargate::default()) .build(|router, _, storage| { router .bank .init_balance(storage, owner, init_native_coins(&test_coins)) .unwrap() }); let mut asset_infos_vec: Vec<_> = test_coins .clone() .into_iter() .filter_map(|coin| Some((coin.clone(), native_asset_info(coin.denom()?)))) .collect(); let token_code_id = app.store_code(token_contract()); test_coins.clone().into_iter().for_each(|coin| { if let Some((name, decimals)) = coin.cw20_init_data() { let token_addr = Self::init_token(&mut app, token_code_id, name, decimals, &owner); asset_infos_vec.push((coin, token_asset_info(token_addr))) } }); let coin_registry_id = app.store_code(store_coin_registry_code()); let coin_registry_address = app .instantiate_contract( coin_registry_id, Addr::unchecked(owner.to_string()), &astroport::native_coin_registry::InstantiateMsg { owner: owner.to_string(), }, &[], "Coin registry", None, ) .unwrap(); app.execute_contract( Addr::unchecked(owner.to_string()), coin_registry_address.clone(), &astroport::native_coin_registry::ExecuteMsg::Add { native_coins: vec![ ("one".to_string(), 6), ("three".to_string(), 6), ("five".to_string(), 6), ("uusd".to_string(), 6), ("ibc/usd".to_string(), 6), ("uluna".to_string(), 6), ], }, &[], ) .unwrap(); let pair_code_id = app.store_code(pair_contract()); let factory_code_id = app.store_code(factory_contract()); let init_msg = astroport::factory::InstantiateMsg { fee_address: None, pair_configs: vec![PairConfig { code_id: pair_code_id, maker_fee_bps: 5000, total_fee_bps: swap_fee.unwrap_or(5u16), pair_type: PairType::Stable {}, is_disabled: false, is_generator_disabled: false, permissioned: false, }], token_code_id, generator_address: None, owner: owner.to_string(), whitelist_code_id: 234u64, coin_registry_address: coin_registry_address.to_string(), tracker_config: None, }; let factory = app.instantiate_contract( factory_code_id, owner.clone(), &init_msg, &[], "FACTORY", None, )?; let asset_infos = asset_infos_vec .clone() .into_iter() .map(|(_, asset_info)| asset_info) .collect_vec(); let init_pair_msg = astroport::factory::ExecuteMsg::CreatePair { pair_type: PairType::Stable {}, asset_infos: asset_infos.clone(), init_params: Some(to_json_binary(&StablePoolParams { amp, owner: None }).unwrap()), }; app.execute_contract(owner.clone(), factory.clone(), &init_pair_msg, &[])?; let resp: PairInfo = app.wrap().query_wasm_smart( &factory, &astroport::factory::QueryMsg::Pair { asset_infos }, )?; Ok(Self { app, owner: owner.clone(), assets: asset_infos_vec.into_iter().collect(), factory, pair_addr: resp.contract_addr, lp_token: resp.liquidity_token, amp, }) } pub fn provide_liquidity( &mut self, sender: &Addr, assets: &[Asset], min_lp_to_receive: Option, ) -> AnyResult { let funds = assets.mock_coins_sent(&mut self.app, sender, &self.pair_addr, SendType::Allowance); let msg = ExecuteMsg::ProvideLiquidity { assets: assets.to_vec(), slippage_tolerance: None, auto_stake: None, receiver: None, min_lp_to_receive, }; self.app .execute_contract(sender.clone(), self.pair_addr.clone(), &msg, &funds) } pub fn withdraw_liquidity( &mut self, sender: &Addr, amount: u128, assets: Vec, min_assets_to_receive: Option>, ) -> AnyResult { self.app.execute_contract( sender.clone(), self.pair_addr.clone(), &ExecuteMsg::WithdrawLiquidity { assets, min_assets_to_receive, }, &[coin(amount, self.lp_token.to_string())], ) } pub fn swap( &mut self, sender: &Addr, offer_asset: &Asset, ask_asset_info: Option, ) -> AnyResult { match &offer_asset.info { AssetInfo::Token { contract_addr } => { let msg = Cw20ExecuteMsg::Send { contract: self.pair_addr.to_string(), amount: offer_asset.amount, msg: to_json_binary(&Cw20HookMsg::Swap { ask_asset_info, belief_price: None, max_spread: None, to: None, }) .unwrap(), }; self.app .execute_contract(sender.clone(), contract_addr.clone(), &msg, &[]) } AssetInfo::NativeToken { .. } => { let funds = offer_asset.mock_coin_sent( &mut self.app, sender, &self.pair_addr, SendType::None, ); let msg = ExecuteMsg::Swap { offer_asset: offer_asset.clone(), ask_asset_info, belief_price: None, max_spread: None, to: None, }; self.app .execute_contract(sender.clone(), self.pair_addr.clone(), &msg, &funds) } } } pub fn simulate_swap( &self, offer_asset: &Asset, ask_asset_info: Option, ) -> StdResult { self.app.wrap().query_wasm_smart( &self.pair_addr, &QueryMsg::Simulation { offer_asset: offer_asset.clone(), ask_asset_info, }, ) } pub fn simulate_reverse_swap( &self, ask_asset: &Asset, offer_asset_info: Option, ) -> StdResult { self.app.wrap().query_wasm_smart( &self.pair_addr, &QueryMsg::ReverseSimulation { ask_asset: ask_asset.clone(), offer_asset_info, }, ) } pub fn query_prices(&self) -> StdResult { self.app .wrap() .query_wasm_smart(&self.pair_addr, &QueryMsg::CumulativePrices {}) } pub fn query_observe(&self, seconds_ago: u64) -> StdResult { self.app .wrap() .query_wasm_smart(&self.pair_addr, &QueryMsg::Observe { seconds_ago }) } fn init_token( app: &mut TestApp, token_code: u64, name: String, decimals: u8, owner: &Addr, ) -> Addr { let init_balance = INIT_BALANCE * 10u128.pow(decimals as u32); app.instantiate_contract( token_code, owner.clone(), &astroport::token::InstantiateMsg { symbol: name.to_string(), name, decimals, initial_balances: vec![Cw20Coin { address: owner.to_string(), amount: Uint128::from(init_balance), }], mint: None, marketing: None, }, &[], "{name}_token", None, ) .unwrap() } pub fn token_balance(&self, token_addr: &Addr, user: &Addr) -> u128 { let resp: BalanceResponse = self .app .wrap() .query_wasm_smart( token_addr, &Cw20QueryMsg::Balance { address: user.to_string(), }, ) .unwrap(); resp.balance.u128() } pub fn native_balance(&self, denom: impl Into, user: &Addr) -> u128 { self.app .wrap() .query_balance(user, denom) .unwrap() .amount .u128() } pub fn coin_balance(&self, coin: &TestCoin, user: &Addr) -> u128 { match &self.assets[coin] { AssetInfo::Token { contract_addr } => self.token_balance(contract_addr, user), AssetInfo::NativeToken { denom } => self.native_balance(denom, user), } } pub fn give_me_money(&mut self, assets: &[Asset], recipient: &Addr) { let funds = assets.mock_coins_sent(&mut self.app, &self.owner, recipient, SendType::Transfer); if !funds.is_empty() { self.app .send_tokens(self.owner.clone(), recipient.clone(), &funds) .unwrap(); } } pub fn observe_price(&self, seconds_ago: u64) -> StdResult { self.app .wrap() .query_wasm_smart::( &self.pair_addr, &QueryMsg::Observe { seconds_ago }, ) .map(|val| val.price) } } #[derive(Clone, Copy)] pub enum SendType { Allowance, Transfer, None, } pub trait AssetExt { fn mock_coin_sent( &self, app: &mut TestApp, user: &Addr, spender: &Addr, typ: SendType, ) -> Vec; } impl AssetExt for Asset { fn mock_coin_sent( &self, app: &mut TestApp, user: &Addr, spender: &Addr, typ: SendType, ) -> Vec { let mut funds = vec![]; match &self.info { AssetInfo::Token { contract_addr } if !self.amount.is_zero() => { let msg = match typ { SendType::Allowance => Cw20ExecuteMsg::IncreaseAllowance { spender: spender.to_string(), amount: self.amount, expires: None, }, SendType::Transfer => Cw20ExecuteMsg::Transfer { recipient: spender.to_string(), amount: self.amount, }, _ => unimplemented!(), }; app.execute_contract(user.clone(), contract_addr.clone(), &msg, &[]) .unwrap(); } AssetInfo::NativeToken { denom } if !self.amount.is_zero() => { funds = vec![coin(self.amount.u128(), denom)]; } _ => {} } funds } } pub trait AssetsExt { fn mock_coins_sent( &self, app: &mut TestApp, user: &Addr, spender: &Addr, typ: SendType, ) -> Vec; } impl AssetsExt for &[Asset] { fn mock_coins_sent( &self, app: &mut TestApp, user: &Addr, spender: &Addr, typ: SendType, ) -> Vec { let mut funds = vec![]; for asset in self.iter() { funds.extend(asset.mock_coin_sent(app, user, spender, typ)); } funds } } pub trait AppExtension { fn next_block(&mut self, time: u64); } impl AppExtension for TestApp { fn next_block(&mut self, time: u64) { self.update_block(|block| { block.time = block.time.plus_seconds(time); block.height += 1 }); } }