Added a check in available function to return user's staked balance i…

source/staking/contracts/Staking.sol

@@ -4,6 +4,7 @@ pragma solidity ^0.8.17;
 import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "@openzeppelin/contracts/access/Ownable.sol";
+import "@openzeppelin/contracts/utils/math/Math.sol";
 import "./interfaces/IStaking.sol";
 
 /**
@@ -227,13 +228,17 @@ contract Staking is IStaking, Ownable {
    */
   function available(address _account) public view override returns (uint256) {
     Stake storage _selected = stakes[_account];
-    uint256 _available = (_selected.balance *
-      (block.timestamp - _selected.timestamp)) /
-      (_selected.maturity - _selected.timestamp);
-    if (_available >= _selected.balance) {
+    if (_selected.maturity == _selected.timestamp) {
       return _selected.balance;
     } else {
-      return _available;
+      uint256 _available = (_selected.balance *
+        (block.timestamp - _selected.timestamp)) /
+        (_selected.maturity - _selected.timestamp);
+      if (_available >= _selected.balance) {
+        return _selected.balance;
+      } else {
+        return _available;
+      }
     }
   }
 
@@ -276,10 +281,12 @@ contract Staking is IStaking, Ownable {
     if (_amount > available(_account)) revert AmountInvalid(_amount);
     uint256 nowAvailable = available(_account);
     _selected.balance = _selected.balance - _amount;
-    _selected.timestamp =
+    _selected.timestamp = Math.min(
       block.timestamp -
-      (((10000 - ((10000 * _amount) / nowAvailable)) *
-        (block.timestamp - _selected.timestamp)) / 10000);
+        (((10000 - ((10000 * _amount) / nowAvailable)) *
+          (block.timestamp - _selected.timestamp)) / 10000),
+      _selected.maturity
+    );
     stakingToken.safeTransfer(_account, _amount);
     emit Transfer(_account, address(0), _amount);
   }

source/staking/test/Staking.js

@@ -339,6 +339,109 @@ describe('Staking Unit', () => {
       expect(userStakes.balance).to.equal(90)
     })
 
+    it('successful successive maximum unstaking', async () => {
+      await token.mock.transferFrom.returns(true)
+      await token.mock.transfer.returns(true)
+      await staking.connect(account1).stake('100')
+
+      // move 10 seconds forward - 10% unstakeable
+      await ethers.provider.send('evm_increaseTime', [10])
+      await ethers.provider.send('evm_mine')
+
+      await staking.connect(account1).unstake('11')
+      const userStakes = await staking
+        .connect(account1)
+        .getStakes(account1.address)
+
+      expect(userStakes.balance).to.equal(89)
+
+      // move 20 seconds forward - 20% unstakeable
+      await ethers.provider.send('evm_increaseTime', [20])
+      await ethers.provider.send('evm_mine')
+      await staking.connect(account1).unstake('21')
+      const userStakes2 = await staking
+        .connect(account1)
+        .getStakes(account1.address)
+
+      expect(userStakes2.balance).to.equal(68)
+    })
+
+    it('successful partial unstaking after stake maturity', async () => {
+      await token.mock.transferFrom.returns(true)
+      await token.mock.transfer.returns(true)
+      await staking.connect(account1).stake('100')
+
+      // move 1000 seconds forward - 100% unstakeable
+      await ethers.provider.send('evm_increaseTime', [999])
+      await ethers.provider.send('evm_mine')
+
+      await staking.connect(account1).unstake('10')
+      const userStakes = await staking
+        .connect(account1)
+        .getStakes(account1.address)
+
+      expect(userStakes.balance).to.equal(90)
+
+      await staking.connect(account1).unstake('10')
+      const userStakes2 = await staking
+        .connect(account1)
+        .getStakes(account1.address)
+
+      expect(userStakes2.balance).to.equal(80)
+    })
+
+    it('successful unstake after partial unstake and restake', async () => {
+      await token.mock.transferFrom.returns(true)
+      await token.mock.transfer.returns(true)
+      await staking.connect(account1).stake('100')
+
+      // move 1000 seconds forward - 100% unstakeable
+      await ethers.provider.send('evm_increaseTime', [1000])
+      await ethers.provider.send('evm_mine')
+
+      // Unstake 10
+      await staking.connect(account1).unstake('10')
+      const userStakes = await staking
+        .connect(account1)
+        .getStakes(account1.address)
+
+      expect(userStakes.balance).to.equal(90)
+
+      // Stake 100
+      await staking.connect(account1).stake('100')
+
+      const userStakes2 = await staking
+        .connect(account1)
+        .getStakes(account1.address)
+
+      expect(userStakes2.balance).to.equal(190)
+
+      // Unstake remainder of unstakable amount
+      await staking.connect(account1).unstake('90')
+
+      const userStakes3 = await staking
+        .connect(account1)
+        .getStakes(account1.address)
+
+      expect(userStakes3.balance).to.equal(100)
+
+      // Cannot unstake the new amount immediately
+      await expect(staking.connect(account1).unstake('10'))
+        .to.be.revertedWith('AmountInvalid')
+        .withArgs(10)
+
+      // move 10 seconds forward - 10% unstakeable
+      await ethers.provider.send('evm_increaseTime', [10])
+      await ethers.provider.send('evm_mine')
+
+      await staking.connect(account1).unstake('10')
+      const userStakes4 = await staking
+        .connect(account1)
+        .getStakes(account1.address)
+
+      expect(userStakes4.balance).to.equal(90)
+    })
+
     it('successful extended stake and successful unstaking', async () => {
       await token.mock.transferFrom.returns(true)
       await token.mock.transfer.returns(true)