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

@@ -2756,3 +2756,47 @@ fn update_rewards_disable_rewards_flow_test() {
         :result, expected_error: StakingError::REWARDS_ALREADY_UPDATED.describe(),
     );
 }
+
+/// 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, disable_rewards: false);
+    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

@@ -1518,6 +1518,12 @@ pub(crate) impl SystemStakerImpl of SystemStakerTrait {
             .update_rewards(staker_address: staker.staker.address, :disable_rewards);
     }
 
+    fn start_consensus_rewards(self: SystemState) {
+        let epoch_id = self.staking.get_current_epoch() + K.into();
+        self.set_consensus_rewards_first_epoch(:epoch_id);
+        self.advance_k_epochs();
+    }
+
     #[feature("safe_dispatcher")]
     fn safe_attest(self: SystemState, staker: Staker) -> Result<(), Array<felt252>> {
         let block_hash = Zero::zero();