test: add token state after upgrade flow

Author: arad-starkware

src/flow_test/flow_ideas.md

@@ -62,6 +62,3 @@ more ideas:
 - Member from V0, change balance at V1, no action at V2, test curr balance
 - Member from V2, change balance, upgrade, test curr balance
 - Cover all ifs with migration from: V0, V1, V2.
-
-## k=1 -> k=2 token
-- enable token A and disable token B, next epoch upgrade, test views and rewards.

src/flow_test/flows.cairo

@@ -35,9 +35,9 @@ use staking::test_utils::constants::{
 };
 use staking::test_utils::{
     calculate_pool_member_rewards, calculate_staker_btc_pool_rewards_v2,
-    calculate_staker_strk_rewards_v2, calculate_staker_strk_rewards_with_balances_v2,
-    calculate_staker_strk_rewards_with_balances_v3, calculate_strk_pool_rewards_v1,
-    calculate_strk_pool_rewards_v2, calculate_strk_pool_rewards_v3,
+    calculate_staker_btc_pool_rewards_v3, calculate_staker_strk_rewards_v2,
+    calculate_staker_strk_rewards_with_balances_v2, calculate_staker_strk_rewards_with_balances_v3,
+    calculate_strk_pool_rewards_v1, calculate_strk_pool_rewards_v2, calculate_strk_pool_rewards_v3,
     calculate_strk_pool_rewards_with_pool_balance_v2, compute_rewards_per_unit,
     declare_pool_contract, declare_pool_eic_contract, declare_staking_contract,
     load_from_simple_map, load_one_felt, strk_pool_update_rewards_v0, to_amount_18_decimals,
@@ -8822,3 +8822,153 @@ pub(crate) impl BalancesDelayFlowImpl of FlowTrait<BalancesDelayFlow> {
         assert!(btc_delegator_rewards == expected_btc_pool_rewards);
     }
 }
+
+/// Flow:
+/// Add tokens A and B
+/// Delegate for both tokens
+/// Enable token A, disable token B
+/// Advance epoch
+/// Upgrade - test views and rewards
+/// Start consensus rewards - test rewards
+#[derive(Drop, Copy)]
+pub(crate) struct EnableDisableTokenBeforeAfterUpgradeFlow {
+    pub(crate) token_a: Option<Token>,
+    pub(crate) token_b: Option<Token>,
+    pub(crate) staker: Option<Staker>,
+    pub(crate) pool_a: Option<ContractAddress>,
+    pub(crate) pool_b: Option<ContractAddress>,
+    pub(crate) delegator_a: Option<Delegator>,
+    pub(crate) delegator_b: Option<Delegator>,
+    pub(crate) delegation_amount: Option<Amount>,
+}
+pub(crate) impl EnableDisableTokenBeforeAfterUpgradeFlowImpl of FlowTrait<
+    EnableDisableTokenBeforeAfterUpgradeFlow,
+> {
+    fn setup_v2(ref self: EnableDisableTokenBeforeAfterUpgradeFlow, ref system: SystemState) {
+        let token_a = system.deploy_new_btc_token(name: "TOKEN_A", decimals: TEST_BTC_DECIMALS);
+        let token_b = system.deploy_new_btc_token(name: "TOKEN_B", decimals: TEST_BTC_DECIMALS);
+        system.staking.add_token(token_address: token_a.contract_address());
+        system.staking.add_token(token_address: token_b.contract_address());
+        system.staking.enable_token(token_address: token_b.contract_address());
+
+        let stake_amount = system.staking.get_min_stake();
+        let delegation_amount = TEST_MIN_BTC_FOR_REWARDS;
+        let staker = system.new_staker(amount: stake_amount);
+        system.stake(:staker, amount: stake_amount, pool_enabled: true, commission: 200);
+        let pool_a = system
+            .set_open_for_delegation(:staker, token_address: token_a.contract_address());
+        let pool_b = system
+            .set_open_for_delegation(:staker, token_address: token_b.contract_address());
+        let delegator_a = system.new_btc_delegator(amount: delegation_amount, token: token_a);
+        let delegator_b = system.new_btc_delegator(amount: delegation_amount, token: token_b);
+        system
+            .delegate_btc(
+                delegator: delegator_a, pool: pool_a, amount: delegation_amount, token: token_a,
+            );
+        system
+            .delegate_btc(
+                delegator: delegator_b, pool: pool_b, amount: delegation_amount, token: token_b,
+            );
+
+        system.advance_epoch();
+        system.staking.disable_token(token_address: token_b.contract_address());
+        system.staking.enable_token(token_address: token_a.contract_address());
+
+        system.advance_epoch();
+        system.set_staker_for_migration(staker_address: staker.staker.address);
+        self.token_a = Option::Some(token_a);
+        self.token_b = Option::Some(token_b);
+        self.staker = Option::Some(staker);
+        self.pool_a = Option::Some(pool_a);
+        self.pool_b = Option::Some(pool_b);
+        self.delegator_a = Option::Some(delegator_a);
+        self.delegator_b = Option::Some(delegator_b);
+        self.delegation_amount = Option::Some(delegation_amount);
+    }
+
+    fn test(self: EnableDisableTokenBeforeAfterUpgradeFlow, ref system: SystemState) {
+        let token_a = self.token_a.unwrap();
+        let token_b = self.token_b.unwrap();
+        let staker = self.staker.unwrap();
+        let pool_a = self.pool_a.unwrap();
+        let pool_b = self.pool_b.unwrap();
+        let delegator_a = self.delegator_a.unwrap();
+        let delegator_b = self.delegator_b.unwrap();
+        let delegation_amount = self.delegation_amount.unwrap();
+        let staker_info = system.staker_info_v1(:staker);
+        let commission = staker_info.pool_info.unwrap().commission;
+        let stake_amount = staker_info.amount_own;
+
+        // Test get_tokens view.
+        let tokens = system.staking.dispatcher().get_tokens();
+        let expected_tokens = array![
+            (STRK_TOKEN_ADDRESS, true), (system.btc_token.contract_address(), true),
+            (token_a.contract_address(), true), (token_b.contract_address(), false),
+        ]
+            .span();
+        assert!(tokens == expected_tokens);
+
+        // Test rewards - only token A.
+        system.advance_block_into_attestation_window(:staker);
+        system.attest(:staker);
+        system.advance_epoch();
+        let staker_rewards = system.staker_claim_rewards(:staker);
+        let (expected_staker_rewards, _) = calculate_staker_strk_rewards_v2(
+            :staker_info,
+            staking_contract: system.staking.address,
+            minting_curve_contract: system.minting_curve.address,
+        );
+        assert!(expected_staker_rewards.is_non_zero());
+        let rewards_a = system.delegator_claim_rewards(delegator: delegator_a, pool: pool_a);
+        let rewards_b = system.delegator_claim_rewards(delegator: delegator_b, pool: pool_b);
+        let (expected_commission_rewards, expected_pool_rewards) =
+            calculate_staker_btc_pool_rewards_v2(
+            pool_balance: delegation_amount,
+            :commission,
+            staking_contract: system.staking.address,
+            minting_curve_contract: system.minting_curve.address,
+            token_address: token_a.contract_address(),
+        );
+        assert!(expected_commission_rewards.is_non_zero());
+        assert!(expected_pool_rewards.is_non_zero());
+
+        assert!(staker_rewards == expected_staker_rewards + expected_commission_rewards);
+        assert!(rewards_a == expected_pool_rewards);
+        assert!(rewards_b.is_zero());
+
+        system.start_consensus_rewards();
+        system.update_rewards(:staker, disable_rewards: false);
+        system.advance_epoch();
+        let staker_rewards = system.staker_claim_rewards(:staker);
+        let rewards_a = system.delegator_claim_rewards(delegator: delegator_a, pool: pool_a);
+        let rewards_b = system.delegator_claim_rewards(delegator: delegator_b, pool: pool_b);
+
+        let (expected_staker_rewards, _) = calculate_staker_strk_rewards_with_balances_v3(
+            amount_own: stake_amount,
+            pool_amount: Zero::zero(),
+            :commission,
+            staking_contract: system.staking.address,
+            minting_curve_contract: system.minting_curve.address,
+        );
+        assert!(expected_staker_rewards.is_non_zero());
+        let (expected_commission_rewards, expected_pool_rewards) =
+            calculate_staker_btc_pool_rewards_v3(
+            normalized_pool_balance: NormalizedAmountTrait::from_native_amount(
+                amount: delegation_amount, decimals: TEST_BTC_DECIMALS,
+            ),
+            normalized_staker_total_btc_balance: NormalizedAmountTrait::from_native_amount(
+                amount: delegation_amount, decimals: TEST_BTC_DECIMALS,
+            ),
+            :commission,
+            staking_contract: system.staking.address,
+            minting_curve_contract: system.minting_curve.address,
+            token_address: token_a.contract_address(),
+        );
+        assert!(expected_commission_rewards.is_non_zero());
+        assert!(expected_pool_rewards.is_non_zero());
+
+        assert!(staker_rewards == expected_staker_rewards + expected_commission_rewards);
+        assert!(rewards_a == expected_pool_rewards);
+        assert!(rewards_b.is_zero());
+    }
+}

src/flow_test/fork_test.cairo

@@ -562,3 +562,19 @@ fn balances_delay_flow_test() {
     };
     test_flow_mainnet(ref :flow);
 }
+
+#[test]
+#[fork("MAINNET_LATEST")]
+fn enable_disable_token_before_after_upgrade_flow_test() {
+    let mut flow = flows::EnableDisableTokenBeforeAfterUpgradeFlow {
+        token_a: Option::None,
+        token_b: Option::None,
+        staker: Option::None,
+        pool_a: Option::None,
+        pool_b: Option::None,
+        delegator_a: Option::None,
+        delegator_b: Option::None,
+        delegation_amount: Option::None,
+    };
+    test_flow_mainnet(ref :flow);
+}