test: delegator_claim_rewards_after_change_balance_flow_test (#75)

Author: noa-starkware

src/flow_test/flow_ideas.md

@@ -12,14 +12,9 @@
 - update rewards for 2 different blocks in the same epoch - should be same rewards.
 - Change epoch len in blocks - rewards should be changed.
 - disable_rewards = true, advance block, disable_rewards = false, advance block, disable_rewards = true, test rewards.
-- update rewards for some different blocks in the same epoch, test rewards of pool member.
 - staker change balance, attest, change balance, attest, set_v3, change balance, update_rewards, change_balance, update_rewards, test rewards.
 - with member from previous versions.
 
-## k=1 -> k=2 balances
-- delegator claim rewards when last change is in epoch + 2, then advance epochs and claim again to see no missing rewards
-- same as above, also when there is change in epoch + 1
-
 ## k=1 -> k=2 Migration Member
 - find sigma: Enter V0, change in V1, catch all ifs.
 more ideas:

src/flow_test/test.cairo

@@ -24,7 +24,8 @@ use staking::test_utils::constants::{
 use staking::test_utils::{
     StakingInitConfig, calculate_staker_btc_pool_rewards_v2,
     calculate_staker_strk_rewards_with_balances_v2, calculate_staker_strk_rewards_with_balances_v3,
-    calculate_strk_pool_rewards_with_pool_balance_v2, compute_rewards_per_unit,
+    calculate_strk_pool_rewards_with_pool_balance_v2,
+    calculate_strk_pool_rewards_with_pool_balance_v3, compute_rewards_per_unit,
     custom_decimals_token, deploy_mock_erc20_decimals_contract,
 };
 use starkware_utils::errors::{Describable, ErrorDisplay};
@@ -3571,3 +3572,119 @@ fn delegate_claim_after_claim_with_rewards_flow_test() {
     system.advance_epoch();
     assert!(system.delegator_claim_rewards(:delegator, :pool).is_zero());
 }
+
+
+/// Flow:
+/// Stake + delegate.
+/// Advance K epochs.
+/// Attest.
+/// Advance epoch.
+/// Increase delegate.
+/// Attest.
+/// Delegator claim rewards.
+/// Advance epoch and attest.
+/// Delegator claim rewards.
+/// Advance epoch and attest.
+/// Delegator claim rewards.
+/// Advance epoch and attest.
+/// Delegator claim rewards.
+/// Start consensus rewards.
+/// Increase delegate.
+/// update_rewards.
+/// Advance block and update_rewards.
+/// Delegator claim rewards.
+/// Advance epoch.
+/// Delegator claim rewards.
+/// update_rewards.
+/// Advance block and update_rewards.
+/// Advance epoch.
+/// update_rewards.
+/// Delegator claim rewards.
+/// Advance block and update_rewards.
+/// Advance epoch.
+/// Delegator claim rewards.
+#[test]
+fn delegator_claim_rewards_after_change_balance_flow_test() {
+    let cfg: StakingInitConfig = Default::default();
+    let mut system = SystemConfigTrait::basic_stake_flow_cfg(:cfg).deploy();
+    let min_stake = system.staking.get_min_stake();
+    let stake_amount = min_stake * 2;
+    let delegation_amount = stake_amount / 2;
+    let staker = system.new_staker(amount: stake_amount);
+    let commission = 200;
+    system.stake(:staker, amount: stake_amount, pool_enabled: true, :commission);
+    let pool = system.staking.get_pool(:staker);
+    let delegator = system.new_delegator(amount: 3 * delegation_amount);
+    system.delegate(:delegator, :pool, amount: delegation_amount);
+    system.advance_k_epochs();
+
+    let mut pool_balance = delegation_amount;
+    system.advance_block_custom_and_attest(:staker, stake: stake_amount + pool_balance);
+    let expected_rewards = calculate_strk_pool_rewards_with_pool_balance_v2(
+        staker_address: staker.staker.address,
+        staking_contract: system.staking.address,
+        minting_curve_contract: system.minting_curve.address,
+        :pool_balance,
+    );
+    system.advance_epoch();
+    system.add_to_delegation_pool(:delegator, :pool, amount: delegation_amount);
+    system.advance_block_custom_and_attest(:staker, stake: stake_amount + pool_balance);
+    let rewards = system.delegator_claim_rewards(:delegator, :pool);
+    assert!(rewards == expected_rewards);
+    system.advance_epoch();
+    system.advance_block_custom_and_attest(:staker, stake: stake_amount + pool_balance);
+    let rewards = system.delegator_claim_rewards(:delegator, :pool);
+    assert!(rewards == expected_rewards);
+    system.advance_epoch();
+    pool_balance += delegation_amount;
+    system.advance_block_custom_and_attest(:staker, stake: stake_amount + pool_balance);
+    let rewards = system.delegator_claim_rewards(:delegator, :pool);
+    assert!(rewards == expected_rewards);
+    system.advance_epoch();
+    let expected_rewards = calculate_strk_pool_rewards_with_pool_balance_v2(
+        staker_address: staker.staker.address,
+        staking_contract: system.staking.address,
+        minting_curve_contract: system.minting_curve.address,
+        :pool_balance,
+    );
+    system.advance_block_custom_and_attest(:staker, stake: stake_amount + pool_balance);
+    let rewards = system.delegator_claim_rewards(:delegator, :pool);
+    assert!(rewards == expected_rewards);
+
+    system.start_consensus_rewards();
+    system.add_to_delegation_pool(:delegator, :pool, amount: delegation_amount);
+    system.update_rewards(:staker, disable_rewards: false);
+    let expected_rewards_v3 = calculate_strk_pool_rewards_with_pool_balance_v3(
+        staker_address: staker.staker.address,
+        staking_contract: system.staking.address,
+        minting_curve_contract: system.minting_curve.address,
+        :pool_balance,
+    );
+    advance_block_number_global(blocks: 1);
+    system.update_rewards(:staker, disable_rewards: false);
+    let rewards = system.delegator_claim_rewards(:delegator, :pool);
+    assert!(rewards == expected_rewards);
+    system.advance_epoch();
+    let rewards = system.delegator_claim_rewards(:delegator, :pool);
+    assert!(rewards == expected_rewards_v3 * 2);
+    system.update_rewards(:staker, disable_rewards: false);
+    advance_block_number_global(blocks: 1);
+    system.update_rewards(:staker, disable_rewards: false);
+    system.advance_epoch();
+    pool_balance += delegation_amount;
+    system.update_rewards(:staker, disable_rewards: false);
+    let rewards = system.delegator_claim_rewards(:delegator, :pool);
+    assert!(rewards == expected_rewards_v3 * 2);
+    advance_block_number_global(blocks: 1);
+    system.update_rewards(:staker, disable_rewards: false);
+    system.advance_epoch();
+    let expected_rewards_v3 = calculate_strk_pool_rewards_with_pool_balance_v3(
+        staker_address: staker.staker.address,
+        staking_contract: system.staking.address,
+        minting_curve_contract: system.minting_curve.address,
+        :pool_balance,
+    );
+    let rewards = system.delegator_claim_rewards(:delegator, :pool);
+    // Allow small rounding errors.
+    assert!(wide_abs_diff(rewards, expected_rewards_v3 * 2) <= 1);
+}

src/test_utils.cairo

@@ -1172,6 +1172,25 @@ pub(crate) fn calculate_strk_pool_rewards_v3(
     pool_rewards
 }
 
+/// Calculate strk pool rewards for one epoch for the given pool balance.
+pub(crate) fn calculate_strk_pool_rewards_with_pool_balance_v3(
+    staker_address: ContractAddress,
+    staking_contract: ContractAddress,
+    minting_curve_contract: ContractAddress,
+    pool_balance: Amount,
+) -> Amount {
+    let staking_dispatcher = IStakingDispatcher { contract_address: staking_contract };
+    let staker_info = staking_dispatcher.staker_info_v1(:staker_address);
+    let (_, pool_rewards) = calculate_staker_strk_rewards_with_balances_v3(
+        amount_own: staker_info.amount_own,
+        pool_amount: pool_balance,
+        commission: staker_info.pool_info.unwrap().commission,
+        :staking_contract,
+        :minting_curve_contract,
+    );
+    pool_rewards
+}
+
 /// Returns (staker_commission_rewards, BTC_pool_rewards) for the specified
 /// `normalized_pool_balance`, `normalized_staker_total_btc_balance` and `commission`.
 ///