diff --git a/src/BaseStaking.sol b/src/BaseStaking.sol
index 32e6ffd..a7e2b6a 100644
--- a/src/BaseStaking.sol
+++ b/src/BaseStaking.sol
@@ -207,7 +207,6 @@ abstract contract BaseStaking is OwnableUpgradeable, ERC20VotesUpgradeable {
     //////////////////////////////////////////////////////////////*/
 
     /// @notice Deposit SHU into the contract
-    /// @param user The user address
     /// @param amount The amount of SHU to deposit
     function _deposit(uint256 amount) internal {
         // Calculate the amount of shares to mint
diff --git a/src/Staking.sol b/src/Staking.sol
index 498ced9..e51888d 100644
--- a/src/Staking.sol
+++ b/src/Staking.sol
@@ -242,20 +242,8 @@ contract Staking is BaseStaking {
                 );
             }
 
-            uint256 maxWithdraw = keyperStake.amount - minStake;
-            require(amount <= maxWithdraw, WithdrawAmountTooHigh());
-
-            // The unstake can't never result in a keyper SHU staked < minStake
-            require(
-                _maxWithdraw(keyper, keyperStake.amount) >= amount,
-                WithdrawAmountTooHigh()
-            );
-        } else {
-            // doesn't include the min stake and locked staked as the keyper is not a keyper anymore
-            require(
-                convertToAssets(balanceOf(keyper)) >= amount,
-                WithdrawAmountTooHigh()
-            );
+            uint256 maxWithdrawAvailable = keyperStake.amount - minStake;
+            require(amount <= maxWithdrawAvailable, WithdrawAmountTooHigh());
         }
 
         // Decrease the amount from the stake
@@ -307,31 +295,12 @@ contract Staking is BaseStaking {
     /// @return amount The maximum amount of assets that a keyper can withdraw
     function maxWithdraw(
         address keyper
-    ) public view override returns (uint256) {
-        return _maxWithdraw(keyper, 0);
-    }
-
-    /*//////////////////////////////////////////////////////////////
-                          INTERNAL FUNCTIONS
-    //////////////////////////////////////////////////////////////*/
-
-    /// @notice Get the maximum amount of assets that a keyper can withdraw
-    ///         after unlocking a certain amount
-    ///         - if the keyper has no shares, the function will revert
-    ///         - if the keyper sSHU balance is less or equal than the minimum
-    ///         stake or the total locked amount, the function will return 0
-    /// @param keyper The keyper address
-    /// @param unlockedAmount The amount of unlocked assets
-    /// @return amount The maximum amount of assets that a keyper can withdraw after unlocking a certain amount
-    function _maxWithdraw(
-        address keyper,
-        uint256 unlockedAmount
-    ) internal view virtual returns (uint256 amount) {
+    ) public view override returns (uint256 amount) {
         uint256 assets = convertToAssets(balanceOf(keyper));
         require(assets > 0, UserHasNoShares());
 
         unchecked {
-            uint256 locked = totalLocked[keyper] - unlockedAmount;
+            uint256 locked = totalLocked[keyper];
             uint256 compare = locked >= minStake ? locked : minStake;
 
             // need the first branch as convertToAssets rounds down
diff --git a/test/RewardsDistributor.t.sol b/test/RewardsDistributor.t.sol
index 1170d10..d25ae86 100644
--- a/test/RewardsDistributor.t.sol
+++ b/test/RewardsDistributor.t.sol
@@ -224,7 +224,11 @@ contract OwnableFunctions is RewardsDistributorTest {
 
         uint256 balanceBefore = govToken.balanceOf(_to);
 
-        rewardsDistributor.withdrawFunds(_to, _amount);
+        rewardsDistributor.withdrawFunds(
+            address(rewardsDistributor.rewardToken()),
+            _to,
+            _amount
+        );
 
         assertEq(govToken.balanceOf(_to), balanceBefore + _amount);
     }
@@ -243,7 +247,11 @@ contract OwnableFunctions is RewardsDistributorTest {
             )
         );
         vm.prank(_anyone);
-        rewardsDistributor.withdrawFunds(_to, _amount);
+        rewardsDistributor.withdrawFunds(
+            address(rewardsDistributor.rewardToken()),
+            _to,
+            _amount
+        );
     }
 }
 
diff --git a/test/Staking.t.sol b/test/Staking.t.sol
index 29c0765..70cbf54 100644
--- a/test/Staking.t.sol
+++ b/test/Staking.t.sol
@@ -712,9 +712,8 @@ contract Stake is StakingTest {
         );
 
         // bob unstake maximum he can unstake
-        uint256 maxBobCanWithdraw = staking.exposed_maxWithdraw(bob, bobStake);
         vm.prank(bob);
-        staking.unstake(bob, bobStakeId, maxBobCanWithdraw);
+        staking.unstake(bob, bobStakeId, bobStake);
 
         uint256 bobBalance = govToken.balanceOf(bob);
 
diff --git a/test/helpers/StakingHarness.sol b/test/helpers/StakingHarness.sol
index 0a99125..dee2aae 100644
--- a/test/helpers/StakingHarness.sol
+++ b/test/helpers/StakingHarness.sol
@@ -8,13 +8,6 @@ contract StakingHarness is Staking {
         return nextStakeId;
     }
 
-    function exposed_maxWithdraw(
-        address keyper,
-        uint256 unlockedAmount
-    ) external view virtual returns (uint256) {
-        return _maxWithdraw(keyper, unlockedAmount);
-    }
-
     function exposed_previewWithdraw(
         uint256 amount
     ) external view returns (uint256) {