From 0cbfb52bb0d39b7822fb7ae9abcadc3194324ae6 Mon Sep 17 00:00:00 2001
From: Rafal Czajkowski <rafal.czajkowski@keep.network>
Date: Thu, 26 Oct 2023 19:05:18 +0200
Subject: [PATCH] Add initial impl of the unstaking

Add `unstake` function to the `TokenStaking` contract - this function
reduces stake amount by the provided amount and withdraws tokens to the
owner.
---
 core/contracts/staking/TokenStaking.sol | 18 ++++++++++-
 core/contracts/test/TestToken.sol       |  2 +-
 core/test/staking/TokenStaking.test.ts  | 41 +++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/core/contracts/staking/TokenStaking.sol b/core/contracts/staking/TokenStaking.sol
index 86a583b04..cb48a5196 100644
--- a/core/contracts/staking/TokenStaking.sol
+++ b/core/contracts/staking/TokenStaking.sol
@@ -14,6 +14,7 @@ contract TokenStaking is IReceiveApproval {
     using SafeERC20 for IERC20;
 
     event Staked(address indexed staker, uint256 amount);
+    event Unstaked(address indexed staker, uint256 amount);
 
     IERC20 internal immutable token;
 
@@ -58,7 +59,7 @@ contract TokenStaking is IReceiveApproval {
         // governable params.
         return 1;
     }
-
+    
     /// @notice Returns maximum amount of staking tokens.
     function maximumStake() public pure returns (uint256) {
         // TODO: Fetch this param from "parameters" contract that stores
@@ -66,6 +67,21 @@ contract TokenStaking is IReceiveApproval {
         return 100 ether;
     }
 
+    /// @notice Reduces stake amount by the provided amount and
+    ///         withdraws tokens to the owner.
+    /// @param amount Amount to unstake and withdraw.
+    function unstake(uint256 amount) external {
+        require((amount > 0), "Amount can not be zero");
+        
+        uint256 balance = balanceOf[msg.sender];
+        require(balance >= amount, "Insufficient funds");
+
+        balanceOf[msg.sender] -= amount;
+
+        emit Unstaked(msg.sender, amount);
+        token.safeTransfer(msg.sender, amount);
+    }
+
     function _stake(address staker, uint256 amount) private {
         require(amount >= minimumStake(), "Amount is less than minimum");
         require(amount <= maximumStake(), "Amount is greater than maxium");
diff --git a/core/contracts/test/TestToken.sol b/core/contracts/test/TestToken.sol
index c061cf278..57832b4f1 100644
--- a/core/contracts/test/TestToken.sol
+++ b/core/contracts/test/TestToken.sol
@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: GPL-3.0-or-later
 
-pragma solidity 0.8.20;
+pragma solidity ^0.8.20;
 
 import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
 import "../shared/IReceiveApproval.sol";
diff --git a/core/test/staking/TokenStaking.test.ts b/core/test/staking/TokenStaking.test.ts
index 157a5fc31..07fa1ce5f 100644
--- a/core/test/staking/TokenStaking.test.ts
+++ b/core/test/staking/TokenStaking.test.ts
@@ -102,4 +102,45 @@ describe("TokenStaking", () => {
       })
     })
   })
+
+  describe("unstaking", () => {
+    const amountToStake = WeiPerEther * 10n
+
+    beforeEach(async () => {
+      // Stake tokens.
+      await token
+        .connect(tokenHolder)
+        .approveAndCall(await tokenStaking.getAddress(), amountToStake, "0x")
+    })
+
+    it("should unstake tokens", async () => {
+      const staker = await tokenHolder.getAddress()
+      const stakingBalance = await tokenStaking.balanceOf(staker)
+      const balanceBeforeUnstaking = await token.balanceOf(staker)
+
+      await expect(tokenStaking.connect(tokenHolder).unstake(stakingBalance))
+        .to.emit(tokenStaking, "Unstaked")
+        .withArgs(staker, stakingBalance)
+
+      expect(await token.balanceOf(staker)).to.be.equal(
+        balanceBeforeUnstaking + stakingBalance,
+      )
+      expect(await tokenStaking.balanceOf(staker)).to.be.eq(0)
+    })
+
+    it("should revert if the unstaked amount is equal 0", async () => {
+      await expect(
+        tokenStaking.connect(tokenHolder).unstake(0),
+      ).to.be.revertedWith("Amount can not be zero")
+    })
+
+    it("should revert if the user wants to unstake more tokens than currently staked", async () => {
+      const staker = await tokenHolder.getAddress()
+      const stakingBalance = await tokenStaking.balanceOf(staker)
+
+      await expect(
+        tokenStaking.connect(tokenHolder).unstake(stakingBalance + 10n),
+      ).to.be.revertedWith("Insufficient funds")
+    })
+  })
 })