From 2287d3c0a5ca30e42247e30ac1ccbc2158784e8c 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 | 16 +++++++++++
 core/test/staking/TokenStaking.test.ts  | 38 +++++++++++++++++++++++++
 2 files changed, 54 insertions(+)

diff --git a/core/contracts/staking/TokenStaking.sol b/core/contracts/staking/TokenStaking.sol
index a9051040d..a6b61dd5e 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 account, uint256 amount);
+    event Unstaked(address indexed account, uint256 amount);
 
     IERC20 internal immutable token;
 
@@ -51,6 +52,21 @@ contract TokenStaking is IReceiveApproval {
         _stake(msg.sender, amount);
     }
 
+    /// @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 account, uint256 amount) private {
         require(amount > 0, "Amount is less than minimum");
         require(account != address(0), "Can not be the zero address");
diff --git a/core/test/staking/TokenStaking.test.ts b/core/test/staking/TokenStaking.test.ts
index 076313950..5e7f305e2 100644
--- a/core/test/staking/TokenStaking.test.ts
+++ b/core/test/staking/TokenStaking.test.ts
@@ -86,4 +86,42 @@ describe("TokenStaking", () => {
       })
     })
   })
+
+  describe("unstaking", () => {
+    beforeEach(async () => {
+      const tokenBalance = await token.balanceOf(await tokenHolder.getAddress())
+
+      // Stake tokens.
+      await token
+        .connect(tokenHolder)
+        .approveAndCall(await tokenStaking.getAddress(), tokenBalance, "0x")
+    })
+
+    it("should unstake tokens", async () => {
+      const staker = await tokenHolder.getAddress()
+      const stakingBalance = await tokenStaking.balanceOf(staker)
+
+      await expect(tokenStaking.connect(tokenHolder).unstake(stakingBalance))
+        .to.emit(tokenStaking, "Unstaked")
+        .withArgs(staker, stakingBalance)
+
+      expect(await token.balanceOf(staker)).to.be.equal(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")
+    })
+  })
 })