test: Add tests for other reward token types

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,

contracts/reward-distributor/tests/integration_test.rs

@@ -1,6 +1,7 @@
 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;
@@ -9,7 +10,7 @@ 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,
@@ -64,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,
@@ -84,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);
@@ -99,6 +102,7 @@ fn test_correct_distribute() {
         DEPS_PATH,
         UNOPTIMIZED_PATH,
         treasury_addr.address(),
+        TestRewardType::VaultToken,
         &admin,
         emission_per_second,
         rewards_start_time,
@@ -144,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();
@@ -160,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,

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,

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 }

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;
@@ -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<Uint128>,
         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_token_info: RewardInfo::VaultAddr(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
     }