Stake: Implement first test for the withdrawable rewards in chunks

Phoenix-Protocol-Group · Nov 22, 2024 · 65bde48 · 65bde48
1 parent 6b8b644
commit 65bde48
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 13 deletions.
diff --git a/contracts/stake/src/contract.rs b/contracts/stake/src/contract.rs
@@ -305,7 +305,7 @@ impl StakingTrait for Staking {
         let start_day = stakes.last_reward_time / crate::distribution::SECONDS_PER_DAY;
 
         // Calculate pending rewards in chunks
-        let pending_rewards =
+        let (pending_rewards, last_reward_day) =
             calculate_pending_rewards_chunked(&env, &reward_token, &stakes, start_day, chunk_size);
 
         // Transfer pending rewards to the user
@@ -318,8 +318,7 @@ impl StakingTrait for Staking {
         }
 
         // Explicitly document that chunk_size is inclusive of start_day
-        let last_processed_day = start_day + chunk_size as u64 - 1;
-        stakes.last_reward_time = last_processed_day * crate::distribution::SECONDS_PER_DAY;
+        stakes.last_reward_time = last_reward_day;
         save_stakes(&env, &sender, &stakes);
     }
 
@@ -418,7 +417,7 @@ impl StakingTrait for Staking {
 
         // Iterate over all distributions and calculate withdrawable rewards
         for asset in get_distributions(&env) {
-            let pending_reward =
+            let (pending_reward, _) =
                 calculate_pending_rewards_chunked(&env, &asset, &stakes, start_day, chunk_size);
 
             rewards.push_back(WithdrawableReward {

diff --git a/contracts/stake/src/distribution.rs b/contracts/stake/src/distribution.rs
@@ -127,7 +127,7 @@ pub fn calculate_pending_rewards_chunked(
     user_info: &BondingInfo,
     start_day: u64,
     chunk_size: u32,
-) -> i128 {
+) -> (i128, u64) {
     let current_timestamp = env.ledger().timestamp();
     let last_claim_time = user_info.last_reward_time;
 
@@ -143,6 +143,7 @@ pub fn calculate_pending_rewards_chunked(
     }
 
     let mut pending_rewards: i128 = 0;
+    let mut last_reward_day = 0u64;
 
     for reward_day in reward_keys
         .into_iter()
@@ -163,8 +164,9 @@ pub fn calculate_pending_rewards_chunked(
                     pending_rewards += (user_share as i128) * multiplier;
                 }
             }
+            last_reward_day = reward_day;
         }
     }
 
-    pending_rewards
+    (pending_rewards, last_reward_day)
 }
diff --git a/contracts/stake/src/tests/distribution.rs b/contracts/stake/src/tests/distribution.rs
@@ -1260,12 +1260,6 @@ fn add_distribution_and_distribute_reward_in_chunks() {
     staking.withdraw_rewards_chunks(&user, &reward_token.address, &chunk_size);
     assert_eq!(reward_token.balance(&user), 49_980);
 
-    // Ensure `last_reward_time` was updated correctly
-    assert_eq!(
-        staking.query_staked(&user).last_reward_time,
-        (chunk_size - 1) as u64 * SECONDS_PER_DAY
-    );
-
     // Query withdrawable rewards for the second chunk
     let withdrawable_chunk_2 = staking.query_withdrawable_rewards_ch(
         &user,
@@ -1299,7 +1293,13 @@ fn add_distribution_and_distribute_reward_in_chunks() {
     assert_eq!(
         withdrawable_chunk_3,
         WithdrawableRewardsResponse {
-            rewards: vec![&env]
+            rewards: vec![
+                &env,
+                WithdrawableReward {
+                    reward_address: reward_token.address.clone(),
+                    reward_amount: 0
+                }
+            ]
         }
     );
 }