test: add another balances delay flow

Author: arad-starkware

src/flow_test/flow_ideas.md

@@ -34,7 +34,6 @@ more ideas:
 - staker enter in V2, advance epoch, advance epoch, upgrade to V3, attest, update balance,advance epoch, attest, advance epoch, attest, test rewards
 - staker enter in V2, advance epoch, update balance, upgrade to V3, attest, advance epoch, attest, advance epoch, attest, test rewards
 - staker enter in V2, advance epoch, upgrade to V3, advance epoch, attest,
-- staker in V2, update balance staker+update balance pool, upgrade, attest in current epoch, attest in next epoch, attest in next next epoch
 - staker in V2, advance epoch, update balance staker+update balance pool, advance epoch, update balance staker+update balance pool, upgrade, update balance staker+update balance pool, attest in current epoch, attest in next epoch, attest in next next epoch
 
 ## pool member balance at curr epoch migration

src/flow_test/flows.cairo

@@ -8716,6 +8716,179 @@ pub(crate) impl MemberChangeBalanceClaimRewardsOneRewardsFlowImpl of MultiVersio
     }
 }
 
+/// Flow:
+/// Epoch 0:
+///     Staker stake
+/// Epoch 1:
+///     Staker increase stake
+///     Delegators delegate STRK and BTC
+///     Upgrade
+///     Test total staking power - according to Epoch 0
+///     Test staker balance with attestation rewards - according to Epoch 0
+/// Epoch 2:
+///     Test delegator balances with attestation rewards - according to Epoch 0
+///     Test total staking power - according to Epoch 1
+///     Test staker balance with attestation rewards - according to Epoch 1
+/// Epoch 3:
+///     Test delegator balances with attestation rewards - according to Epoch 1
+#[derive(Drop, Copy)]
+pub(crate) struct BalancesDelayFlow {
+    pub(crate) staker: Option<Staker>,
+    pub(crate) stake_amount: Option<Amount>,
+    pub(crate) commission: Option<Commission>,
+    pub(crate) strk_delegated_amount: Option<Amount>,
+    pub(crate) btc_delegated_amount: Option<Amount>,
+    pub(crate) strk_delegator: Option<Delegator>,
+    pub(crate) btc_delegator: Option<Delegator>,
+    pub(crate) strk_pool: Option<ContractAddress>,
+    pub(crate) btc_pool: Option<ContractAddress>,
+}
+pub(crate) impl BalancesDelayFlowImpl of FlowTrait<BalancesDelayFlow> {
+    fn setup_v2(ref self: BalancesDelayFlow, ref system: SystemState) {
+        let stake_amount = system.staking.get_min_stake();
+        let strk_delegated_amount = STRK_CONFIG.min_for_rewards;
+        let btc_delegated_amount = TEST_MIN_BTC_FOR_REWARDS;
+        let staker = system.new_staker(amount: stake_amount * 3);
+        let commission = 200;
+        let strk_delegator = system.new_delegator(amount: strk_delegated_amount * 2);
+        let btc_delegator = system
+            .new_btc_delegator(amount: btc_delegated_amount * 2, token: system.btc_token);
+
+        system.stake(:staker, amount: stake_amount, pool_enabled: true, :commission);
+        let strk_pool = system.staking.get_pool(:staker);
+        let btc_pool = system
+            .set_open_for_delegation(:staker, token_address: system.btc_token.contract_address());
+
+        system.advance_epoch(); // Epoch 0 - > 1
+        system.increase_stake(:staker, amount: stake_amount);
+        system.delegate(delegator: strk_delegator, pool: strk_pool, amount: strk_delegated_amount);
+        system
+            .delegate_btc(
+                delegator: btc_delegator,
+                pool: btc_pool,
+                amount: btc_delegated_amount,
+                token: system.btc_token,
+            );
+
+        system.set_staker_for_migration(staker_address: staker.staker.address);
+        self.staker = Option::Some(staker);
+        self.stake_amount = Option::Some(stake_amount);
+        self.commission = Option::Some(commission);
+        self.strk_delegated_amount = Option::Some(strk_delegated_amount);
+        self.btc_delegated_amount = Option::Some(btc_delegated_amount);
+        self.strk_delegator = Option::Some(strk_delegator);
+        self.btc_delegator = Option::Some(btc_delegator);
+        self.strk_pool = Option::Some(strk_pool);
+        self.btc_pool = Option::Some(btc_pool);
+    }
+
+    #[feature("safe_dispatcher")]
+    fn test(self: BalancesDelayFlow, ref system: SystemState) {
+        let staker = self.staker.unwrap();
+        let stake_amount = self.stake_amount.unwrap();
+        let commission = self.commission.unwrap();
+        let strk_delegated_amount = self.strk_delegated_amount.unwrap();
+        let btc_delegated_amount = self.btc_delegated_amount.unwrap();
+        let strk_delegator = self.strk_delegator.unwrap();
+        let btc_delegator = self.btc_delegator.unwrap();
+        let strk_pool = self.strk_pool.unwrap();
+        let btc_pool = self.btc_pool.unwrap();
+
+        let staking_contract = system.staking.address;
+        let minting_curve_contract = system.minting_curve.address;
+
+        // Test total staking power - according to Epoch 0.
+        let mut staker_stake = stake_amount;
+        let total_staking_power = system.staking.get_current_total_staking_power_v2();
+        let expected_total_staking_power = (
+            NormalizedAmountTrait::from_strk_native_amount(staker_stake), Zero::zero(),
+        );
+        assert!(total_staking_power == expected_total_staking_power);
+
+        // Calculate expected rewards - according to Epoch 0.
+        let (expected_staker_rewards, _) = calculate_staker_strk_rewards_with_balances_v2(
+            amount_own: staker_stake,
+            pool_amount: Zero::zero(),
+            :commission,
+            :staking_contract,
+            :minting_curve_contract,
+        );
+        assert!(expected_staker_rewards.is_non_zero());
+
+        // Test staker balance with attestation rewards - according to Epoch 0.
+        system
+            .advance_block_into_attestation_window_custom_stake(
+                staker_address: staker.staker.address, stake: staker_stake,
+            );
+        system.attest(:staker);
+        let staker_rewards = system.staker_claim_rewards(:staker);
+        assert!(staker_rewards == expected_staker_rewards);
+
+        // Advance epoch - test delegator balances with attestation rewards - according to Epoch 0.
+        system.advance_epoch(); // Epoch 1 - > 2
+        let strk_delegator_rewards = system
+            .delegator_claim_rewards(delegator: strk_delegator, pool: strk_pool);
+        let btc_delegator_rewards = system
+            .delegator_claim_rewards(delegator: btc_delegator, pool: btc_pool);
+        assert!(strk_delegator_rewards == Zero::zero());
+        assert!(btc_delegator_rewards == Zero::zero());
+
+        // Update variables for Epoch 1.
+        staker_stake += stake_amount;
+        let mut strk_pool_balance = strk_delegated_amount;
+        let mut btc_pool_balance = btc_delegated_amount;
+
+        // Test total staking power - according to Epoch 1.
+        let total_staking_power = system.staking.get_current_total_staking_power_v2();
+        let expected_total_staking_power = (
+            NormalizedAmountTrait::from_strk_native_amount(staker_stake + strk_pool_balance),
+            NormalizedAmountTrait::from_native_amount(btc_pool_balance, TEST_BTC_DECIMALS),
+        );
+        assert!(total_staking_power == expected_total_staking_power);
+
+        // Calculate expected rewards - according to Epoch 1.
+        let prev_staker_rewards = expected_staker_rewards;
+        let (expected_staker_rewards, expected_strk_pool_rewards) =
+            calculate_staker_strk_rewards_with_balances_v2(
+            amount_own: staker_stake,
+            pool_amount: strk_pool_balance,
+            :commission,
+            :staking_contract,
+            :minting_curve_contract,
+        );
+        assert!(expected_staker_rewards != prev_staker_rewards);
+        assert!(expected_strk_pool_rewards.is_non_zero());
+        let (expected_btc_commission_rewards, expected_btc_pool_rewards) =
+            calculate_staker_btc_pool_rewards_v2(
+            pool_balance: btc_pool_balance,
+            :commission,
+            :staking_contract,
+            :minting_curve_contract,
+            token_address: system.btc_token.contract_address(),
+        );
+        assert!(expected_btc_commission_rewards.is_non_zero());
+        assert!(expected_btc_pool_rewards.is_non_zero());
+
+        // Test staker balance with attestation rewards - according to Epoch 1.
+        system
+            .advance_block_into_attestation_window_custom_stake(
+                staker_address: staker.staker.address, stake: staker_stake + strk_pool_balance,
+            );
+        system.attest(:staker);
+        let staker_rewards = system.staker_claim_rewards(:staker);
+        assert!(staker_rewards == expected_staker_rewards + expected_btc_commission_rewards);
+
+        // Advance epoch - test delegator balances with attestation rewards - according to Epoch 1.
+        system.advance_epoch(); // Epoch 2 - > 3
+        let strk_delegator_rewards = system
+            .delegator_claim_rewards(delegator: strk_delegator, pool: strk_pool);
+        let btc_delegator_rewards = system
+            .delegator_claim_rewards(delegator: btc_delegator, pool: btc_pool);
+        assert!(strk_delegator_rewards == expected_strk_pool_rewards);
+        assert!(btc_delegator_rewards == expected_btc_pool_rewards);
+    }
+}
+
 /// Flow:
 /// Epoch 0:
 ///     Staker stake
@@ -8738,7 +8911,7 @@ pub(crate) impl MemberChangeBalanceClaimRewardsOneRewardsFlowImpl of MultiVersio
 /// Epoch 4:
 ///     Test delegator balances with attestation rewards - according to Epoch 1 (post-upgrade)
 #[derive(Drop, Copy)]
-pub(crate) struct BalancesDelayFlow {
+pub(crate) struct BalancesDelayBalanceChangeFlow {
     pub(crate) staker: Option<Staker>,
     pub(crate) stake_amount: Option<Amount>,
     pub(crate) commission: Option<Commission>,
@@ -8749,8 +8922,8 @@ pub(crate) struct BalancesDelayFlow {
     pub(crate) strk_pool: Option<ContractAddress>,
     pub(crate) btc_pool: Option<ContractAddress>,
 }
-pub(crate) impl BalancesDelayFlowImpl of FlowTrait<BalancesDelayFlow> {
-    fn setup_v2(ref self: BalancesDelayFlow, ref system: SystemState) {
+pub(crate) impl BalancesDelayBalanceChangeFlowImpl of FlowTrait<BalancesDelayBalanceChangeFlow> {
+    fn setup_v2(ref self: BalancesDelayBalanceChangeFlow, ref system: SystemState) {
         let stake_amount = system.staking.get_min_stake();
         let strk_delegated_amount = STRK_CONFIG.min_for_rewards;
         let btc_delegated_amount = TEST_MIN_BTC_FOR_REWARDS;
@@ -8789,7 +8962,7 @@ pub(crate) impl BalancesDelayFlowImpl of FlowTrait<BalancesDelayFlow> {
     }
 
     #[feature("safe_dispatcher")]
-    fn test(self: BalancesDelayFlow, ref system: SystemState) {
+    fn test(self: BalancesDelayBalanceChangeFlow, ref system: SystemState) {
         let staker = self.staker.unwrap();
         let stake_amount = self.stake_amount.unwrap();
         let commission = self.commission.unwrap();

src/flow_test/fork_test.cairo

@@ -612,6 +612,23 @@ fn delegator_rewards_v0_find_sigma_flow_test() {
     test_flow_mainnet(ref :flow);
 }
 
+#[test]
+#[fork("MAINNET_LATEST")]
+fn balances_delay_balance_change_flow_test() {
+    let mut flow = flows::BalancesDelayBalanceChangeFlow {
+        staker: Option::None,
+        stake_amount: Option::None,
+        commission: Option::None,
+        strk_delegated_amount: Option::None,
+        btc_delegated_amount: Option::None,
+        strk_delegator: Option::None,
+        btc_delegator: Option::None,
+        strk_pool: Option::None,
+        btc_pool: Option::None,
+    };
+    test_flow_mainnet(ref :flow);
+}
+
 #[test]
 #[fork("MAINNET_LATEST")]
 fn enable_disable_token_before_after_upgrade_flow_test() {