TokenGrant is able to compute withdrawableAmount and withdraw it

contracts/grant/GrantStakingPolicy.sol

@@ -9,7 +9,8 @@ interface IGrantStakingPolicy {
         uint256 duration,
         uint256 start,
         uint256 cliff,
-        uint256 withdrawn) external view returns (uint256);
+        uint256 withdrawn
+    ) external view returns (uint256);
 }
 
-// TODO: add more policies
+// TODO: add more policies

contracts/grant/TokenGrant.sol

@@ -3,26 +3,36 @@
 pragma solidity 0.8.4;
 
 import "./GrantStakingPolicy.sol";
+import "../token/T.sol";
+
+import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 
 contract TokenGrant {
+    using SafeERC20 for T;
+
+    T public token;
 
     address public grantee;
     bool public revocable;
     uint256 public amount;
     uint256 public duration;
     uint256 public start;
     uint256 public cliff;
-    
+
     uint256 public withdrawn;
     uint256 public staked;
 
-    uint256 public revokedAt;
-    uint256 public revokedAmount;
-    uint256 public revokedWithdrawn;
-
     IGrantStakingPolicy public stakingPolicy;
 
+    event Withdrawn(uint256 amount);
+
+    modifier onlyGrantee() {
+        require(msg.sender == grantee, "Not authorized");
+        _;
+    }
+
     function initialize(
+        T _token,
         address _grantee,
         bool _revocable,
         uint256 _amount,
@@ -31,29 +41,62 @@ contract TokenGrant {
         uint256 _cliff,
         IGrantStakingPolicy _stakingPolicy
     ) public {
+        token = _token;
         grantee = _grantee;
         revocable = _revocable;
         amount = _amount;
         duration = _duration;
         start = _start;
         cliff = _cliff;
         stakingPolicy = _stakingPolicy;
+
+        token.safeTransferFrom(msg.sender, address(this), amount);
+    }
+
+    function stake(uint256 amountToStake) external onlyGrantee {
+        staked += amountToStake;
+
+        // TODO: implement
+    }
+
+    function withdraw() external onlyGrantee {
+        uint256 withdrawable = withdrawableAmount();
+        require(withdrawable > 0, "There is nothing to withdraw");
+
+        emit Withdrawn(withdrawable);
+        withdrawn += withdrawable;
+        token.safeTransfer(grantee, withdrawable);
     }
 
     function unlockedAmount() public view returns (uint256) {
-        if (block.timestamp < start) { // start reached?
+        /* solhint-disable-next-line not-rely-on-time */
+        if (block.timestamp < start) {
             return 0;
         }
 
-        if (block.timestamp < cliff) { // cliff reached?
-            return 0; 
+        /* solhint-disable-next-line not-rely-on-time */
+        if (block.timestamp < cliff) {
+            return 0;
         }
 
+        /* solhint-disable-next-line not-rely-on-time */
         uint256 timeElapsed = block.timestamp - start;
 
         bool unlockingPeriodFinished = timeElapsed >= duration;
-        if (unlockingPeriodFinished) { return amount; }
+        if (unlockingPeriodFinished) {
+            return amount;
+        }
 
-        return amount * timeElapsed / duration;
+        return (amount * timeElapsed) / duration;
+    }
+
+    function withdrawableAmount() public view returns (uint256) {
+        uint256 unlocked = unlockedAmount();
+
+        if (withdrawn + staked >= unlocked) {
+            return 0;
+        } else {
+            return unlocked - withdrawn - staked;
+        }
     }
-}
+}

test/grant/TokenGrant.test.js

@@ -3,13 +3,20 @@ const {
   ZERO_ADDRESS,
   to1e18,
   lastBlockTime,
+  increaseTime,
+  pastEvents,
 } = require("../helpers/contract-test-helpers")
 
 describe("TokenGrant", () => {
+  let token
   let grantee
 
   beforeEach(async () => {
-    ;[grantee] = await ethers.getSigners()
+    ;[deployer, thirdParty, grantee] = await ethers.getSigners()
+
+    const T = await ethers.getContractFactory("T")
+    token = await T.deploy()
+    await token.deployed()
   })
 
   describe("unlockedAmount", () => {
@@ -27,7 +34,7 @@ describe("TokenGrant", () => {
     context("before the schedule start", () => {
       it("should return no tokens unlocked", async () => {
         const start = now + 60
-        const cliff = now
+        const cliff = now + 60
         const grant = await createGrant(false, amount, duration, start, cliff)
         expect(await grant.unlockedAmount()).to.equal(0)
       })
@@ -101,22 +108,219 @@ describe("TokenGrant", () => {
         )
       })
     })
+
+    context("when some tokens are staked", () => {
+      it("should return token amount unlocked so far", async () => {
+        const start = now - 3600 // one hour earlier
+        const cliff = now - 3600
+        const grant = await createGrant(false, amount, duration, start, cliff)
+        await grant.connect(grantee).withdraw()
+        expect(await grant.unlockedAmount()).is.closeTo(
+          to1e18(23), // 3600 / 15552000 * 100k = ~23 tokens
+          assertionPrecision
+        )
+      })
+    })
+
+    context("when some tokens were withdrawn", () => {
+      it("should return token amount unlocked so far", async () => {
+        const start = now - 3600 // one hour earlier
+        const cliff = now - 3600
+        const grant = await createGrant(false, amount, duration, start, cliff)
+        await grant.connect(grantee).stake(to1e18(20))
+        expect(await grant.unlockedAmount()).is.closeTo(
+          to1e18(23), // 3600 / 15552000 * 100k = ~23 tokens
+          assertionPrecision
+        )
+      })
+    })
+  })
+
+  describe("withdrawableAmount", () => {
+    const assertionPrecision = to1e18(1) // +- 1 token
+
+    const amount = to1e18(100000) // 100k tokens
+    const duration = 15552000 // 180 days
+
+    let grant
+
+    beforeEach(async () => {
+      const now = await lastBlockTime()
+      const start = now - 7200 // two hours earlier
+      const cliff = now - 7200
+      grant = await createGrant(false, amount, duration, start, cliff)
+    })
+
+    context("when no tokens were staked or withdrawn", () => {
+      it("should return tokens unlocked so far", async () => {
+        expect(await grant.withdrawableAmount()).is.closeTo(
+          to1e18(46), // 7200 / 15552000 * 100k = ~46 tokens
+          assertionPrecision
+        )
+      })
+    })
+
+    context("when some tokens were staked", () => {
+      it("should return tokens unlocked so far minus staked tokens", async () => {
+        await grant.connect(grantee).stake(to1e18(5))
+        expect(await grant.withdrawableAmount()).is.closeTo(
+          to1e18(41), // 7200 / 15552000 * 100k - 5 = ~41 tokens
+          assertionPrecision
+        )
+      })
+    })
+
+    context("when some tokens were withdrawn", () => {
+      it("should return tokens unlocked so far minus withdrawn tokens", async () => {
+        await grant.connect(grantee).withdraw()
+        await increaseTime(3600)
+        expect(await grant.withdrawableAmount()).is.closeTo(
+          to1e18(23), // 3600 / 15552000 * 100k = ~23 tokens
+          assertionPrecision
+        )
+      })
+    })
+
+    context("when tokens were withdrawn multiple times", () => {
+      it("should return tokens unlocked so far minus withdrawn tokens", async () => {
+        await grant.connect(grantee).withdraw()
+        await increaseTime(7200)
+        await grant.connect(grantee).withdraw()
+        await increaseTime(7200)
+
+        expect(await grant.withdrawableAmount()).is.closeTo(
+          to1e18(46), // 7200 / 15552000 * 100k = ~46 tokens
+          assertionPrecision
+        )
+      })
+    })
+
+    context("when tokens were staked and withdrawn", () => {
+      it("should return tokens unlocked so far minus withdrawn and staked tokens", async () => {
+        await grant.connect(grantee).withdraw()
+        await increaseTime(7200)
+        await grant.connect(grantee).stake(to1e18(20))
+        expect(await grant.withdrawableAmount()).is.closeTo(
+          to1e18(26), // 7200 / 15552000 * 100k - 20 = ~26 tokens
+          assertionPrecision
+        )
+      })
+    })
+
+    context("when tokens were staked and withdrawn multiple times", () => {
+      it("should return tokens unlocked so far minus withdrawn and staked tokens", async () => {
+        await grant.connect(grantee).withdraw()
+        await increaseTime(7200)
+        await grant.connect(grantee).withdraw()
+        await grant.connect(grantee).stake(to1e18(10))
+        await increaseTime(7200)
+        expect(await grant.withdrawableAmount()).is.closeTo(
+          to1e18(36), // 7200 / 15552000 * 100k - 10 = ~36 tokens
+          assertionPrecision
+        )
+      })
+    })
+  })
+
+  describe("withdraw", () => {
+    const assertionPrecision = to1e18(1) // +- 1 token
+
+    const amount = to1e18(200000) // 200k tokens
+    const duration = 7776000 // 90 days
+
+    let grant
+
+    beforeEach(async () => {
+      const now = await lastBlockTime()
+      const start = now - 3888000 // 45 days earlier
+      const cliff = now - 3888000
+      grant = await createGrant(false, amount, duration, start, cliff)
+    })
+
+    context("when called by a third party", () => {
+      it("should revert", async () => {
+        await expect(grant.connect(thirdParty).withdraw()).to.be.revertedWith(
+          "Not authorized"
+        )
+      })
+    })
+
+    context("when called by a grant creator", () => {
+      it("should revert", async () => {
+        await expect(grant.connect(deployer).withdraw()).to.be.revertedWith(
+          "Not authorized"
+        )
+      })
+    })
+
+    context("when called by grantee", () => {
+      context("when there are no withdrawable tokens", () => {
+        it("should revert", async () => {
+          await grant.connect(grantee).stake(amount)
+          await expect(grant.connect(grantee).withdraw()).to.be.revertedWith(
+            "There is nothing to withdraw"
+          )
+        })
+      })
+
+      context("when there are withdrawable tokens", () => {
+        let tx
+
+        beforeEach(async () => {
+          // 3888000/7776000 * 200k = 100k
+          tx = await grant.connect(grantee).withdraw()
+        })
+
+        it("should increase withdrawn amount", async () => {
+          expect(await grant.withdrawn()).to.be.closeTo(
+            to1e18(100000),
+            assertionPrecision
+          )
+        })
+
+        it("should transfer tokens to grantee", async () => {
+          expect(await token.balanceOf(grantee.address)).to.be.closeTo(
+            to1e18(100000),
+            assertionPrecision
+          )
+          expect(await token.balanceOf(grant.address)).to.be.closeTo(
+            to1e18(100000),
+            assertionPrecision
+          )
+        })
+
+        it("should emit Withdrawn event", async () => {
+          const events = pastEvents(await tx.wait(), grant, "Withdrawn")
+          expect(events.length).to.equal(1)
+          expect(events[0].args["amount"]).to.be.closeTo(
+            to1e18(100000),
+            assertionPrecision
+          )
+        })
+      })
+    })
   })
 
   async function createGrant(revocable, amount, duration, start, cliff) {
     const TokenGrant = await ethers.getContractFactory("TokenGrant")
     const tokenGrant = await TokenGrant.deploy()
     await tokenGrant.deployed()
 
-    await tokenGrant.initialize(
-      grantee.address,
-      revocable,
-      amount,
-      duration,
-      start,
-      cliff,
-      ZERO_ADDRESS
-    )
+    await token.connect(deployer).mint(deployer.address, amount)
+    await token.connect(deployer).approve(tokenGrant.address, amount)
+
+    await tokenGrant
+      .connect(deployer)
+      .initialize(
+        token.address,
+        grantee.address,
+        revocable,
+        amount,
+        duration,
+        start,
+        cliff,
+        ZERO_ADDRESS
+      )
 
     return tokenGrant
   }