increase unstake coverage

src/Staking.sol

@@ -10,7 +10,6 @@ import {ERC20VotesUpgradeable} from "@openzeppelin-upgradeable/contracts/token/E
 import {FixedPointMathLib} from "@solmate/utils/FixedPointMathLib.sol";
 import {IRewardsDistributor} from "./interfaces/IRewardsDistributor.sol";
 
-// TODO should be pausable?
 // TODO is this vulnerable to first deposit attack?
 // TODO check calculations
 contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
@@ -181,9 +180,6 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
         uint256 _minStake
     ) public initializer {
         // TODO set name and symbol
-        // Does nothing but calls anyway for consistency
-        __ERC20Votes_init();
-        __Ownable2Step_init();
 
         // Transfer ownership to the DAO contract
         _transferOwnership(newOwner);
@@ -312,12 +308,14 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
                 );
             }
 
+            // TODO branch never reached
             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
+            // TODO branch never reached
             require(
                 convertToAssets(balanceOf(keyper)) >= amount,
                 WithdrawAmountTooHigh()
@@ -421,12 +419,34 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
         }
     }
 
+    /*//////////////////////////////////////////////////////////////
+                              TRANSFER LOGIC
+    //////////////////////////////////////////////////////////////*/
+
+    /// @notice Transfer is disabled
+    function transfer(address, uint256) public pure override returns (bool) {
+        revert TransferDisabled();
+    }
+
+    /// @notice Transfer is disabled
+    function transferFrom(
+        address,
+        address,
+        uint256
+    ) public pure override returns (bool) {
+        revert TransferDisabled();
+    }
+
+    /*//////////////////////////////////////////////////////////////
+                              VIEW FUNCTIONS
+    //////////////////////////////////////////////////////////////*/
+
     /// @notice Calculates the maximum amount of assets that a keyper can withdraw,
     ///         factoring in the principal and any compounded rewards.
     ///         This function subtracts the minimum required stake and includes any amounts
     ///         currently locked. As a result, the maximum withdrawable amount might be less
     ///         than the total withdrawable at the current block timestamp.
-    ///         TODO revisirt natspec
+    ///         TODO revisit natspec
     /// @param keyper The keyper address
     /// @return The maximum amount of assets that a keyper can withdraw
     function maxWithdraw(address keyper) public view virtual returns (uint256) {
@@ -467,28 +487,6 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
         }
     }
 
-    /*//////////////////////////////////////////////////////////////
-                              TRANSFER LOGIC
-    //////////////////////////////////////////////////////////////*/
-
-    /// @notice Transfer is disabled
-    function transfer(address, uint256) public pure override returns (bool) {
-        revert TransferDisabled();
-    }
-
-    /// @notice Transfer is disabled
-    function transferFrom(
-        address,
-        address,
-        uint256
-    ) public pure override returns (bool) {
-        revert TransferDisabled();
-    }
-
-    /*//////////////////////////////////////////////////////////////
-                              VIEW FUNCTIONS
-    //////////////////////////////////////////////////////////////*/
-
     /// @notice Get the total amount of shares the assets are worth
     /// @param assets The amount of assets
     function convertToShares(

test/Staking.t.sol

@@ -462,6 +462,22 @@ contract Stake is StakingTest {
         assertEq(timestamp2, vm.getBlockTimestamp(), "Wrong timestamp");
     }
 
+    function testFuzz_stakeReturnsStakeId(
+        address _depositor,
+        uint256 _amount
+    ) public {
+        _amount = _boundToRealisticStake(_amount);
+
+        _mintGovToken(_depositor, _amount);
+        _setKeyper(_depositor, true);
+
+        vm.assume(_depositor != address(0));
+
+        uint256 stakeId = _stake(_depositor, _amount);
+
+        assertGt(stakeId, 0, "Wrong stake id");
+    }
+
     function testFuzz_increaseTotalLockedWhenStaking(
         address _depositor,
         uint256 _amount
@@ -1070,6 +1086,82 @@ contract Unstake is StakingTest {
         );
     }
 
+    function testFuzz_UnstakeOnlyAmountSpecified(
+        address _depositor,
+        uint256 _amount1,
+        uint256 _amount2,
+        uint256 _jump
+    ) public {
+        _amount1 = _boundToRealisticStake(_amount1);
+        _amount2 = _boundToRealisticStake(_amount2);
+        _jump = _boundUnlockedTime(_jump);
+
+        vm.assume(_amount1 > _amount2);
+        _jump = _boundUnlockedTime(_jump);
+
+        _mintGovToken(_depositor, _amount1);
+
+        _setKeyper(_depositor, true);
+
+        uint256 stakeId = _stake(_depositor, _amount1);
+        assertEq(govToken.balanceOf(_depositor), 0, "Wrong balance");
+
+        _jumpAhead(_jump);
+
+        vm.prank(_depositor);
+        staking.unstake(_depositor, stakeId, _amount2);
+
+        assertEq(govToken.balanceOf(_depositor), _amount2, "Wrong balance");
+
+        uint256[] memory stakeIds = staking.getKeyperStakeIds(_depositor);
+        assertEq(stakeIds.length, 1, "Wrong stake ids");
+
+        (uint256 amount, , ) = staking.stakes(stakeIds[0]);
+
+        assertEq(amount, _amount1 - _amount2, "Wrong amount");
+    }
+
+    function testFuzz_RevertIf_StakeIsStillLocked(
+        address _depositor,
+        uint256 _amount,
+        uint256 _jump
+    ) public {
+        _amount = _boundToRealisticStake(_amount);
+        _jump = bound(_jump, vm.getBlockTimestamp(), LOCK_PERIOD);
+
+        _mintGovToken(_depositor, _amount);
+        _setKeyper(_depositor, true);
+
+        uint256 stakeIndex = _stake(_depositor, _amount);
+
+        _jumpAhead(_jump);
+
+        vm.prank(_depositor);
+        vm.expectRevert(Staking.StakeIsStillLocked.selector);
+        staking.unstake(_depositor, stakeIndex, 0);
+    }
+
+    function testFuzz_RevertIf_StakeIsStillLockedAfterLockPeriodChangedToLessThanCurrent(
+        address _depositor,
+        uint256 _amount,
+        uint256 _jump
+    ) public {
+        _amount = _boundToRealisticStake(_amount);
+        _jump = bound(_jump, vm.getBlockTimestamp(), LOCK_PERIOD);
+
+        _mintGovToken(_depositor, _amount);
+        _setKeyper(_depositor, true);
+
+        uint256 stakeIndex = _stake(_depositor, _amount);
+
+        staking.setLockPeriod(_jump);
+        _jumpAhead(_jump - 1);
+
+        vm.prank(_depositor);
+        vm.expectRevert(Staking.StakeIsStillLocked.selector);
+        staking.unstake(_depositor, stakeIndex, 0);
+    }
+
     function testFuzz_RevertIf_StakeDoesNotBelongToKeyper(
         address _depositor1,
         address _depositor2,