fix: update rewards before direct claim

contracts/staking/StakingV01.sol

@@ -357,6 +357,7 @@ contract StakingV01 is
         address user = msg.sender;
         StakeInfo storage info = _info[stakeId];
 
+        _updateRewards(info);
         _claimRewards(user, info);
     }
 

tests/test_staking_unstake_claim_rewards.py

@@ -364,6 +364,72 @@ def test_stake_and_claim_rewards(
     assert dip_reserves_before - rewards == dip_reserves_after
 
 
+def test_stake_and_claim_rewards_before_unstake_v01(
+    mockInstance: MockInstance,
+    mockRegistry: MockInstanceRegistry,
+    usd2: USD2,
+    proxyAdmin: OwnableProxyAdmin,
+    proxyAdminOwner: Account,
+    chainRegistryV01: ChainRegistryV01,
+    registryOwner: Account,
+    dip: DIP,
+    instanceOperator: Account,
+    stakingV01Beta: StakingV01,
+    stakingOwner: Account,
+    staker: Account,
+    theOutsider: Account
+):
+    bundle_nft = create_mock_bundle_setup(
+        mockInstance,
+        mockRegistry,
+        usd2,
+        proxyAdmin,
+        proxyAdminOwner,
+        chainRegistryV01,
+        registryOwner,
+        theOutsider)
+
+    reward_rate = stakingV01Beta.toRate(20, -2)
+    stakingV01Beta.setRewardRate(reward_rate, {'from': stakingOwner})
+
+    reward_reserves = 10000 * 10 ** dip.decimals()
+    dip.approve(stakingV01Beta, reward_reserves, {'from': instanceOperator})
+    stakingV01Beta.refillRewardReserves(reward_reserves, {'from': instanceOperator})
+
+    staking_amount = 100000 * 10 ** dip.decimals()
+    prepare_staker(staker, staking_amount, dip, instanceOperator, stakingV01Beta)
+
+    staking_tx = stakingV01Beta.createStake(
+        bundle_nft,
+        staking_amount,
+        {'from': staker})
+    nft_id = staking_tx.events['LogStakingStaked']['id']
+
+    quarter_year = int(stakingV01Beta.YEAR_DURATION() / 4)
+    chain.sleep(quarter_year)
+    chain.mine(1)
+
+    stake_info = stakingV01Beta.getInfo(nft_id)
+    rewards_increment = stakingV01Beta.calculateRewardsIncrement(stake_info)
+    assert rewards_increment > 0
+    assert stake_info.dict()['rewardBalance'] == 0
+
+    dip_balance_before = dip.balanceOf(staker)
+    reward_balance_before = stakingV01Beta.rewardBalance()
+
+    claim_tx = stakingV01Beta.claimRewards(nft_id, {'from': staker})
+
+    assert 'LogStakingRewardsUpdated' in claim_tx.events
+    assert 'LogStakingRewardsClaimed' in claim_tx.events
+
+    rewards_claimed = claim_tx.events['LogStakingRewardsClaimed']['amount']
+    assert rewards_claimed > 0
+    assert delta_is_tiny(rewards_claimed, rewards_increment)
+    assert dip.balanceOf(staker) == dip_balance_before + rewards_claimed
+    assert stakingV01Beta.getInfo(nft_id).dict()['rewardBalance'] == 0
+    assert stakingV01Beta.rewardBalance() == reward_balance_before
+
+
 def test_stake_unstake_and_claim_rewards(
     mockInstance: MockInstance,
     mockRegistry: MockInstanceRegistry,