test: add update_rewards strk pool flow

Author: arad-starkware

src/flow_test/flow_ideas.md

@@ -3,7 +3,6 @@
 - Get staker info while staker has zero balance.
 
 ## `update_rewards`
-- staker with only strk pool.
 - staker with only btc pool.
 - staker with empty pool (STRK + BTC).
 - staker with 2 btc pools with different decimals.

src/flow_test/test.cairo

@@ -17,7 +17,7 @@ use staking::test_utils::constants::{
 };
 use staking::test_utils::{
     StakingInitConfig, calculate_staker_btc_pool_rewards_v2,
-    calculate_staker_strk_rewards_with_balances_v2,
+    calculate_staker_strk_rewards_with_balances_v2, calculate_staker_strk_rewards_with_balances_v3,
     calculate_strk_pool_rewards_with_pool_balance_v2, compute_rewards_per_unit,
     custom_decimals_token, deploy_mock_erc20_decimals_contract,
 };
@@ -2619,3 +2619,47 @@ fn get_stakers_staking_power_100_flow_test() {
         .span();
     assert!(stakers == expected_stakers);
 }
+
+/// Flow:
+/// Start consensus rewards.
+/// Staker stake with STRK pool.
+/// Delegator delegate STRK.
+/// Advance K epochs.
+/// update_rewards.
+/// Test staker rewards.
+/// Test delegator rewards.
+/// Advance epoch.
+/// Test delegator rewards.
+#[test]
+fn update_rewards_strk_pool_flow_test() {
+    let cfg: StakingInitConfig = Default::default();
+    let mut system = SystemConfigTrait::basic_stake_flow_cfg(:cfg).deploy();
+    system.start_consensus_rewards();
+    let stake_amount = system.staking.get_min_stake();
+    let delegation_amount = STRK_CONFIG.min_for_rewards;
+    let staker = system.new_staker(amount: stake_amount);
+    let commission = 200;
+    system.stake(:staker, amount: stake_amount, pool_enabled: true, :commission);
+    let pool = system.staking.get_pool(:staker);
+    let delegator = system.new_delegator(amount: delegation_amount);
+    system.delegate(:delegator, :pool, amount: delegation_amount);
+    system.advance_k_epochs();
+
+    system.update_rewards(:staker);
+    let (expected_staker_rewards, expected_pool_rewards) =
+        calculate_staker_strk_rewards_with_balances_v3(
+        amount_own: stake_amount,
+        pool_amount: delegation_amount,
+        :commission,
+        staking_contract: system.staking.address,
+        minting_curve_contract: system.minting_curve.address,
+    );
+    let staker_rewards = system.staker_claim_rewards(:staker);
+    let delegator_rewards = system.delegator_claim_rewards(:delegator, :pool);
+    assert!(staker_rewards == expected_staker_rewards);
+    assert!(delegator_rewards == Zero::zero());
+
+    system.advance_epoch();
+    let delegator_rewards = system.delegator_claim_rewards(:delegator, :pool);
+    assert!(delegator_rewards == expected_pool_rewards);
+}

src/flow_test/utils.cairo

@@ -39,10 +39,11 @@ use staking::staking::interface::{
     IStakingConsensusDispatcher, IStakingConsensusSafeDispatcher, IStakingDispatcher,
     IStakingDispatcherTrait, IStakingMigrationDispatcher, IStakingMigrationDispatcherTrait,
     IStakingMigrationSafeDispatcher, IStakingPauseDispatcher, IStakingPauseDispatcherTrait,
-    IStakingPoolDispatcher, IStakingPoolSafeDispatcher, IStakingSafeDispatcher,
-    IStakingSafeDispatcherTrait, IStakingTokenManagerDispatcher,
-    IStakingTokenManagerDispatcherTrait, IStakingTokenManagerSafeDispatcher,
-    IStakingTokenManagerSafeDispatcherTrait, StakerInfoV1, StakerInfoV1Trait, StakerPoolInfoV2,
+    IStakingPoolDispatcher, IStakingPoolSafeDispatcher, IStakingRewardsManagerDispatcher,
+    IStakingRewardsManagerDispatcherTrait, IStakingSafeDispatcher, IStakingSafeDispatcherTrait,
+    IStakingTokenManagerDispatcher, IStakingTokenManagerDispatcherTrait,
+    IStakingTokenManagerSafeDispatcher, IStakingTokenManagerSafeDispatcherTrait, StakerInfoV1,
+    StakerInfoV1Trait, StakerPoolInfoV2,
 };
 use staking::staking::objects::{
     EpochInfo, EpochInfoTrait, NormalizedAmount, StakerVersion, StakerVersionTrait,
@@ -338,6 +339,10 @@ pub(crate) impl StakingImpl of StakingTrait {
         IStakingPoolSafeDispatcher { contract_address: self.address }
     }
 
+    fn rewards_manager_dispatcher(self: StakingState) -> IStakingRewardsManagerDispatcher nopanic {
+        IStakingRewardsManagerDispatcher { contract_address: self.address }
+    }
+
     fn set_roles(self: StakingState) {
         set_account_as_upgrade_governor(
             contract: self.address,
@@ -1492,6 +1497,23 @@ pub(crate) impl SystemStakerImpl of SystemStakerTrait {
         self.attestation.unwrap().dispatcher().attest(:block_hash);
     }
 
+    fn start_consensus_rewards(self: SystemState) {
+        let epoch_id = self.staking.get_current_epoch() + K.into();
+        cheat_caller_address_once(
+            contract_address: self.staking.address, caller_address: self.staking.roles.app_governor,
+        );
+        let config_dispatcher = IStakingConfigDispatcher { contract_address: self.staking.address };
+        config_dispatcher.set_consensus_rewards_first_epoch(:epoch_id);
+        self.advance_k_epochs();
+    }
+
+    fn update_rewards(self: SystemState, staker: Staker) {
+        self
+            .staking
+            .rewards_manager_dispatcher()
+            .update_rewards(staker_address: staker.staker.address, disable_rewards: false);
+    }
+
     #[feature("safe_dispatcher")]
     fn safe_attest(self: SystemState, staker: Staker) -> Result<(), Array<felt252>> {
         let block_hash = Zero::zero();