Return claimed reward amount

src/UniStaker.sol

@@ -566,17 +566,20 @@ contract UniStaker is INotifiableRewardReceiver, Multicall, EIP712, Nonces {
 
   /// @notice Claim reward tokens the message sender has earned as a stake beneficiary. Tokens are
   /// sent to the message sender.
-  function claimReward() external {
-    _claimReward(msg.sender);
+  /// @return Amount of reward tokens claimed.
+  function claimReward() external returns (uint256) {
+    return _claimReward(msg.sender);
   }
 
   /// @notice Claim earned reward tokens for a beneficiary, using a signature to validate the
   /// beneficiary's intent. Tokens are sent to the beneficiary.
   /// @param _beneficiary Address of the beneficiary who will receive the reward.
   /// @param _deadline The timestamp at which the signature should expire.
   /// @param _signature Signature of the beneficiary authorizing this reward claim.
+  /// @return Amount of reward tokens claimed.
   function claimRewardOnBehalf(address _beneficiary, uint256 _deadline, bytes memory _signature)
     external
+    returns (uint256)
   {
     _revertIfPastDeadline(_deadline);
     _revertIfSignatureIsNotValidNow(
@@ -588,7 +591,7 @@ contract UniStaker is INotifiableRewardReceiver, Multicall, EIP712, Nonces {
       ),
       _signature
     );
-    _claimReward(_beneficiary);
+    return _claimReward(_beneficiary);
   }
 
   /// @notice Called by an authorized rewards notifier to alert the staking contract that a new
@@ -794,18 +797,20 @@ contract UniStaker is INotifiableRewardReceiver, Multicall, EIP712, Nonces {
   }
 
   /// @notice Internal convenience method which claims earned rewards.
+  /// @return Amount of reward tokens claimed.
   /// @dev This method must only be called after proper authorization has been completed.
   /// @dev See public claimReward methods for additional documentation.
-  function _claimReward(address _beneficiary) internal {
+  function _claimReward(address _beneficiary) internal returns (uint256) {
     _checkpointGlobalReward();
     _checkpointReward(_beneficiary);
 
     uint256 _reward = scaledUnclaimedRewardCheckpoint[_beneficiary] / SCALE_FACTOR;
-    if (_reward == 0) return;
+    if (_reward == 0) return 0;
     scaledUnclaimedRewardCheckpoint[_beneficiary] = 0;
     emit RewardClaimed(_beneficiary, _reward);
 
     SafeERC20.safeTransfer(REWARD_TOKEN, _beneficiary, _reward);
+    return _reward;
   }
 
   /// @notice Checkpoints the global reward per token accumulator.

test/UniStaker.t.sol

@@ -4905,6 +4905,33 @@ contract ClaimReward is UniStakerRewardsTest {
     assertEq(rewardToken.balanceOf(_depositor), _earned);
   }
 
+  function testFuzz_ReturnsRewardAmount(
+    address _depositor,
+    address _delegatee,
+    uint256 _stakeAmount,
+    uint256 _rewardAmount,
+    uint256 _durationPercent
+  ) public {
+    vm.assume(_depositor != address(uniStaker));
+
+    (_stakeAmount, _rewardAmount) = _boundToRealisticStakeAndReward(_stakeAmount, _rewardAmount);
+    _durationPercent = bound(_durationPercent, 0, 100);
+
+    // A user deposits staking tokens
+    _boundMintAndStake(_depositor, _stakeAmount, _delegatee);
+    // The contract is notified of a reward
+    _mintTransferAndNotifyReward(_rewardAmount);
+    // A portion of the duration passes
+    _jumpAheadByPercentOfRewardDuration(_durationPercent);
+
+    uint256 _earned = uniStaker.unclaimedReward(_depositor);
+
+    vm.prank(_depositor);
+    uint256 _claimedAmount = uniStaker.claimReward();
+
+    assertEq(_earned, _claimedAmount);
+  }
+
   function testFuzz_ResetsTheRewardsEarnedByTheUser(
     address _depositor,
     address _delegatee,
@@ -5008,6 +5035,56 @@ contract ClaimRewardOnBehalf is UniStakerRewardsTest {
     assertEq(rewardToken.balanceOf(_beneficiary), _earned);
   }
 
+  function testFuzz_ReturnsClaimedRewardAmount(
+    uint256 _beneficiaryPrivateKey,
+    address _sender,
+    uint256 _depositAmount,
+    uint256 _durationPercent,
+    uint256 _rewardAmount,
+    address _delegatee,
+    address _depositor,
+    uint256 _currentNonce,
+    uint256 _deadline
+  ) public {
+    vm.assume(_delegatee != address(0) && _depositor != address(0) && _sender != address(0));
+    _beneficiaryPrivateKey = bound(_beneficiaryPrivateKey, 1, 100e18);
+    address _beneficiary = vm.addr(_beneficiaryPrivateKey);
+
+    UniStaker.DepositIdentifier _depositId;
+    (_depositAmount, _rewardAmount) = _boundToRealisticStakeAndReward(_depositAmount, _rewardAmount);
+    _durationPercent = bound(_durationPercent, 0, 100);
+
+    // A user deposits staking tokens
+    (_depositAmount, _depositId) =
+      _boundMintAndStake(_depositor, _depositAmount, _delegatee, _beneficiary);
+    // The contract is notified of a reward
+    _mintTransferAndNotifyReward(_rewardAmount);
+    // A portion of the duration passes
+    _jumpAheadByPercentOfRewardDuration(_durationPercent);
+    _deadline = bound(_deadline, block.timestamp, type(uint256).max);
+
+    uint256 _earned = uniStaker.unclaimedReward(_beneficiary);
+
+    stdstore.target(address(uniStaker)).sig("nonces(address)").with_key(_beneficiary).checked_write(
+      _currentNonce
+    );
+
+    bytes32 _message = keccak256(
+      abi.encode(
+        uniStaker.CLAIM_REWARD_TYPEHASH(), _beneficiary, uniStaker.nonces(_beneficiary), _deadline
+      )
+    );
+
+    bytes32 _messageHash =
+      keccak256(abi.encodePacked("\x19\x01", EIP712_DOMAIN_SEPARATOR, _message));
+    bytes memory _signature = _sign(_beneficiaryPrivateKey, _messageHash);
+
+    vm.prank(_sender);
+    uint256 _claimedAmount = uniStaker.claimRewardOnBehalf(_beneficiary, _deadline, _signature);
+
+    assertEq(_earned, _claimedAmount);
+  }
+
   function testFuzz_RevertIf_WrongNonceIsUsed(
     uint256 _beneficiaryPrivateKey,
     address _sender,