From 9c0cf1b9549c424ff840a4d649eb1d205dd3db6c Mon Sep 17 00:00:00 2001
From: Mark Watney <mark@marsprotocol.foundation>
Date: Fri, 21 Jun 2024 19:06:56 +0800
Subject: [PATCH] improvements

---
 .../tests/test_claim_astro_lp_rewards.rs      | 108 ++++--------------
 1 file changed, 21 insertions(+), 87 deletions(-)

diff --git a/contracts/incentives/tests/tests/test_claim_astro_lp_rewards.rs b/contracts/incentives/tests/tests/test_claim_astro_lp_rewards.rs
index e8e235c3..ee408e83 100644
--- a/contracts/incentives/tests/tests/test_claim_astro_lp_rewards.rs
+++ b/contracts/incentives/tests/tests/test_claim_astro_lp_rewards.rs
@@ -6,7 +6,10 @@ use cosmwasm_std::{
 use mars_incentives::{contract::execute, query, state::ASTRO_TOTAL_LP_DEPOSITS, ContractError};
 use mars_testing::{assert_eq_vec, MarsMockQuerier};
 use mars_types::{
-    credit_manager::{ActionAmount, ActionCoin},
+    credit_manager::{
+        ActionAmount::{self, Exact},
+        ActionCoin,
+    },
     incentives::ExecuteMsg,
 };
 
@@ -436,7 +439,7 @@ fn assert_only_credit_manager() {
 }
 
 #[test]
-fn unstake_correctly_updates_lp_states() {
+fn lp_states_update_correctly() {
     // SETUP
     let env = mock_env();
     let mut deps: OwnedDeps<MemoryStorage, MockApi, MarsMockQuerier> = th_setup();
@@ -478,8 +481,8 @@ fn unstake_correctly_updates_lp_states() {
     // User a = 100
     // User b = 200
     assert_eq!(
-        ASTRO_TOTAL_LP_DEPOSITS.may_load(&deps.storage, lp_denom).unwrap(),
-        Some(Uint128::new(300u128))
+        ASTRO_TOTAL_LP_DEPOSITS.load(&deps.storage, lp_denom).unwrap(),
+        Uint128::new(300u128)
     );
     assert_eq!(
         query::query_staked_astro_lp_position(
@@ -513,7 +516,7 @@ fn unstake_correctly_updates_lp_states() {
         user_a_id.to_string(),
         ActionCoin {
             denom: lp_denom.to_string(),
-            amount: mars_types::credit_manager::ActionAmount::Exact(Uint128::new(50u128)),
+            amount: Exact(Uint128::new(50u128)),
         },
     )
     .unwrap();
@@ -522,8 +525,8 @@ fn unstake_correctly_updates_lp_states() {
     // User a = 50
     // User b = 200
     assert_eq!(
-        ASTRO_TOTAL_LP_DEPOSITS.may_load(&deps.storage, lp_denom).unwrap(),
-        Some(Uint128::new(250u128))
+        ASTRO_TOTAL_LP_DEPOSITS.load(&deps.storage, lp_denom).unwrap(),
+        Uint128::new(250u128)
     );
     assert_eq!(
         query::query_staked_astro_lp_position(
@@ -566,8 +569,8 @@ fn unstake_correctly_updates_lp_states() {
     // User a = 50
     // User b = 0
     assert_eq!(
-        ASTRO_TOTAL_LP_DEPOSITS.may_load(&deps.storage, lp_denom).unwrap(),
-        Some(Uint128::new(50u128))
+        ASTRO_TOTAL_LP_DEPOSITS.load(&deps.storage, lp_denom).unwrap(),
+        Uint128::new(50u128)
     );
     assert_eq!(
         query::query_staked_astro_lp_position(
@@ -593,91 +596,22 @@ fn unstake_correctly_updates_lp_states() {
         .amount,
         Uint128::new(0u128)
     );
-}
-
-#[test]
-fn stake_correctly_updates_lp_states() {
-    // SETUP
-    let env = mock_env();
-    let mut deps: OwnedDeps<MemoryStorage, MockApi, MarsMockQuerier> = th_setup();
-
-    // Users
-    let user_a_id = "1";
-    let user_b_id = "2";
-
-    let credit_manager = Addr::unchecked("credit_manager");
-    let astroport_incentives_addr = Addr::unchecked("astroport_incentives");
-    deps.querier.set_astroport_incentives_address(astroport_incentives_addr.clone());
-
-    let lp_denom = "uusd/ubtc";
-
-    // State:
-    // - LP in incentives = 0
-    // - Rewards available = 0
-    assert_eq!(ASTRO_TOTAL_LP_DEPOSITS.may_load(&deps.storage, lp_denom).unwrap(), None);
-
-    deposit_for_user(
-        deps.as_mut(),
-        env.clone(),
-        credit_manager.as_str(),
-        user_a_id.to_string(),
-        Coin::new(100u128, lp_denom),
-    ).unwrap();
 
     deposit_for_user(
         deps.as_mut(),
         env.clone(),
         credit_manager.as_str(),
         user_b_id.to_string(),
-        Coin::new(200u128, lp_denom),
-    ).unwrap();
-
-    // LP in incentives = 300
-    // User a = 100
-    // User b = 200
-    assert_eq!(
-        ASTRO_TOTAL_LP_DEPOSITS.may_load(&deps.storage, lp_denom).unwrap(),
-        Some(Uint128::new(300u128))
-    );
-    assert_eq!(
-        query::query_staked_astro_lp_position(
-            deps.as_ref(),
-            env.clone(),
-            user_a_id.to_string(),
-            lp_denom.to_string()
-        )
-        .unwrap()
-        .lp_coin
-        .amount,
-        Uint128::new(100u128)
-    );
-    assert_eq!(
-        query::query_staked_astro_lp_position(
-            deps.as_ref(),
-            env.clone(),
-            user_b_id.to_string(),
-            lp_denom.to_string()
-        )
-        .unwrap()
-        .lp_coin
-        .amount,
-        Uint128::new(200u128)
-    );
-
-    deposit_for_user(
-        deps.as_mut(),
-        env.clone(),
-        credit_manager.as_str(),
-        user_a_id.to_string(),
         Coin::new(50u128, lp_denom),
-    ).unwrap();
+    )
+    .unwrap();
 
-    // LP in incentives = 350
-    // User a = 150
-    // User b = 200
+    // LP in incentives = 100
+    // User a = 50
+    // User b = 50
     assert_eq!(
-        ASTRO_TOTAL_LP_DEPOSITS.may_load(&deps.storage, lp_denom).unwrap(),
-        Some(Uint128::new(350u128))
+        ASTRO_TOTAL_LP_DEPOSITS.load(&deps.storage, lp_denom).unwrap(),
+        Uint128::new(100u128)
     );
     assert_eq!(
         query::query_staked_astro_lp_position(
@@ -689,7 +623,7 @@ fn stake_correctly_updates_lp_states() {
         .unwrap()
         .lp_coin
         .amount,
-        Uint128::new(150u128)
+        Uint128::new(50u128)
     );
     assert_eq!(
         query::query_staked_astro_lp_position(
@@ -701,6 +635,6 @@ fn stake_correctly_updates_lp_states() {
         .unwrap()
         .lp_coin
         .amount,
-        Uint128::new(200u128)
+        Uint128::new(50u128)
     );
 }