From fed6800c2cc7a618e2cc5d536ff7c275af239c15 Mon Sep 17 00:00:00 2001 From: Tuan Tran Date: Thu, 23 Nov 2023 21:32:56 +0700 Subject: [PATCH] feat(query): allow contract owner to add supported pool manually, refactor tests --- src/contract.rs | 215 ++++++++++++++++++++++++++++++--------- src/error.rs | 2 + src/integration_tests.rs | 62 ++++++++--- src/msg.rs | 31 ++++-- src/state.rs | 35 ++++++- 5 files changed, 269 insertions(+), 76 deletions(-) diff --git a/src/contract.rs b/src/contract.rs index 8cf53bb..9fd0657 100644 --- a/src/contract.rs +++ b/src/contract.rs @@ -8,9 +8,10 @@ use cw2::set_contract_version; use crate::error::ContractError; use crate::msg::{ExecuteMsg, InstantiateMsg, QueryMsg}; use crate::state::{State, STATE}; -use astroport::pair; -use cosmwasm_std::StdError; -use cosmwasm_std::WasmMsg; +use astroport::pair::{self}; +use cosmwasm_std::{Addr, Order, StdError, WasmMsg}; + +use self::query::*; // version info for migration info const CONTRACT_NAME: &str = "crates.io:lq-express-sm"; const CONTRACT_VERSION: &str = env!("CARGO_PKG_VERSION"); @@ -42,6 +43,11 @@ pub fn execute( ) -> Result { match msg { ExecuteMsg::Astro { pair_address } => execute::astro_exec(deps, info, pair_address), + ExecuteMsg::AddSupportedPool { + pool_address, + token_1, + token_2, + } => execute::add_supported_pool(deps, info, pool_address, token_1, token_2), } } @@ -73,7 +79,6 @@ fn handle_swap_reply(_deps: DepsMut, _env: Env, msg: Reply) -> StdResult Result { - let inj_amount = cw_utils::must_pay(&info, "inj").unwrap().u128(); + let pool_infos: Vec<_> = POOL_INFO + .idx + .address + .prefix(pair_address.clone()) + .range(deps.storage, None, None, Order::Ascending) + .take(1) + .flatten() + .collect(); + if pool_infos.is_empty() { + return Err(ContractError::PoolNotExist {}); + } + let pair_info = pool_infos + .first() + .map(|pool_info| (pool_info.1.token_1.clone(), pool_info.1.token_2.clone())) + .unwrap(); + // Based on cw_utils::must_pay implementation + let coin = cw_utils::one_coin(&info)?; + let offer_asset = coin.denom.clone(); + let amount: u128 = if offer_asset != pair_info.0 && offer_asset != pair_info.1 { + return Err(ContractError::Payment( + cw_utils::PaymentError::MissingDenom(coin.denom.to_string()), + )); + } else { + coin.amount.into() + }; + let asked_asset = pair_info.1.clone(); // Pair of hINJ-INJ on testnet + let swap_astro_msg = pair::ExecuteMsg::Swap { - offer_asset: Asset::native("inj", inj_amount), + offer_asset: Asset::native(&offer_asset, amount), ask_asset_info: None, belief_price: None, max_spread: Some(Decimal::percent(50)), @@ -114,84 +147,166 @@ pub mod execute { let exec_cw20_mint_msg = WasmMsg::Execute { contract_addr: pair_address.clone(), msg: to_json_binary(&swap_astro_msg)?, - funds: coins(inj_amount, "inj"), + funds: coins(amount, &offer_asset), }; + assert!(offer_asset == "inj"); let submessage = SubMsg::reply_on_success(exec_cw20_mint_msg, SWAP_REPLY_ID); let res = Response::new() .add_submessage(submessage) .add_attribute("action", "swap") .add_attribute("pair", pair_address) - .add_attribute("offer_asset", "hinj") - .add_attribute("ask_asset_info", "inj"); + .add_attribute("offer_asset", offer_asset) + .add_attribute("ask_asset_info", asked_asset); Ok(res) } + pub fn add_supported_pool( + deps: DepsMut, + info: MessageInfo, + pool_address: String, + token_1: String, + token_2: String, + ) -> Result { + // Check authorization + if info.sender != STATE.load(deps.storage)?.owner { + return Err(ContractError::Unauthorized {}); + } + let pool_info = PoolInfo { + address: Addr::unchecked(pool_address), + token_1: token_1.clone(), + token_2: token_2.clone(), + }; + POOL_INFO.save( + deps.storage, + [token_1.clone(), token_2.clone()].join("_").as_str(), + &pool_info, + )?; + + Ok(Response::new() + .add_attribute("method", "add_supported_pool") + .add_attribute("pool_address", pool_info.address) + .add_attribute("token 1", token_1) + .add_attribute("token2", token_2)) + } } #[cfg_attr(not(feature = "library"), entry_point)] -pub fn query(_deps: Deps, _env: Env, msg: QueryMsg) -> StdResult { - match msg {} +pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult { + match msg { + QueryMsg::GetPair { pool_address } => to_json_binary(&query_pair(deps, env, pool_address)?), + QueryMsg::GetPoolAddr { token_1, token_2 } => { + to_json_binary(&query_pool_addr(deps, env, token_1, token_2)?) + } + } } -pub mod query {} +pub mod query { + use crate::msg::{GetPairResponse, GetPoolAddrResponse}; + use crate::state::POOL_INFO; + use crate::{contract::*, error}; + use cosmwasm_std::StdResult; + + pub fn query_pair(deps: Deps, env: Env, pool_address: String) -> StdResult { + let pair: Vec<_> = POOL_INFO + .idx + .address + .prefix(pool_address) + .range(deps.storage, None, None, Order::Ascending) + .flatten() + .collect(); + if pair.is_empty() { + return Err(StdError::GenericErr { + msg: "pool address not found".to_string(), + }); + } + let (token_1, token_2): (String, String) = pair + .first() + .iter() + .map(|pair| (pair.1.token_1.clone(), pair.1.token_2.clone())) + .unzip(); + let resp = GetPairResponse { token_1, token_2 }; + Ok(resp) + } + pub fn query_pool_addr( + deps: Deps, + env: Env, + token_1: String, + token_2: String, + ) -> StdResult { + let token_1 = token_1.to_lowercase(); + let token_2 = token_2.to_lowercase(); + + let pools = POOL_INFO + .idx + .pair + .prefix((token_1, token_2)) + .range(deps.storage, None, None, Order::Ascending) + .flatten() + .collect::>(); + + if pools.is_empty() { + return Err(StdError::GenericErr { + msg: "No pool exist for this pair yet".to_string(), + }); + } + let pool_addresses = pools + .iter() + .map(|pool_info| pool_info.1.address.to_string().clone()) + .collect::>(); + return Ok(GetPoolAddrResponse { pool_addresses }); + } +} #[cfg(test)] mod tests { + use crate::msg::GetPairResponse; + use super::*; use cosmwasm_std::coins; - use cosmwasm_std::testing::{mock_dependencies, mock_env, mock_info}; use cosmwasm_std::Addr; use cw_multi_test::{App, ContractWrapper, Executor}; #[test] - fn proper_initialization() { - let mut deps = mock_dependencies(); - - let msg = InstantiateMsg {}; - let info = mock_info("creator", &coins(1000, "earth")); - - // we can just call .unwrap() to assert this was a success - let res = instantiate(deps.as_mut(), mock_env(), info, msg).unwrap(); - assert_eq!(0, res.messages.len()); - } - - #[test] - fn astro_test() { - let mut app = App::new(|router, _, storage| { - router - .bank - .init_balance(storage, &Addr::unchecked("user"), coins(1000, "uinj")) - .unwrap() - }); + fn test_add_pool() { + let mut app = App::default(); + let owner = Addr::unchecked("owner"); let code = ContractWrapper::new(execute, instantiate, query); let code_id = app.store_code(Box::new(code)); + let addr = app .instantiate_contract( code_id, - Addr::unchecked("user"), + owner.clone(), &InstantiateMsg {}, &[], "Contract", None, ) .unwrap(); - let _ = app - .execute_contract( - Addr::unchecked("user"), - addr, - &ExecuteMsg::Astro { - pair_address: "pair".to_string(), + let msg = ExecuteMsg::AddSupportedPool { + pool_address: "pool1".to_string(), + token_1: "inj".to_string(), + token_2: "atom".to_string(), + }; + app.execute_contract(owner.clone(), addr.clone(), &msg, &[]) + .unwrap(); + + app.update_block(|b| b.height += 1); + + let resp: GetPairResponse = app + .wrap() + .query_wasm_smart( + addr.clone(), + &QueryMsg::GetPair { + pool_address: "pool1".to_string(), }, - &coins(10, "uinj"), ) .unwrap(); - // let wasm = resp.events.iter().find(|ev| ev.ty == "wasm").unwrap(); - // assert_eq!( - // wasm.attributes - // .iter() - // .find(|attr| attr.key == "action") - // .unwrap() - // .value, - // "swap" - // ); + assert_eq!( + resp, + GetPairResponse { + token_1: "inj".to_string(), + token_2: "atom".to_string() + } + ) } } diff --git a/src/error.rs b/src/error.rs index 8b16c29..349f4fe 100644 --- a/src/error.rs +++ b/src/error.rs @@ -16,4 +16,6 @@ pub enum ContractError { ExceedMintableBlock {}, #[error("Exceed maximum mintable amount")] ExceedMaximumMintableAmount {}, + #[error("Pool not exists")] + PoolNotExist {}, } diff --git a/src/integration_tests.rs b/src/integration_tests.rs index e51d02c..fa49641 100644 --- a/src/integration_tests.rs +++ b/src/integration_tests.rs @@ -2,6 +2,7 @@ #[cfg(test)] mod tests { use crate::helpers::CwTemplateContract; + use crate::msg::ExecuteMsg; use crate::msg::InstantiateMsg; use astroport::asset::{Asset, AssetInfo, PairInfo}; use astroport::factory::{InstantiateMsg as FactoryInitMsg, PairConfig, PairType}; @@ -99,6 +100,10 @@ mod tests { denom: "ttt".to_string(), amount: Uint128::new(10_000_000_000_000), }, + Coin { + denom: "abc".to_string(), + amount: Uint128::new(10_000_000_000_000), + }, ], ) .unwrap(); @@ -196,10 +201,10 @@ mod tests { let msg = PairInitMsg { asset_infos: [ AssetInfo::NativeToken { - denom: "ttt".to_string(), + denom: NATIVE_DENOM.to_string(), }, AssetInfo::NativeToken { - denom: NATIVE_DENOM.to_string(), + denom: "ttt".to_string(), }, ] .to_vec(), @@ -254,15 +259,15 @@ mod tests { assets: vec![ Asset { info: AssetInfo::NativeToken { - denom: "ttt".to_string(), + denom: "inj".to_string(), }, - amount: ttt_amount, + amount: inj_amount, }, Asset { info: AssetInfo::NativeToken { - denom: "inj".to_string(), + denom: "ttt".to_string(), }, - amount: inj_amount, + amount: ttt_amount, }, ], slippage_tolerance, @@ -314,6 +319,19 @@ mod tests { let ttt_amount = Uint128::new(1_000_000_000000); let inj_offer = Uint128::new(1_000000); + // Add supported pool + app.execute_contract( + owner.clone(), + my_contract.addr(), + &ExecuteMsg::AddSupportedPool { + pool_address: pair_contract.addr().to_string(), + token_1: "inj".into(), + token_2: "ttt".into(), + }, + &[], + ) + .unwrap(); + let (msg, coins) = provide_liquidity_msg(ttt_amount, inj_amount, None, None, &token_contract); @@ -341,11 +359,11 @@ mod tests { Some(PairInfo { asset_infos: vec![ AssetInfo::NativeToken { - denom: "ttt".to_string() + denom: "inj".to_string() }, AssetInfo::NativeToken { - denom: "inj".to_string() - } + denom: "ttt".to_string() + }, ], contract_addr: Addr::unchecked("contract4"), liquidity_token: Addr::unchecked("contract5"), @@ -358,7 +376,7 @@ mod tests { pair_contract.addr(), &PairQueryMsg::AssetBalanceAt { asset_info: AssetInfo::NativeToken { - denom: "inj".to_owned(), + denom: "inj".to_string(), }, block_height: app.block_info().height.into(), }, @@ -372,7 +390,7 @@ mod tests { pair_contract.addr(), &PairQueryMsg::AssetBalanceAt { asset_info: AssetInfo::NativeToken { - denom: "ttt".to_owned(), + denom: "ttt".to_string(), }, block_height: app.block_info().height.into(), }, @@ -382,7 +400,7 @@ mod tests { let swap_msg = PairExecuteMsg::Swap { offer_asset: Asset { info: AssetInfo::NativeToken { - denom: "inj".to_owned(), + denom: "inj".to_string(), }, amount: Uint128::new(1_000000), }, @@ -392,7 +410,7 @@ mod tests { to: None, }; let send_funds = vec![Coin { - denom: "inj".to_owned(), + denom: "inj".to_string(), amount: Uint128::new(1_000000), }]; app.execute_contract(owner.clone(), pair_contract.addr(), &swap_msg, &send_funds) @@ -405,7 +423,7 @@ mod tests { pair_contract.addr(), &PairQueryMsg::AssetBalanceAt { asset_info: AssetInfo::NativeToken { - denom: "ttt".to_owned(), + denom: "ttt".to_string(), }, block_height: app.block_info().height.into(), }, @@ -431,7 +449,7 @@ mod tests { pair_address: pair_contract.addr().into_string(), }; let send_funds = vec![Coin { - denom: "inj".to_owned(), + denom: "inj".to_string(), amount: Uint128::new(1_000000), }]; app.execute_contract(owner.clone(), my_contract.addr(), &my_swap_msg, &send_funds) @@ -511,6 +529,18 @@ mod tests { .unwrap() .amount > 0u128.into() - ) + ); + + // if wrong pair, return error + let my_swap_msg = crate::msg::ExecuteMsg::Astro { + pair_address: pair_contract.addr().into_string(), + }; + let send_funds = vec![Coin { + denom: "abc".to_owned(), + amount: Uint128::new(1_000000), + }]; + assert!(app + .execute_contract(owner, pair_contract.addr(), &my_swap_msg, &send_funds) + .is_err()) } } diff --git a/src/msg.rs b/src/msg.rs index 8603f8e..2ffa271 100644 --- a/src/msg.rs +++ b/src/msg.rs @@ -5,19 +5,34 @@ pub struct InstantiateMsg {} #[cw_serde] pub enum ExecuteMsg { - Astro { pair_address: String }, + Astro { + pair_address: String, + }, + AddSupportedPool { + pool_address: String, + token_1: String, + token_2: String, + }, } #[cw_serde] #[derive(QueryResponses)] pub enum QueryMsg { // GetCount returns the current count as a json-encoded number - // #[returns(GetCountResponse)] - // GetCount {}, + #[returns(GetPairResponse)] + GetPair { pool_address: String }, + #[returns(GetPoolAddrResponse)] + GetPoolAddr { token_1: String, token_2: String }, } -// We define a custom struct for each query response -// #[cw_serde] -// pub struct GetCountResponse { -// pub count: i32, -// } +//We define a custom struct for each query response +#[cw_serde] +pub struct GetPairResponse { + pub token_1: String, + pub token_2: String, +} + +#[cw_serde] +pub struct GetPoolAddrResponse { + pub pool_addresses: Vec, +} diff --git a/src/state.rs b/src/state.rs index 48cf02f..ee20f82 100644 --- a/src/state.rs +++ b/src/state.rs @@ -2,11 +2,42 @@ use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use cosmwasm_std::Addr; -use cw_storage_plus::Item; +use cw_storage_plus::{Index, IndexList, IndexedMap, Item, MultiIndex}; #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)] pub struct State { pub owner: Addr, } - +pub struct InfoIndexes<'a> { + pub address: MultiIndex<'a, String, PoolInfo, String>, + pub pair: MultiIndex<'a, (String, String), PoolInfo, String>, +} +impl<'a> IndexList for InfoIndexes<'a> { + fn get_indexes(&'_ self) -> Box> + '_> { + let v: Vec<&dyn Index> = vec![&self.address, &self.pair]; + Box::new(v.into_iter()) + } +} +pub const fn infos<'a>() -> IndexedMap<'a, &'a str, PoolInfo, InfoIndexes<'a>> { + let indexes = InfoIndexes { + address: MultiIndex::new( + |_pk: &[u8], d: &PoolInfo| d.address.to_string(), + "infos", + "infos__address", + ), + pair: MultiIndex::new( + |_pk: &[u8], d: &PoolInfo| (d.token_1.clone(), d.token_2.clone()), + "infos", + "infos__pair", + ), + }; + IndexedMap::new("infos", indexes) +} +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)] +pub struct PoolInfo { + pub address: Addr, + pub token_1: String, + pub token_2: String, +} pub const STATE: Item = Item::new("state"); +pub const POOL_INFO: IndexedMap<&str, PoolInfo, InfoIndexes> = infos();