test: add balances delay flow (#65)

<!-- Reviewable:start -->
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/starkware-libs/starknet-staking/65)
<!-- Reviewable:end -->

<!-- CURSOR_SUMMARY -->
---

> [!NOTE]
> Adds a new BalancesDelayFlow and mainnet fork test to verify staking power and STRK/BTC rewards timing across epochs pre/post-upgrade.
> 
> - **Flows (`src/flow_test/flows.cairo`)**:
>   - **`BalancesDelayFlow`**: New flow exercising delayed balance effects across epochs and an upgrade.
>     - Sets up staker, STRK and BTC delegators; performs stake/delegate increases; triggers upgrade; and asserts:
>       - Total staking power per epoch (Epoch 0 baseline; Epoch 1 pre-/post-upgrade).
>       - Staker rewards via attestations and delegator rewards (STRK/BTC) across Epochs 1–4.
> - **Tests (`src/flow_test/fork_test.cairo`)**:
>   - Adds `balances_delay_flow_test` (MAINNET_LATEST) to run the new flow.
> 
> <sup>Written by [Cursor Bugbot](https://cursor.com/dashboard?tab=bugbot) for commit 5f4a056. This will update automatically on new commits. Configure [here](https://cursor.com/dashboard?tab=bugbot).</sup>
<!-- /CURSOR_SUMMARY -->

Author: arad-starkware

src/flow_test/flow_ideas.md

@@ -50,7 +50,6 @@ more ideas:
 - staker enter in V0, advance epoch, update balance, advance epoch, update balance, upgrade to V3, attest
 - staker enter in V1, advance epoch, update balance, advance epoch, update balance, upgrade to V3, 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, 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
 - 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

@@ -8570,3 +8570,255 @@ pub(crate) impl MemberChangeBalanceClaimRewardsOneRewardsFlowImpl of MultiVersio
         );
     }
 }
+
+/// Flow:
+/// Epoch 0:
+///     Staker stake
+/// Epoch 1:
+///     Staker increase stake
+///     Delegators delegate STRK and BTC
+///     Upgrade
+///     Staker increase stake
+///     Delegators increase delegation STRK and BTC
+///     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 (pre-upgrade)
+///     Test staker balance with attestation rewards - according to Epoch 1 (pre-upgrade)
+/// Epoch 3:
+///     Test delegator balances with attestation rewards - according to Epoch 1 (pre-upgrade)
+///     Test total staking power - according to Epoch 1 (post-upgrade)
+///     Test staker balance with attestation rewards - according to Epoch 1 (post-upgrade)
+/// Epoch 4:
+///     Test delegator balances with attestation rewards - according to Epoch 1 (post-upgrade)
+#[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;
+
+        // Increase stake and delegations.
+        system.increase_stake(:staker, amount: stake_amount);
+        system
+            .increase_delegate(
+                delegator: strk_delegator, pool: strk_pool, amount: strk_delegated_amount,
+            );
+        system
+            .increase_delegate_btc(
+                delegator: btc_delegator,
+                pool: btc_pool,
+                amount: btc_delegated_amount,
+                token: system.btc_token,
+            );
+
+        // 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 (pre-upgrade).
+        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 (pre-upgrade).
+        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 (pre-upgrade).
+        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 (pre-upgrade).
+        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
+        // (pre-upgrade).
+        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);
+
+        // Update variables for Epoch 1 (post-upgrade).
+        staker_stake += stake_amount;
+        strk_pool_balance += strk_delegated_amount;
+        btc_pool_balance += btc_delegated_amount;
+
+        // Test total staking power - according to Epoch 1 (post-upgrade).
+        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 (post-upgrade).
+        let prev_staker_rewards = expected_staker_rewards;
+        let prev_strk_pool_rewards = expected_strk_pool_rewards;
+        let prev_btc_commission_rewards = expected_btc_commission_rewards;
+        let prev_btc_pool_rewards = expected_btc_pool_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 != prev_strk_pool_rewards);
+        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 != prev_btc_commission_rewards);
+        assert!(expected_btc_pool_rewards != prev_btc_pool_rewards);
+
+        // Test staker balance with attestation rewards - according to Epoch 1 (post-upgrade).
+        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
+        // (post-upgrade).
+        system.advance_epoch(); // Epoch 3 - > 4
+        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);
+    }
+}

src/flow_test/fork_test.cairo

@@ -545,3 +545,20 @@ fn delegator_change_balance_before_upgrade_rewards_consensus_flow_test() {
     };
     test_flow_mainnet(ref :flow);
 }
+
+#[test]
+#[fork("MAINNET_LATEST")]
+fn balances_delay_flow_test() {
+    let mut flow = flows::BalancesDelayFlow {
+        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);
+}