Merge pull request #25 from BlockchainLabsNZ/vesting_contract

Feedback related to Vesting
BlockchainLabsNZ · Nov 30, 2017 · 7d46986 · 7d46986
2 parents 2cafb6c + 82d644f
commit 7d46986
Show file tree

Hide file tree

Showing 6 changed files with 230 additions and 7 deletions.
diff --git a/.solcover.js b/.solcover.js
@@ -1,4 +1,4 @@
 module.exports = {
-    skipFiles: ['lifecycle/Migrations.sol'],
-    copyNodeModules: true
-}
+  skipFiles: ["lifecycle/Migrations.sol"],
+  copyNodeModules: true
+};
diff --git a/contracts/PolyMathVesting.sol b/contracts/PolyMathVesting.sol
@@ -9,6 +9,7 @@ contract PolyMathVesting is Ownable {
 
   // Contract holds ERC20 POLY tokens.
   ERC20Basic token;
+  uint256 vestingAmount;
 
   // Important dates for vesting contract.
   uint256 public startTime;
@@ -36,13 +37,16 @@ contract PolyMathVesting is Ownable {
       uint256 _startTime,
       uint256 _cliffTime,
       uint256 _releaseTime,
-      uint256 _period
+      uint256 _period,
+      uint256 _vestingAmount
   ) public {
     require(_token != 0x0);
     require(_startTime > getBlockTimestamp());
     require(_cliffTime >= _startTime);
     require(_releaseTime >= _cliffTime);
+    require(_vestingAmount > 0);
     token = ERC20Basic(_token);
+    vestingAmount = _vestingAmount;
     startTime = _startTime;
     cliffTime = _cliffTime;
     releaseTime = _releaseTime;
@@ -62,12 +66,22 @@ contract PolyMathVesting is Ownable {
 
   function allocate(address _holder, uint256 _amount) public onlyOwner {
     require(!allocationFinished);
+    require(_holder != 0x0);
     allocated = allocated.sub(allocations[_holder]).add(_amount);
     require(allocated <= token.balanceOf(address(this)));
-    allocationFinished = allocated == token.balanceOf(address(this));
+
+    if (allocated == vestingAmount) {
+      finishAllocation();
+    }
+
     allocations[_holder] = _amount;
   }
 
+  function finishAllocation() public onlyOwner {
+    vestingAmount = allocated;
+    allocationFinished = true;
+  }
+
   function collect() public {
     require(getBlockTimestamp() >= cliffTime);
     uint256 balance = allocations[msg.sender];
@@ -96,4 +110,23 @@ contract PolyMathVesting is Ownable {
   function getBlockTimestamp() internal view returns (uint256) {
     return block.timestamp;
   }
+
+  event TokensClaimed(address indexed token, uint256 _amount);
+
+  function claimTokens(address _token) public onlyOwner {
+    if (_token == 0x0) {
+      owner.transfer(this.balance);
+      return;
+    }
+
+    ERC20Basic refundToken = ERC20Basic(_token);
+    uint256 balance = refundToken.balanceOf(this);
+
+    if (refundToken == token) {
+      balance = balance.sub(allocated);
+    }
+
+    refundToken.transfer(owner, balance);
+    TokensClaimed(_token, balance);
+  }
 }
diff --git a/test/helpers/PolyMathVestingMock.sol b/test/helpers/PolyMathVestingMock.sol
@@ -18,8 +18,16 @@ contract PolyMathVestingMock is PolyMathVesting {
       uint256 _startTime,
       uint256 _cliffTime,
       uint256 _releaseTime,
-      uint256 _period
-  ) PolyMathVesting(_token, _startTime, _cliffTime, _releaseTime, _period) {
+      uint256 _period,
+      uint256 _totalTokens
+  ) PolyMathVesting(
+      _token,
+      _startTime,
+      _cliffTime,
+      _releaseTime,
+      _period,
+      _totalTokens
+  ) {
     timeStamp = now;
   }
 }
diff --git a/test/vestingAdvisorsTests.js b/test/vestingAdvisorsTests.js
@@ -32,6 +32,7 @@ contract("PolyMathVesting", async function(
         cliffTime,
         releaseTime,
         period,
+        new BigNumber(25000000).mul(10 ** 18),
         { from: owner }
       );
 
@@ -58,6 +59,14 @@ contract("PolyMathVesting", async function(
         );
       });
 
+      await assertFail(async () => {
+        await polyVestingDeployed.allocate(
+          "0x0",
+          new BigNumber(15000000).mul(10 ** 18),
+          { from: owner }
+        );
+      });
+
       await polyVestingDeployed.allocate(
         advisor2,
         new BigNumber(15000000).mul(10 ** 18),
@@ -139,5 +148,57 @@ contract("PolyMathVesting", async function(
         15000000 * 10 ** 18
       );
     });
+
+    it("Admin can Claim Token excess of tokens", async () => {
+      assert.isFalse(await polyVestingDeployed.allocationFinished());
+
+      await polyVestingDeployed.allocateArray(
+        [advisor1, advisor2],
+        [
+          new BigNumber(10000000).mul(10 ** 18),
+          new BigNumber(10000000).mul(10 ** 18)
+        ],
+        { from: owner }
+      );
+      assert.equal(
+        (await polyVestingDeployed.allocations(advisor1)).toNumber(),
+        10000000 * 10 ** 18
+      );
+      assert.equal(
+        (await polyVestingDeployed.allocations(advisor2)).toNumber(),
+        10000000 * 10 ** 18
+      );
+      await assertFail(async () => {
+        await polyVestingDeployed.finishAllocation();
+      });
+      await polyVestingDeployed.finishAllocation({ from: owner });
+      assert.isTrue(await polyVestingDeployed.allocationFinished());
+      await assertFail(async () => {
+        await polyVestingDeployed.allocate(
+          notAdvisor,
+          new BigNumber(5000000).mul(10 ** 18),
+          { from: owner }
+        );
+      });
+
+      await assertFail(async () => {
+        await polyVestingDeployed.claimTokens(tokenDeployed.address);
+      });
+
+      assert.equal((await tokenDeployed.balanceOf(owner)).toNumber(), 0);
+      await polyVestingDeployed.claimTokens(tokenDeployed.address, {
+        from: owner
+      });
+
+      assert.equal(
+        (await tokenDeployed.balanceOf(owner)).toNumber(),
+        5000000 * 10 ** 18
+      );
+
+      assert.equal(
+        (await tokenDeployed.balanceOf(polyVestingDeployed.address)).toNumber(),
+        20000000 * 10 ** 18
+      );
+    });
   });
 });
diff --git a/test/vestingFoundersTests.js b/test/vestingFoundersTests.js
@@ -33,6 +33,7 @@ contract("PolyMathVesting", async function(
         cliffTime,
         releaseTime,
         period,
+        new BigNumber(150000000).mul(10 ** 18),
         { from: owner }
       );
 
@@ -59,6 +60,14 @@ contract("PolyMathVesting", async function(
         );
       });
 
+      await assertFail(async () => {
+        await polyVestingDeployed.allocate(
+          "0x0",
+          new BigNumber(50000000).mul(10 ** 18),
+          { from: owner }
+        );
+      });
+
       await polyVestingDeployed.allocate(
         founder2,
         new BigNumber(50000000).mul(10 ** 18),
@@ -185,5 +194,56 @@ contract("PolyMathVesting", async function(
         48000000 * 10 ** 18
       );
     });
+
+    it("Admin can Claim Token excess of tokens", async () => {
+      assert.isFalse(await polyVestingDeployed.allocationFinished());
+      await polyVestingDeployed.allocateArray(
+        [founder1, founder2],
+        [
+          new BigNumber(100000000).mul(10 ** 18),
+          new BigNumber(40000000).mul(10 ** 18)
+        ],
+        { from: owner }
+      );
+      assert.equal(
+        (await polyVestingDeployed.allocations(founder1)).toNumber(),
+        100000000 * 10 ** 18
+      );
+      assert.equal(
+        (await polyVestingDeployed.allocations(founder2)).toNumber(),
+        40000000 * 10 ** 18
+      );
+      await assertFail(async () => {
+        await polyVestingDeployed.finishAllocation();
+      });
+      await polyVestingDeployed.finishAllocation({ from: owner });
+      assert.isTrue(await polyVestingDeployed.allocationFinished());
+      await assertFail(async () => {
+        await polyVestingDeployed.allocate(
+          notFounder,
+          new BigNumber(10000000).mul(10 ** 18),
+          { from: owner }
+        );
+      });
+
+      await assertFail(async () => {
+        await polyVestingDeployed.claimTokens(tokenDeployed.address);
+      });
+
+      assert.equal((await tokenDeployed.balanceOf(owner)).toNumber(), 0);
+      await polyVestingDeployed.claimTokens(tokenDeployed.address, {
+        from: owner
+      });
+
+      assert.equal(
+        (await tokenDeployed.balanceOf(owner)).toNumber(),
+        10000000 * 10 ** 18
+      );
+
+      assert.equal(
+        (await tokenDeployed.balanceOf(polyVestingDeployed.address)).toNumber(),
+        140000000 * 10 ** 18
+      );
+    });
   });
 });
diff --git a/test/vestingReserveTests.js b/test/vestingReserveTests.js
@@ -33,6 +33,7 @@ contract("PolyMathVesting", async function(
         cliffTime,
         releaseTime,
         period,
+        new BigNumber(450000000).mul(10 ** 18),
         { from: owner }
       );
 
@@ -44,6 +45,22 @@ contract("PolyMathVesting", async function(
 
     it("Owner of the vesting contract can allocate tokens to an address", async () => {
       assert.isFalse(await polyVestingDeployed.allocationFinished());
+
+      await assertFail(async () => {
+        await polyVestingDeployed.allocate(
+          reserve,
+          new BigNumber(450000000).mul(10 ** 18)
+        );
+      });
+
+      await assertFail(async () => {
+        await polyVestingDeployed.allocate(
+          "0x0",
+          new BigNumber(450000000).mul(10 ** 18),
+          { from: owner }
+        );
+      });
+
       await polyVestingDeployed.allocate(
         reserve,
         new BigNumber(450000000).mul(10 ** 18),
@@ -147,5 +164,49 @@ contract("PolyMathVesting", async function(
         450000000 * 10 ** 18
       );
     });
+
+    it("Admin can Claim Token excess of tokens", async () => {
+      assert.isFalse(await polyVestingDeployed.allocationFinished());
+      await polyVestingDeployed.allocate(
+        reserve,
+        new BigNumber(400000000).mul(10 ** 18),
+        { from: owner }
+      );
+      assert.equal(
+        (await polyVestingDeployed.allocations(reserve)).toNumber(),
+        400000000 * 10 ** 18
+      );
+      await assertFail(async () => {
+        await polyVestingDeployed.finishAllocation();
+      });
+      await polyVestingDeployed.finishAllocation({ from: owner });
+      assert.isTrue(await polyVestingDeployed.allocationFinished());
+      await assertFail(async () => {
+        await polyVestingDeployed.allocate(
+          notReserve,
+          new BigNumber(50000000).mul(10 ** 18),
+          { from: owner }
+        );
+      });
+
+      await assertFail(async () => {
+        await polyVestingDeployed.claimTokens(tokenDeployed.address);
+      });
+
+      assert.equal((await tokenDeployed.balanceOf(owner)).toNumber(), 0);
+      await polyVestingDeployed.claimTokens(tokenDeployed.address, {
+        from: owner
+      });
+
+      assert.equal(
+        (await tokenDeployed.balanceOf(owner)).toNumber(),
+        50000000 * 10 ** 18
+      );
+
+      assert.equal(
+        (await tokenDeployed.balanceOf(polyVestingDeployed.address)).toNumber(),
+        400000000 * 10 ** 18
+      );
+    });
   });
 });