test: staker migration more flows

Author: noa-starkware

src/flow_test/flow_ideas.md

@@ -40,11 +40,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 enter in V0, upgrade to V3, attest
-- staker enter in V0, advance epoch, update balance, upgrade to V3, attest, update balance, advance epoch, attest
-- staker enter in V1, advance epoch, update balance, upgrade to V3, attest, update balance, advance epoch, attest
-- 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, 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

@@ -8822,3 +8822,150 @@ pub(crate) impl BalancesDelayFlowImpl of FlowTrait<BalancesDelayFlow> {
         assert!(btc_delegator_rewards == expected_btc_pool_rewards);
     }
 }
+
+/// Flow:
+/// Staker stake in V0
+/// Upgrade
+/// Advance Epoch
+/// Attest
+/// Increase stake
+/// Advance Epoch
+/// Attest
+/// Advance K epochs
+/// Attest
+#[derive(Drop, Copy)]
+pub(crate) struct StakerV0AttestFlow {
+    pub(crate) staker: Option<Staker>,
+}
+pub(crate) impl StakerV0AttestFlowImpl of FlowTrait<StakerV0AttestFlow> {
+    fn setup(ref self: StakerV0AttestFlow, ref system: SystemState) {
+        let amount = system.staking.get_min_stake();
+        let staker = system.new_staker(amount: 2 * amount);
+        system.stake(:staker, :amount, pool_enabled: false, commission: 0);
+        self.staker = Option::Some(staker);
+        system.set_staker_for_migration(staker_address: staker.staker.address);
+    }
+    fn test(self: StakerV0AttestFlow, ref system: SystemState) {
+        let staker = self.staker.unwrap();
+        let amount = system.staker_info_v1(:staker).amount_own;
+        system.advance_epoch();
+        system.advance_block_custom_and_attest(:staker, stake: amount);
+        let (expected_rewards, _) = calculate_staker_strk_rewards_with_balances_v2(
+            amount_own: amount,
+            pool_amount: 0,
+            commission: 0,
+            staking_contract: system.staking.address,
+            minting_curve_contract: system.minting_curve.address,
+        );
+        assert!(system.staker_claim_rewards(:staker) == expected_rewards);
+        system.increase_stake(:staker, :amount);
+        system.advance_epoch();
+        system.advance_block_custom_and_attest(:staker, stake: amount);
+        assert!(system.staker_claim_rewards(:staker) == expected_rewards);
+        system.advance_k_epochs_and_attest(:staker);
+        let (expected_rewards, _) = calculate_staker_strk_rewards_with_balances_v2(
+            amount_own: 2 * amount,
+            pool_amount: 0,
+            commission: 0,
+            staking_contract: system.staking.address,
+            minting_curve_contract: system.minting_curve.address,
+        );
+        assert!(system.staker_claim_rewards(:staker) == expected_rewards);
+    }
+}
+
+/// Flow:
+/// Staker stake in V1
+/// Advance Epoch
+/// Increase stake
+/// Advance Epoch
+/// Upgrade
+/// Attest
+/// Increase stake
+/// Advance Epoch
+/// Attest
+/// Advance K epochs
+/// Attest
+#[derive(Drop, Copy)]
+pub(crate) struct StakerV1AttestFlow {
+    pub(crate) staker: Option<Staker>,
+}
+pub(crate) impl StakerV1AttestFlowImpl of FlowTrait<StakerV1AttestFlow> {
+    fn setup_v1(ref self: StakerV1AttestFlow, ref system: SystemState) {
+        let amount = system.staking.get_min_stake();
+        let staker = system.new_staker(amount: 3 * amount);
+        system.stake(:staker, :amount, pool_enabled: false, commission: 0);
+        system.advance_epoch();
+        system.increase_stake(:staker, :amount);
+        system.advance_epoch();
+        self.staker = Option::Some(staker);
+        system.set_staker_for_migration(staker_address: staker.staker.address);
+    }
+    fn test(self: StakerV1AttestFlow, ref system: SystemState) {
+        let staker = self.staker.unwrap();
+        let amount = system.staker_info_v1(:staker).amount_own / 2;
+        system.advance_block_custom_and_attest(:staker, stake: 2 * amount);
+        let (expected_rewards, _) = calculate_staker_strk_rewards_with_balances_v2(
+            amount_own: 2 * amount,
+            pool_amount: 0,
+            commission: 0,
+            staking_contract: system.staking.address,
+            minting_curve_contract: system.minting_curve.address,
+        );
+        assert!(system.staker_claim_rewards(:staker) == expected_rewards);
+        system.increase_stake(:staker, :amount);
+        system.advance_epoch();
+        system.advance_block_custom_and_attest(:staker, stake: 2 * amount);
+        assert!(system.staker_claim_rewards(:staker) == expected_rewards);
+        system.advance_k_epochs_and_attest(:staker);
+        let (expected_rewards, _) = calculate_staker_strk_rewards_with_balances_v2(
+            amount_own: 3 * amount,
+            pool_amount: 0,
+            commission: 0,
+            staking_contract: system.staking.address,
+            minting_curve_contract: system.minting_curve.address,
+        );
+        assert!(system.staker_claim_rewards(:staker) == expected_rewards);
+    }
+}
+
+/// Flow:
+/// Staker stake in V1
+/// Advance Epoch
+/// Increase stake
+/// Advance Epoch
+/// Increase stake
+/// Advance Epoch
+/// Upgrade
+/// Attest
+#[derive(Drop, Copy)]
+pub(crate) struct StakerV1ChangeBalanceAttestFlow {
+    pub(crate) staker: Option<Staker>,
+}
+pub(crate) impl StakerV1ChangeBalanceAttestFlowImpl of FlowTrait<StakerV1ChangeBalanceAttestFlow> {
+    fn setup_v1(ref self: StakerV1ChangeBalanceAttestFlow, ref system: SystemState) {
+        let amount = system.staking.get_min_stake();
+        let staker = system.new_staker(amount: 3 * amount);
+        system.stake(:staker, :amount, pool_enabled: false, commission: 0);
+        system.advance_epoch();
+        system.increase_stake(:staker, :amount);
+        system.advance_epoch();
+        system.increase_stake(:staker, :amount);
+        system.advance_epoch();
+        self.staker = Option::Some(staker);
+        system.set_staker_for_migration(staker_address: staker.staker.address);
+    }
+    fn test(self: StakerV1ChangeBalanceAttestFlow, ref system: SystemState) {
+        let staker = self.staker.unwrap();
+        let stake_amount = system.staker_info_v1(:staker).amount_own;
+        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: 0,
+            commission: 0,
+            staking_contract: system.staking.address,
+            minting_curve_contract: system.minting_curve.address,
+        );
+        assert!(system.staker_claim_rewards(:staker) == expected_rewards);
+    }
+}

src/flow_test/fork_test.cairo

@@ -562,3 +562,24 @@ fn balances_delay_flow_test() {
     };
     test_flow_mainnet(ref :flow);
 }
+
+#[test]
+#[fork("MAINNET_LATEST")]
+fn staker_v0_attest_flow_test() {
+    let mut flow = flows::StakerV0AttestFlow { staker: Option::None };
+    test_flow_mainnet(ref :flow);
+}
+
+#[test]
+#[fork("MAINNET_LATEST")]
+fn staker_v1_attest_flow_test() {
+    let mut flow = flows::StakerV1AttestFlow { staker: Option::None };
+    test_flow_mainnet(ref :flow);
+}
+
+#[test]
+#[fork("MAINNET_LATEST")]
+fn staker_v1_change_balance_attest_flow_test() {
+    let mut flow = flows::StakerV1ChangeBalanceAttestFlow { staker: Option::None };
+    test_flow_mainnet(ref :flow);
+}