test: staker_change_balance_twice_same_epoch_flow_test (#77)

Author: noa-starkware

src/flow_test/flow_ideas.md

@@ -20,12 +20,10 @@
 - 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
 - delegate, advance epoch and get rewards for the pool, claim - zero rewards for the delegate
-- delegate, advance epoch, delegate, advance epoch, claim rewards - only for the first delegation, advance epoch, claim rewards - for all
 - delegator claim after claim
 - delegator claim after claim when exists checkpoint with the current epoch of the first claim
 - staker change balance in each epoch (increase, intent, delegate increase, delegate intent, delegate exit) and attest in each epoch - test rewards (also some epochs with no balance change and some epochs with no attest)
 - staker has multiple pool with multiple delegator each, change balance (staker, strk delegate, btc delegate) and attest in many epochs and test rewards both staker and members
-- test staker claim rewards with more than one balance change in an epoch.
 - member enter, in the same epoch or one epoch after claim (balance is zero).
 
 ## k=1 -> k=2 Migration Member

src/flow_test/test.cairo

@@ -3029,3 +3029,152 @@ fn enable_disable_btc_tokens_flow_test() {
         .span();
     assert!(tokens == expected_tokens);
 }
+
+/// Flow:
+/// Staker stake
+/// Advance K epochs
+/// Staker increase stake
+/// Staker increase stake again
+/// Staker attest and claim rewards - test rewards.
+/// Advance epoch
+/// Staker attest and claim rewards - test rewards.
+/// Advance epoch
+/// Staker attest and claim rewards - test rewards.
+/// Start consensus rewards
+/// Staker increase stake
+/// Staker increase stake again
+/// Update rewards and claim rewards - test rewards.
+/// Advance epoch
+/// Update rewards and claim rewards - test rewards.
+/// Advance epoch
+/// Update rewards and claim rewards - test rewards.
+#[test]
+fn staker_change_balance_twice_same_epoch_flow_test() {
+    let cfg: StakingInitConfig = Default::default();
+    let mut system = SystemConfigTrait::basic_stake_flow_cfg(:cfg).deploy();
+    let staking_contract = system.staking.address;
+    let minting_curve_contract = system.minting_curve.address;
+    let stake_amount = system.staking.get_min_stake();
+    let staker = system.new_staker(amount: stake_amount * 3);
+    let commission = 200;
+    system.stake(:staker, amount: stake_amount, pool_enabled: true, :commission);
+    system.advance_k_epochs();
+    system.increase_stake(:staker, amount: stake_amount / 2);
+    system.increase_stake(:staker, amount: stake_amount / 2);
+
+    system.advance_block_custom_and_attest(:staker, stake: stake_amount);
+    let (expected_rewards, _) = calculate_staker_strk_rewards_with_balances_v2(
+        amount_own: stake_amount,
+        pool_amount: Zero::zero(),
+        :commission,
+        :staking_contract,
+        :minting_curve_contract,
+    );
+    assert!(expected_rewards.is_non_zero());
+    assert!(system.staker_claim_rewards(:staker) == expected_rewards);
+
+    system.advance_epoch();
+    system.advance_block_custom_and_attest(:staker, stake: stake_amount);
+    assert!(system.staker_claim_rewards(:staker) == expected_rewards);
+
+    system.advance_epoch();
+    system.advance_block_custom_and_attest(:staker, stake: stake_amount * 2);
+    let (expected_rewards, _) = calculate_staker_strk_rewards_with_balances_v2(
+        amount_own: stake_amount * 2,
+        pool_amount: Zero::zero(),
+        :commission,
+        :staking_contract,
+        :minting_curve_contract,
+    );
+    assert!(expected_rewards.is_non_zero());
+    assert!(system.staker_claim_rewards(:staker) == expected_rewards);
+
+    system.start_consensus_rewards();
+    system.increase_stake(:staker, amount: stake_amount / 2);
+    system.increase_stake(:staker, amount: stake_amount / 2);
+
+    system.update_rewards(:staker, disable_rewards: false);
+    let (expected_rewards, _) = calculate_staker_strk_rewards_with_balances_v3(
+        amount_own: stake_amount * 2,
+        pool_amount: Zero::zero(),
+        :commission,
+        :staking_contract,
+        :minting_curve_contract,
+    );
+    assert!(expected_rewards.is_non_zero());
+    assert!(system.staker_claim_rewards(:staker) == expected_rewards);
+
+    system.advance_epoch();
+    system.update_rewards(:staker, disable_rewards: false);
+    assert!(system.staker_claim_rewards(:staker) == expected_rewards);
+
+    system.advance_epoch();
+    system.update_rewards(:staker, disable_rewards: false);
+    let (expected_rewards, _) = calculate_staker_strk_rewards_with_balances_v3(
+        amount_own: stake_amount * 3,
+        pool_amount: Zero::zero(),
+        :commission,
+        :staking_contract,
+        :minting_curve_contract,
+    );
+    assert!(expected_rewards.is_non_zero());
+    assert!(system.staker_claim_rewards(:staker) == expected_rewards);
+}
+
+/// Flow:
+/// Delegate
+/// Advance epoch
+/// Increase delegate
+/// Advance epoch
+/// Attest
+/// Advance epoch
+/// Claim rewards - only for the first delegation
+/// Attest
+/// Advance epoch
+/// Claim rewards - for all
+#[test]
+fn test_delegate_increase_delegate_flow_test() {
+    let cfg: StakingInitConfig = Default::default();
+    let mut system = SystemConfigTrait::basic_stake_flow_cfg(:cfg).deploy();
+    let staking_contract = system.staking.address;
+    let minting_curve_contract = system.minting_curve.address;
+    let stake_amount = system.staking.get_min_stake();
+    let staker = system.new_staker(amount: stake_amount);
+    let commission = 200;
+    system.stake(:staker, amount: stake_amount, pool_enabled: true, :commission);
+    system.advance_k_epochs();
+    let pool = system.staking.get_pool(:staker);
+    let delegator = system.new_delegator(amount: 2 * stake_amount);
+    system.delegate(:delegator, :pool, amount: stake_amount);
+    system.advance_epoch();
+    system.increase_delegate(:delegator, :pool, amount: stake_amount);
+    system.advance_epoch();
+    system.advance_block_custom_and_attest(:staker, stake: 2 * stake_amount);
+    let (expected_staker_rewards, expected_pool_rewards) =
+        calculate_staker_strk_rewards_with_balances_v2(
+        amount_own: stake_amount,
+        pool_amount: stake_amount,
+        :commission,
+        :staking_contract,
+        :minting_curve_contract,
+    );
+    assert!(expected_staker_rewards.is_non_zero());
+    assert!(expected_pool_rewards.is_non_zero());
+    assert!(system.staker_claim_rewards(:staker) == expected_staker_rewards);
+    system.advance_epoch();
+    assert!(system.delegator_claim_rewards(:delegator, :pool) == expected_pool_rewards);
+    system.advance_block_custom_and_attest(:staker, stake: 3 * stake_amount);
+    let (expected_staker_rewards, expected_pool_rewards) =
+        calculate_staker_strk_rewards_with_balances_v2(
+        amount_own: stake_amount,
+        pool_amount: 2 * stake_amount,
+        :commission,
+        :staking_contract,
+        :minting_curve_contract,
+    );
+    assert!(expected_staker_rewards.is_non_zero());
+    assert!(expected_pool_rewards.is_non_zero());
+    assert!(system.staker_claim_rewards(:staker) == expected_staker_rewards);
+    system.advance_epoch();
+    assert!(system.delegator_claim_rewards(:delegator, :pool) == expected_pool_rewards);
+}