test: add accrue_rewards_across_versions flow

Author: arad-starkware

src/flow_test/flow_ideas.md

@@ -16,7 +16,6 @@
 - with member from previous versions.
 
 ## k=1 -> k=2 Migration Staker
-- staker enter in V0, attest in V1, update balance in V1, attest in V1, attest in V2, update balance in V2, attest in V2, attest in V3, update balance in V3, attest in V3, test rewards.
 - 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
 

src/flow_test/flows.cairo

@@ -9946,3 +9946,124 @@ pub(crate) impl EnableDisableTokenBeforeAfterUpgradeFlowImpl of FlowTrait<
         assert!(rewards_b.is_zero());
     }
 }
+
+/// Flow:
+/// Staker stake
+/// Accrue rewards
+/// Upgrade to V1
+/// Attest
+/// Increase stake
+/// Attest
+/// Upgrade to V2
+/// Attest
+/// Increase stake
+/// Attest
+/// Upgrade to V3
+/// Attest
+/// Increase stake
+/// Attest
+/// Start consensus rewards
+/// update_rewards
+/// Increase stake
+/// update_rewards
+/// Test rewards
+#[derive(Drop, Copy)]
+pub(crate) struct AccrueRewardsDifferentBalancesAcrossVersionsFlow {
+    pub(crate) staker: Option<Staker>,
+    pub(crate) stake_amount: Option<Amount>,
+    pub(crate) rewards: Option<Amount>,
+}
+pub(crate) impl AccrueRewardsDifferentBalancesAcrossVersionsFlowImpl of FlowTrait<
+    AccrueRewardsDifferentBalancesAcrossVersionsFlow,
+> {
+    fn setup(ref self: AccrueRewardsDifferentBalancesAcrossVersionsFlow, ref system: SystemState) {
+        let stake_amount = system.staking.get_min_stake();
+        let staker = system.new_staker(amount: stake_amount * 6);
+
+        system.stake(:staker, amount: stake_amount, pool_enabled: false, commission: Zero::zero());
+        system.accrue_rewards(:staker);
+
+        system.set_staker_for_migration(staker_address: staker.staker.address);
+        self.staker = Option::Some(staker);
+        self.stake_amount = Option::Some(stake_amount);
+    }
+
+    fn setup_v1(
+        ref self: AccrueRewardsDifferentBalancesAcrossVersionsFlow, ref system: SystemState,
+    ) {
+        let staker = self.staker.unwrap();
+        let stake_amount = self.stake_amount.unwrap();
+
+        system.accrue_rewards(:staker);
+        system.increase_stake(:staker, amount: stake_amount);
+        system.accrue_rewards(:staker);
+    }
+
+    fn setup_v2(
+        ref self: AccrueRewardsDifferentBalancesAcrossVersionsFlow, ref system: SystemState,
+    ) {
+        let staker = self.staker.unwrap();
+        let stake_amount = self.stake_amount.unwrap();
+
+        system.accrue_rewards(:staker);
+        system.increase_stake(:staker, amount: stake_amount);
+        system.accrue_rewards(:staker);
+
+        let rewards = system.staker_info_v1(:staker).unclaimed_rewards_own;
+        assert!(rewards.is_non_zero());
+        self.rewards = Option::Some(rewards);
+    }
+
+    fn test(self: AccrueRewardsDifferentBalancesAcrossVersionsFlow, ref system: SystemState) {
+        let staker = self.staker.unwrap();
+        let stake_amount = self.stake_amount.unwrap();
+        let mut total_expected_rewards = self.rewards.unwrap();
+
+        system.accrue_rewards(:staker);
+        let (expected_rewards, _) = calculate_staker_strk_rewards_with_balances_v2(
+            amount_own: stake_amount * 3,
+            pool_amount: Zero::zero(),
+            commission: Zero::zero(),
+            staking_contract: system.staking.address,
+            minting_curve_contract: system.minting_curve.address,
+        );
+        total_expected_rewards += expected_rewards;
+
+        system.increase_stake(:staker, amount: stake_amount);
+        system.accrue_rewards(:staker);
+        let (expected_rewards, _) = calculate_staker_strk_rewards_with_balances_v2(
+            amount_own: stake_amount * 4,
+            pool_amount: Zero::zero(),
+            commission: Zero::zero(),
+            staking_contract: system.staking.address,
+            minting_curve_contract: system.minting_curve.address,
+        );
+        total_expected_rewards += expected_rewards;
+
+        system.start_consensus_rewards();
+        system.update_rewards(:staker, disable_rewards: false);
+        let (expected_rewards, _) = calculate_staker_strk_rewards_with_balances_v3(
+            amount_own: stake_amount * 5,
+            pool_amount: Zero::zero(),
+            commission: Zero::zero(),
+            staking_contract: system.staking.address,
+            minting_curve_contract: system.minting_curve.address,
+        );
+        total_expected_rewards += expected_rewards;
+
+        system.increase_stake(:staker, amount: stake_amount);
+        system.advance_k_epochs();
+        system.update_rewards(:staker, disable_rewards: false);
+        let (expected_rewards, _) = calculate_staker_strk_rewards_with_balances_v3(
+            amount_own: stake_amount * 6,
+            pool_amount: Zero::zero(),
+            commission: Zero::zero(),
+            staking_contract: system.staking.address,
+            minting_curve_contract: system.minting_curve.address,
+        );
+        total_expected_rewards += expected_rewards;
+
+        let staker_rewards = system.staker_claim_rewards(:staker);
+        assert!(staker_rewards == total_expected_rewards);
+    }
+}

src/flow_test/fork_test.cairo

@@ -741,3 +741,12 @@ fn enable_disable_token_before_after_upgrade_flow_test() {
     };
     test_flow_mainnet(ref :flow);
 }
+
+#[test]
+#[fork("MAINNET_LATEST")]
+fn accrue_rewards_different_balances_across_versions_flow_test() {
+    let mut flow = flows::AccrueRewardsDifferentBalancesAcrossVersionsFlow {
+        staker: Option::None, stake_amount: Option::None, rewards: Option::None,
+    };
+    test_flow_mainnet(ref :flow);
+}