fix: issue with querying rewards for asset not staked

terra-money · Feb 8, 2024 · 9919f7d · 9919f7d
1 parent ec6898d
commit 9919f7d
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 4 deletions.
diff --git a/contracts/alliance-hub/src/contract.rs b/contracts/alliance-hub/src/contract.rs
@@ -39,7 +39,7 @@ const CREATE_REPLY_ID: u64 = 1;
 const CLAIM_REWARD_ERROR_REPLY_ID: u64 = 2;
 
 #[cfg_attr(not(feature = "library"), entry_point)]
-pub fn migrate(_deps: DepsMut, _env: Env, _msg: MigrateMsg) -> StdResult<Response> {
+pub fn migrate(deps: DepsMut, _env: Env, _msg: MigrateMsg) -> StdResult<Response> {
     Ok(Response::default())
 }
 
@@ -238,7 +238,7 @@ fn unstake(deps: DepsMut, info: MessageInfo, asset: Asset) -> Result<Response, C
                     }
                     Ok(balance - asset.amount)
                 }
-                None => Err(ContractError::InsufficientBalance {}),
+                None => Err(ContractError::AssetNotStaked {}),
             }
         },
     )?;
@@ -308,6 +308,14 @@ fn _claim_reward(
     asset: AssetInfo,
 ) -> Result<Uint128, ContractError> {
     let asset_key = AssetInfoKey::from(&asset);
+    // If the user did not stake the asset, do nothing and return 0
+    if BALANCES
+        .load(storage, (user.clone(), asset_key.clone()))
+        .is_err()
+    {
+        return Ok(Uint128::zero());
+    }
+
     let user_reward_rate = USER_ASSET_REWARD_RATE.load(storage, (user.clone(), asset_key.clone()));
     let asset_reward_rate = ASSET_REWARD_RATE.load(storage, asset_key.clone())?;
 

diff --git a/contracts/alliance-hub/src/query.rs b/contracts/alliance-hub/src/query.rs
@@ -79,7 +79,7 @@ fn get_pending_rewards(deps: Deps, asset_query: AssetQuery) -> StdResult<Binary>
     let user_reward_rate = USER_ASSET_REWARD_RATE.load(deps.storage, key.clone())?;
     let asset_reward_rate =
         ASSET_REWARD_RATE.load(deps.storage, AssetInfoKey::from(asset_query.asset.clone()))?;
-    let user_balance = BALANCES.load(deps.storage, key.clone())?;
+    let user_balance = BALANCES.load(deps.storage, key.clone()).unwrap_or(Uint128::zero());
     let unclaimed_rewards = UNCLAIMED_REWARDS
         .load(deps.storage, key)
         .unwrap_or(Uint128::zero());
@@ -131,7 +131,7 @@ fn get_all_pending_rewards(deps: Deps, query: AllPendingRewardsQuery) -> StdResu
             let user_balance = BALANCES.load(
                 deps.storage,
                 (addr.clone(), AssetInfoKey::from(asset.clone())),
-            )?;
+            ).unwrap_or(Uint128::zero());
             let unclaimed_rewards = UNCLAIMED_REWARDS
                 .load(
                     deps.storage,

diff --git a/packages/alliance-protocol/src/error.rs b/packages/alliance-protocol/src/error.rs
@@ -39,4 +39,7 @@ pub enum ContractError {
 
     #[error("Invalid total distribution: {0}")]
     InvalidTotalDistribution(Decimal),
+
+    #[error("Asset not staked")]
+    AssetNotStaked {},
 }