diff --git a/Cargo.lock b/Cargo.lock index 16c6af7..f9dd66c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1972,6 +1972,7 @@ dependencies = [ "cw-dex", "cw-it", "cw-ownable", + "cw20 0.15.1", "locked-astroport-vault", "locked-astroport-vault-test-helpers", "neutron-astroport-reward-distributor", diff --git a/contracts/reward-distributor/src/contract.rs b/contracts/reward-distributor/src/contract.rs index c253af6..2985bf1 100644 --- a/contracts/reward-distributor/src/contract.rs +++ b/contracts/reward-distributor/src/contract.rs @@ -7,7 +7,7 @@ use cw_dex::astroport::AstroportPool; use cw_vault_standard::{VaultContract, VaultContractUnchecked}; use neutron_astroport_reward_distributor::{ Config, ConfigUnchecked, ContractError, ExecuteMsg, InstantiateMsg, InternalMsg, QueryMsg, - StateResponse, CONFIG, LAST_DISTRIBUTED, REWARD_POOL, REWARD_VAULT, + RewardInfo, RewardType, StateResponse, CONFIG, LAST_DISTRIBUTED, REWARD_TOKEN, }; use crate::execute; @@ -25,24 +25,41 @@ pub fn instantiate( cw2::set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?; cw_ownable::initialize_owner(deps.storage, deps.api, Some(&msg.owner))?; - let reward_vault: VaultContract = - VaultContractUnchecked::new(&msg.reward_vault_addr).check(deps.as_ref())?; + let reward_token = match msg.reward_token_info { + RewardInfo::VaultAddr(reward_vault_addr) => { + let reward_vault: VaultContract = + VaultContractUnchecked::new(&reward_vault_addr).check(deps.as_ref())?; - // Validate reward vault base token as CW20 Astroport LP token - let reward_lp_token = deps - .api - .addr_validate(&reward_vault.base_token) - .map_err(|_| StdError::generic_err("Invalid base token of reward vault"))?; + // Validate reward vault base token as CW20 Astroport LP token + let reward_lp_token = deps + .api + .addr_validate(&reward_vault.base_token) + .map_err(|_| StdError::generic_err("Invalid base token of reward vault"))?; - // Query minter of LP token to get reward pool address - let minter_res: MinterResponse = deps.querier.query(&QueryRequest::Wasm(WasmQuery::Smart { - contract_addr: reward_lp_token.to_string(), - msg: to_binary(&Cw20QueryMsg::Minter {})?, - }))?; - let reward_pool_addr = deps.api.addr_validate(&minter_res.minter)?; + // Query minter of LP token to get reward pool address + let minter_res: MinterResponse = + deps.querier.query(&QueryRequest::Wasm(WasmQuery::Smart { + contract_addr: reward_lp_token.to_string(), + msg: to_binary(&Cw20QueryMsg::Minter {})?, + }))?; + let reward_pool_addr = deps.api.addr_validate(&minter_res.minter)?; - // Query reward pool for pool info to create pool object - let reward_pool = AstroportPool::new(deps.as_ref(), reward_pool_addr)?; + // Query reward pool for pool info to create pool object + let reward_pool = AstroportPool::new(deps.as_ref(), reward_pool_addr)?; + + RewardType::Vault { + vault: reward_vault, + pool: reward_pool, + } + } + RewardInfo::AstroportPoolAddr(pool_addr) => { + let reward_pool = + AstroportPool::new(deps.as_ref(), deps.api.addr_validate(&pool_addr)?)?; + + RewardType::LP(reward_pool) + } + RewardInfo::NativeCoin(reward_coin_denom) => RewardType::Coin(reward_coin_denom), + }; // Create config let config: Config = ConfigUnchecked { @@ -54,8 +71,7 @@ pub fn instantiate( CONFIG.save(deps.storage, &config)?; LAST_DISTRIBUTED.save(deps.storage, &env.block.time.seconds())?; - REWARD_POOL.save(deps.storage, &reward_pool)?; - REWARD_VAULT.save(deps.storage, &reward_vault)?; + REWARD_TOKEN.save(deps.storage, &reward_token)?; Ok(Response::default()) } @@ -103,14 +119,11 @@ pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult { } QueryMsg::State {} => { let config = CONFIG.load(deps.storage)?; - let reward_pool = REWARD_POOL.load(deps.storage)?; - let reward_vault = REWARD_VAULT.load(deps.storage)?; let last_distributed = LAST_DISTRIBUTED.load(deps.storage)?; to_binary(&StateResponse { config, - reward_pool, - reward_vault, + reward_token: REWARD_TOKEN.load(deps.storage)?, last_distributed, }) } diff --git a/contracts/reward-distributor/src/execute.rs b/contracts/reward-distributor/src/execute.rs index 15577ae..207cac5 100644 --- a/contracts/reward-distributor/src/execute.rs +++ b/contracts/reward-distributor/src/execute.rs @@ -1,8 +1,10 @@ use apollo_cw_asset::{Asset, AssetInfo, AssetList}; -use cosmwasm_std::{Deps, DepsMut, Env, Event, MessageInfo, Response, Uint128}; +use cosmwasm_std::{ + coins, BankMsg, CosmosMsg, Deps, DepsMut, Env, Event, MessageInfo, Response, Uint128, +}; use cw_dex::traits::Pool as PoolTrait; use neutron_astroport_reward_distributor::{ - ConfigUpdates, ContractError, InternalMsg, CONFIG, LAST_DISTRIBUTED, REWARD_POOL, REWARD_VAULT, + ConfigUpdates, ContractError, InternalMsg, RewardType, CONFIG, LAST_DISTRIBUTED, REWARD_TOKEN, }; pub fn execute_distribute(deps: DepsMut, env: Env) -> Result { @@ -22,55 +24,77 @@ pub fn execute_distribute(deps: DepsMut, env: Env) -> Result { + // Query the vault to see how many base tokens would be returned after + // redeeming. If zero we return Ok, so that update_config does not fail when + // trying to distribute. + let base_token_amount = vault.query_convert_to_assets(&deps.querier, reward_amount)?; + if base_token_amount.is_zero() { + return Ok(Response::new()); + } + + // Check contract's balance of vault tokens and error if not enough. This is + // just so we get a clearer error message rather than the confusing "cannot + // sub 0 with x". + let vault_token_balance = deps + .querier + .query_balance(&env.contract.address, &vault.vault_token)?; + if vault_token_balance.amount < reward_amount { + return Err(ContractError::InsufficientVaultTokenBalance { + vault_token_balance: vault_token_balance.amount, + redeem_amount: reward_amount, + }); + } + + // Redeem rewards from the vault + let redeem_msg = vault.redeem(reward_amount, None)?; + + // Create internal callback msg + let callback_msg = InternalMsg::VaultTokensRedeemed {}.into_cosmos_msg(&env)?; + + res = res.add_message(redeem_msg).add_message(callback_msg); + } + RewardType::LP(pool) => { + // Create message to withdraw liquidity from pool + let lp_tokens = Asset::new(AssetInfo::Cw20(pool.lp_token_addr.clone()), reward_amount); + res = pool.withdraw_liquidity(deps.as_ref(), &env, lp_tokens, AssetList::new())?; + + // Create internal callback msg + let callback_msg = InternalMsg::LpRedeemed {}.into_cosmos_msg(&env)?; + res = res.add_message(callback_msg); + } + RewardType::Coin(reward_coin_denom) => { + // Create message to send coins to distribution address + let send_msg: CosmosMsg = BankMsg::Send { + to_address: config.distribution_addr.to_string(), + amount: coins(reward_amount.u128(), reward_coin_denom), + } + .into(); + res = res.add_message(send_msg); + } } // Set last distributed time to current time LAST_DISTRIBUTED.save(deps.storage, ¤t_time)?; - // Redeem rewards from the vault - let redeem_msg = reward_vault.redeem(redeem_amount, None)?; - - // Create internal callback msg - let callback_msg = InternalMsg::VaultTokensRedeemed {}.into_cosmos_msg(&env)?; - let event = Event::new("apollo/neutron-astroport-reward-distributor/execute_distribute") - .add_attribute("vault_tokens_redeemed", redeem_amount); + .add_attribute("vault_tokens_redeemed", reward_amount); - Ok(Response::default() - .add_message(redeem_msg) - .add_message(callback_msg) - .add_event(event)) + Ok(res.add_event(event)) } pub fn execute_internal_vault_tokens_redeemed( deps: Deps, env: Env, ) -> Result { - let reward_pool = REWARD_POOL.load(deps.storage)?; + let reward_pool = REWARD_TOKEN.load(deps.storage)?.into_pool()?; // Query lp token balance let reward_lp_token = AssetInfo::Cw20(reward_pool.lp_token_addr.clone()); @@ -93,7 +117,7 @@ pub fn execute_internal_vault_tokens_redeemed( pub fn execute_internal_lp_redeemed(deps: Deps, env: Env) -> Result { let config = CONFIG.load(deps.storage)?; - let reward_pool = REWARD_POOL.load(deps.storage)?; + let reward_pool = REWARD_TOKEN.load(deps.storage)?.into_pool()?; // Query contracts balances of pool assets let pool_asset_balances: AssetList = AssetList::query_asset_info_balances( diff --git a/contracts/reward-distributor/tests/access_control.rs b/contracts/reward-distributor/tests/access_control.rs index 8aebc11..1175d21 100644 --- a/contracts/reward-distributor/tests/access_control.rs +++ b/contracts/reward-distributor/tests/access_control.rs @@ -9,7 +9,7 @@ use locked_astroport_vault_test_helpers::robot::LockedAstroportVaultRobot; use neutron_astroport_reward_distributor::{ConfigUpdates, ExecuteMsg, InternalMsg}; use neutron_astroport_reward_distributor_test_helpers as test_helpers; -use test_helpers::robot::RewardDistributorRobot; +use test_helpers::robot::{RewardDistributorRobot, TestRewardType}; use crate::common::{DEPS_PATH, UNOPTIMIZED_PATH}; @@ -30,6 +30,7 @@ fn update_ownership_can_only_be_called_by_admin() { DEPS_PATH, UNOPTIMIZED_PATH, treasury_addr.address(), + TestRewardType::VaultToken, &admin, emission_per_second, rewards_start_time, @@ -67,6 +68,7 @@ fn update_config_can_only_be_called_by_admin() { DEPS_PATH, UNOPTIMIZED_PATH, treasury_addr.address(), + TestRewardType::VaultToken, &admin, emission_per_second, rewards_start_time, @@ -103,6 +105,7 @@ fn internal_msg_can_only_be_called_by_contract() { DEPS_PATH, UNOPTIMIZED_PATH, treasury_addr.address(), + TestRewardType::VaultToken, &admin, emission_per_second, rewards_start_time, diff --git a/contracts/reward-distributor/tests/integration_test.rs b/contracts/reward-distributor/tests/integration_test.rs index 18fd3d1..3b24397 100644 --- a/contracts/reward-distributor/tests/integration_test.rs +++ b/contracts/reward-distributor/tests/integration_test.rs @@ -1,15 +1,16 @@ use common::get_test_runner; use cosmwasm_std::{coin, Uint128}; use cw_it::helpers::Unwrap; +use cw_it::robot::TestRobot; use cw_it::test_tube::Account; use cw_it::traits::CwItRunner; - use locked_astroport_vault::helpers::INITIAL_VAULT_TOKENS_PER_BASE_TOKEN; use locked_astroport_vault_test_helpers::cw_vault_standard_test_helpers::traits::CwVaultStandardRobot; use locked_astroport_vault_test_helpers::robot::LockedAstroportVaultRobot; +use neutron_astroport_reward_distributor::RewardType; use neutron_astroport_reward_distributor_test_helpers as test_helpers; -use test_helpers::robot::RewardDistributorRobot; +use test_helpers::robot::{RewardDistributorRobot, TestRewardType}; use crate::common::{DEPS_PATH, UNOPTIMIZED_PATH}; @@ -29,6 +30,7 @@ fn test_initialization() { DEPS_PATH, UNOPTIMIZED_PATH, treasury_addr.address(), + TestRewardType::VaultToken, &admin, 1000000u128, rewards_start_time, @@ -39,16 +41,10 @@ fn test_initialization() { let config = state.config; assert_eq!(config.emission_per_second, Uint128::from(1000000u128)); assert_eq!(config.distribution_addr, robot.distribution_acc.address()); - assert_eq!(state.reward_pool, robot.reward_pool); - assert_eq!( - state.reward_pool.lp_token_addr.to_string(), - robot.reward_vault_robot.base_token() - ); - assert_eq!(state.reward_vault.addr, robot.reward_vault_robot.vault_addr); - assert_eq!( - state.reward_vault.vault_token, - robot.reward_vault_robot.vault_token() - ); + assert!(matches!( + state.reward_token, + RewardType::Vault { vault, pool } if vault.addr == robot.reward_vault_robot.vault_addr && pool.lp_token_addr == robot.reward_vault_robot.base_token() + )); // Query ownership let ownership = robot.query_ownership(); @@ -70,6 +66,7 @@ fn distribute_errors_when_not_enough_vault_tokens_in_contract() { DEPS_PATH, UNOPTIMIZED_PATH, treasury_addr.address(), + TestRewardType::VaultToken, &admin, emission_per_second, rewards_start_time, @@ -90,7 +87,7 @@ fn distribute_errors_when_not_enough_vault_tokens_in_contract() { } #[test] -fn test_correct_distribute() { +fn test_vault_tokens_correct_distribute() { let owned_runner = get_test_runner(); let runner = owned_runner.as_ref(); let admin = RewardDistributorRobot::default_account(&runner); @@ -105,6 +102,7 @@ fn test_correct_distribute() { DEPS_PATH, UNOPTIMIZED_PATH, treasury_addr.address(), + TestRewardType::VaultToken, &admin, emission_per_second, rewards_start_time, @@ -150,6 +148,116 @@ fn test_correct_distribute() { ); } +#[test] +fn test_lp_tokens_correct_distribute() { + let owned_runner = get_test_runner(); + let runner = owned_runner.as_ref(); + let admin = RewardDistributorRobot::default_account(&runner); + let treasury_addr = runner.init_account(&[]).unwrap(); + let dependencies = LockedAstroportVaultRobot::instantiate_deps(&runner, &admin, DEPS_PATH); + let emission_per_second = 100_000u128; + let init_time = runner.query_block_time_nanos() / 1_000_000_000; + let rewards_start_time = init_time + 5; + let robot = RewardDistributorRobot::instantiate( + &runner, + &dependencies, + DEPS_PATH, + UNOPTIMIZED_PATH, + treasury_addr.address(), + TestRewardType::LpToken, + &admin, + emission_per_second, + rewards_start_time, + ); + + let vault_robot = &robot.reward_vault_robot; + + // Send base tokens to reward distributor + let base_token_balance = vault_robot.query_base_token_balance(admin.address()); + let deposit_amount = base_token_balance / Uint128::new(10); + let distributor_balance_before = robot + .deposit_to_distributor(deposit_amount, Unwrap::Ok, &admin) + .reward_vault_robot + .query_base_token_balance(&robot.reward_distributor_addr); + + // Distribute rewards and check balances + let time_elapsed = 1000u64; + robot + .assert_distribution_acc_balances_eq(&[]) + .distribute(Unwrap::Ok, &admin) + .increase_time(5) // Rewards have started + .assert_distribution_acc_balances_eq(&[]) + .increase_time(time_elapsed) + .distribute(Unwrap::Ok, &admin) + .assert_distribution_acc_balances_eq(&[ + coin(emission_per_second * time_elapsed as u128, "uaxl"), + coin(emission_per_second * time_elapsed as u128, "untrn"), + ]); + + // Vault token balance of reward distributor should have decreased with the + // amount distributed + vault_robot.assert_base_token_balance_eq( + robot.reward_distributor_addr, + distributor_balance_before.u128() - emission_per_second * time_elapsed as u128, + ); +} + +#[test] +fn test_native_token_correct_distribute() { + let owned_runner = get_test_runner(); + let runner = owned_runner.as_ref(); + let admin = RewardDistributorRobot::default_account(&runner); + let treasury_addr = runner.init_account(&[]).unwrap(); + let dependencies = LockedAstroportVaultRobot::instantiate_deps(&runner, &admin, DEPS_PATH); + let emission_per_second = 100_000u128; + let init_time = runner.query_block_time_nanos() / 1_000_000_000; + let rewards_start_time = init_time + 5; + let reward_token_denom = "uaxl"; + let robot = RewardDistributorRobot::instantiate( + &runner, + &dependencies, + DEPS_PATH, + UNOPTIMIZED_PATH, + treasury_addr.address(), + TestRewardType::NativeCoin(reward_token_denom.to_string()), + &admin, + emission_per_second, + rewards_start_time, + ); + + let vault_robot = &robot.reward_vault_robot; + + // Send native tokens to reward distributor + let reward_token_balance = vault_robot.query_base_token_balance(admin.address()); + let deposit_amount = reward_token_balance / Uint128::new(10); + let distributor_balance_before = robot + .deposit_to_distributor(deposit_amount, Unwrap::Ok, &admin) + .reward_vault_robot + .query_native_token_balance(&robot.reward_distributor_addr, reward_token_denom); + + // Distribute rewards and check balances + let time_elapsed = 1000u64; + robot + .assert_distribution_acc_balances_eq(&[]) + .distribute(Unwrap::Ok, &admin) + .increase_time(5) // Rewards have started + .assert_distribution_acc_balances_eq(&[]) + .increase_time(time_elapsed) + .distribute(Unwrap::Ok, &admin) + .assert_distribution_acc_balances_eq(&[coin( + emission_per_second * time_elapsed as u128, + reward_token_denom, + )]); + + // Vault token balance of reward distributor should have decreased with the + // amount distributed + vault_robot.assert_native_token_balance_eq( + robot.reward_distributor_addr, + reward_token_denom, + distributor_balance_before.u128() - emission_per_second * time_elapsed as u128, + ); +} + #[test] fn distribute_does_not_error_when_distributed_vault_token_amount_would_give_zero_base_tokens() { let owned_runner = get_test_runner(); @@ -166,6 +274,7 @@ fn distribute_does_not_error_when_distributed_vault_token_amount_would_give_zero DEPS_PATH, UNOPTIMIZED_PATH, treasury_addr.address(), + TestRewardType::VaultToken, &admin, emission_per_second, rewards_start_time, diff --git a/contracts/reward-distributor/tests/test_config.rs b/contracts/reward-distributor/tests/test_config.rs index 624df90..a744e05 100644 --- a/contracts/reward-distributor/tests/test_config.rs +++ b/contracts/reward-distributor/tests/test_config.rs @@ -9,7 +9,7 @@ use locked_astroport_vault_test_helpers::robot::LockedAstroportVaultRobot; use neutron_astroport_reward_distributor::{Config, ConfigUpdates}; use neutron_astroport_reward_distributor_test_helpers as test_helpers; -use test_helpers::robot::RewardDistributorRobot; +use test_helpers::robot::{RewardDistributorRobot, TestRewardType}; use crate::common::{DEPS_PATH, UNOPTIMIZED_PATH}; @@ -31,6 +31,7 @@ fn update_config_works_correctly() { DEPS_PATH, UNOPTIMIZED_PATH, treasury_addr.address(), + TestRewardType::VaultToken, &admin, emission_per_second, rewards_start_time, @@ -68,6 +69,7 @@ fn update_config_distributes_rewards_if_emission_rate_or_start_time_changes() { DEPS_PATH, UNOPTIMIZED_PATH, treasury_addr.address(), + TestRewardType::VaultToken, &admin, emission_per_second, rewards_start_time, @@ -115,6 +117,7 @@ fn update_config_does_not_distribute_rewards_when_not_changing_emission_rate_or_ DEPS_PATH, UNOPTIMIZED_PATH, treasury_addr.address(), + TestRewardType::VaultToken, &admin, emission_per_second, rewards_start_time, diff --git a/packages/reward-distributor/src/msg.rs b/packages/reward-distributor/src/msg.rs index 35cbe2b..f92db65 100644 --- a/packages/reward-distributor/src/msg.rs +++ b/packages/reward-distributor/src/msg.rs @@ -1,10 +1,21 @@ use cosmwasm_schema::{cw_serde, QueryResponses}; use cosmwasm_std::{to_binary, CosmosMsg, Env, StdResult, Uint128, WasmMsg}; -use cw_dex::astroport::AstroportPool; use cw_ownable::{cw_ownable_execute, cw_ownable_query}; -use cw_vault_standard::VaultContract; -use crate::{Config, ConfigUpdates}; +use crate::{Config, ConfigUpdates, RewardType}; + +/// An enum for the information needed to instantiate the contract depending on +/// the type of reward token used. +#[cw_serde] +pub enum RewardInfo { + /// The address of the vault if the reward token is a vault token + VaultAddr(String), + /// The address of the Astroport pool if the reward token is an Astroport LP + /// token + AstroportPoolAddr(String), + /// The denom of the native coin if the reward token is a native coin + NativeCoin(String), +} #[cw_serde] pub struct InstantiateMsg { @@ -12,8 +23,9 @@ pub struct InstantiateMsg { pub owner: String, /// The emission rate per second pub emission_per_second: Uint128, - /// The address of the vault contract in which rewards are being held - pub reward_vault_addr: String, + /// The info needed to instantiate the contract depending on the type of + /// reward token used + pub reward_token_info: RewardInfo, /// The address that rewards are being distributed to pub distribution_addr: String, /// The unix timestamp at which rewards start being distributed @@ -67,7 +79,6 @@ pub enum QueryMsg { /// The response to a config query pub struct StateResponse { pub config: Config, - pub reward_pool: AstroportPool, - pub reward_vault: VaultContract, + pub reward_token: RewardType, pub last_distributed: u64, } diff --git a/packages/reward-distributor/src/state.rs b/packages/reward-distributor/src/state.rs index f5e6ebb..f2a5db9 100644 --- a/packages/reward-distributor/src/state.rs +++ b/packages/reward-distributor/src/state.rs @@ -1,17 +1,44 @@ +use cosmwasm_schema::cw_serde; +use cosmwasm_std::StdError; use cw_dex::astroport::AstroportPool; use cw_storage_plus::Item; use cw_vault_standard::VaultContract; use crate::config::Config; +/// An enum representing different types of reward tokens +#[cw_serde] +pub enum RewardType { + /// The reward token is a vault token + Vault { + /// The vault contract + vault: VaultContract, + /// The Astroport pool that the vault holds liquidity in + pool: AstroportPool, + }, + /// The reward token is an Astroport LP token + LP(AstroportPool), + /// The reward token is a native coin + Coin(String), +} + +impl RewardType { + pub fn into_pool(self) -> Result { + match self { + RewardType::Vault { vault: _, pool } => Ok(pool), + RewardType::LP(pool) => Ok(pool), + RewardType::Coin(_) => Err(StdError::generic_err( + "Cannot redeem vault tokens from coin reward", + )), + } + } +} + /// Stores the contract's config pub const CONFIG: Item = Item::new("config"); -/// Stores the Astroport pool in which rewards are being held -pub const REWARD_POOL: Item = Item::new("reward_pool"); - -/// Stores the vault contract in which rewards are being held -pub const REWARD_VAULT: Item = Item::new("reward_vault"); +/// Stores the reward token that this contract is distributing +pub const REWARD_TOKEN: Item = Item::new("reward_token"); /// Stores the last timestamp that rewards were distributed pub const LAST_DISTRIBUTED: Item = Item::new("last_distributed"); diff --git a/packages/test-helpers/Cargo.toml b/packages/test-helpers/Cargo.toml index 118a3ef..660417a 100644 --- a/packages/test-helpers/Cargo.toml +++ b/packages/test-helpers/Cargo.toml @@ -23,3 +23,4 @@ neutron-astroport-reward-distributor = { workspace = true } neutron-astroport-reward-distributor-contract = { workspace = true } cw-ownable = { workspace = true } locked-astroport-vault = { workspace = true } +cw20 = { workspace = true } diff --git a/packages/test-helpers/src/robot.rs b/packages/test-helpers/src/robot.rs index 26ee9dc..48dc0dc 100644 --- a/packages/test-helpers/src/robot.rs +++ b/packages/test-helpers/src/robot.rs @@ -1,6 +1,7 @@ use std::str::FromStr; use cosmwasm_std::{Addr, Coin, Coins, Decimal, Uint128}; +use cw20::Cw20ExecuteMsg; use cw_dex::astroport::AstroportPool; use cw_it::astroport::robot::AstroportTestRobot; use cw_it::astroport::utils::AstroportContracts; @@ -22,7 +23,7 @@ use neutron_astroport_reward_distributor::InstantiateMsg; #[cfg(feature = "osmosis-test-tube")] use cw_it::Artifact; -use reward_distributor::{Config, ConfigUpdates, QueryMsg}; +use reward_distributor::{Config, ConfigUpdates, QueryMsg, RewardInfo}; pub const REWARD_DISTRIBUTOR_WASM_NAME: &str = "neutron_astroport_reward_distributor_contract.wasm"; @@ -45,6 +46,15 @@ pub struct RewardDistributorRobot<'a> { pub distribution_acc: SigningAccount, pub reward_pool: AstroportPool, pub reward_vault_robot: LockedAstroportVaultRobot<'a>, + pub reward_type: TestRewardType, +} + +/// An enum representing different types of reward tokens +#[derive(Clone, Debug)] +pub enum TestRewardType { + VaultToken, + LpToken, + NativeCoin(String), } /// A trait with helper functions for testing the reward distributor contract. @@ -85,6 +95,7 @@ impl<'a> RewardDistributorRobot<'a> { dependency_artifacts_dir: &str, artifacts_dir: &str, vault_treasury_addr: String, + reward_type: TestRewardType, admin: &'a SigningAccount, emission_per_second: impl Into, rewards_start_time: u64, @@ -101,15 +112,25 @@ impl<'a> RewardDistributorRobot<'a> { admin, ); + let reward_token_info = match &reward_type { + TestRewardType::VaultToken => { + RewardInfo::VaultAddr(reward_vault_robot.vault_addr.clone()) + } + TestRewardType::LpToken => { + RewardInfo::AstroportPoolAddr(axl_ntrn_pool.pair_addr.to_string()) + } + TestRewardType::NativeCoin(denom) => RewardInfo::NativeCoin(denom.clone()), + }; + // Upload and instantiate reward distributor contract let code = Self::contract(runner, artifacts_dir); let code_id = runner.store_code(code, admin).unwrap(); let distribution_acc = runner.init_account(&[]).unwrap(); - let msg = InstantiateMsg { + let msg: InstantiateMsg = InstantiateMsg { distribution_addr: distribution_acc.address(), emission_per_second: emission_per_second.into(), owner: admin.address(), - reward_vault_addr: reward_vault_robot.vault_addr.clone(), + reward_token_info, rewards_start_time, }; let contract_addr = Wasm::new(runner) @@ -125,27 +146,50 @@ impl<'a> RewardDistributorRobot<'a> { distribution_acc, reward_pool: axl_ntrn_pool, reward_vault_robot, + reward_type, } } pub fn deposit_to_distributor( &self, - base_token_amount: Uint128, + amount: Uint128, unwrap_choice: Unwrap, signer: &SigningAccount, ) -> &Self { - self.reward_vault_robot - .deposit_cw20(base_token_amount, None, unwrap_choice, signer) - .assert_vault_token_balance_eq( - signer.address(), - base_token_amount * INITIAL_VAULT_TOKENS_PER_BASE_TOKEN, - ) - .send_native_tokens( - signer, - &self.reward_distributor_addr, - base_token_amount * INITIAL_VAULT_TOKENS_PER_BASE_TOKEN, - &self.reward_vault_robot.vault_token(), - ); + match &self.reward_type { + TestRewardType::VaultToken => { + // If the reward token is a vault token, we need to deposit base + // tokens to the vault and then deposit the vault tokens to the + // reward distributor. + self.reward_vault_robot + .deposit_cw20(amount, None, unwrap_choice, signer) + .assert_vault_token_balance_eq( + signer.address(), + amount * INITIAL_VAULT_TOKENS_PER_BASE_TOKEN, + ) + .send_native_tokens( + signer, + &self.reward_distributor_addr, + amount * INITIAL_VAULT_TOKENS_PER_BASE_TOKEN, + &self.reward_vault_robot.vault_token(), + ); + } + TestRewardType::LpToken => { + let msg = Cw20ExecuteMsg::Transfer { + recipient: self.reward_distributor_addr.clone(), + amount, + }; + unwrap_choice.unwrap(self.wasm().execute( + self.reward_pool.lp_token_addr.as_str(), + &msg, + &[], + signer, + )); + } + TestRewardType::NativeCoin(denom) => { + self.send_native_tokens(signer, &self.reward_distributor_addr, amount, denom); + } + } self }