From 0e59f643b803cdd7de1fe0fe3caefea614c64dba Mon Sep 17 00:00:00 2001
From: owen05 <liuhongnan53@163.com>
Date: Fri, 9 Apr 2021 20:36:25 +0800
Subject: [PATCH] update

---
 contracts/DODOToken/DODOMineV2/BaseMine.sol   | 14 +++++-
 .../DODOToken/DODOMineV2/RewardVault.sol      |  6 +++
 test/DODOMineV2/erc20Mine.test.ts             | 46 +++++++++----------
 truffle-test.sh                               |  2 +-
 4 files changed, 43 insertions(+), 25 deletions(-)

diff --git a/contracts/DODOToken/DODOMineV2/BaseMine.sol b/contracts/DODOToken/DODOMineV2/BaseMine.sol
index a618038..7a91342 100644
--- a/contracts/DODOToken/DODOMineV2/BaseMine.sol
+++ b/contracts/DODOToken/DODOMineV2/BaseMine.sol
@@ -44,7 +44,7 @@ contract BaseMine is InitializableOwnable {
     event UpdateEndBlock(uint256 indexed i, uint256 endBlock);
     event NewRewardToken(uint256 indexed i, address rewardToken);
     event RemoveRewardToken(address rewardToken);
-
+    event WithdrawLeftOver(address owner, uint256 i);
 
     // ============ View  ============
 
@@ -176,10 +176,22 @@ contract BaseMine is InitializableOwnable {
     {
         _updateReward(address(0), i);
         RewardTokenInfo storage rt = rewardTokenInfos[i];
+        
+        require(block.number < rt.endBlock, "DODOMineV2: ALREADY_CLOSE");
+
         rt.rewardPerBlock = newRewardPerBlock;
         emit UpdateReward(i, newRewardPerBlock);
     }
 
+    function withdrawLeftOver(uint256 i) external onlyOwner {
+        RewardTokenInfo storage rt = rewardTokenInfos[i];
+        require(block.number > rt.endBlock, "DODOMineV2: MINING_NOT_FINISHED");
+
+        IRewardVault(rt.rewardVault).withdrawLeftOver(msg.sender);
+
+        emit WithdrawLeftOver(msg.sender, i);
+    }
+
 
     // ============ Internal  ============
 
diff --git a/contracts/DODOToken/DODOMineV2/RewardVault.sol b/contracts/DODOToken/DODOMineV2/RewardVault.sol
index 33bacf8..84dc10d 100644
--- a/contracts/DODOToken/DODOMineV2/RewardVault.sol
+++ b/contracts/DODOToken/DODOMineV2/RewardVault.sol
@@ -14,6 +14,7 @@ import {IERC20} from "../../intf/IERC20.sol";
 
 interface IRewardVault {
     function reward(address to, uint256 amount) external;
+    function withdrawLeftOver(address to) external; 
 }
 
 contract RewardVault is Ownable {
@@ -28,4 +29,9 @@ contract RewardVault is Ownable {
     function reward(address to, uint256 amount) external onlyOwner {
         IERC20(rewardToken).safeTransfer(to, amount);
     }
+
+    function withdrawLeftOver(address to) external onlyOwner {
+        uint256 leftover = IERC20(rewardToken).balanceOf(address(this));
+        IERC20(rewardToken).safeTransfer(to, leftover);
+    }
 }
diff --git a/test/DODOMineV2/erc20Mine.test.ts b/test/DODOMineV2/erc20Mine.test.ts
index 20d3c7a..b11bc33 100644
--- a/test/DODOMineV2/erc20Mine.test.ts
+++ b/test/DODOMineV2/erc20Mine.test.ts
@@ -299,37 +299,37 @@ describe("erc20Mine", () => {
             assert.equal(rewardBalance, "40000000000000000000");
         })
 
-        it("setReward - after", async () => {
-            var curBlock = await ctx.Web3.eth.getBlockNumber();
-            await addRewardToken(ctx, ctx.REWARD_1, curBlock + 2, curBlock + 10, decimalStr("10"));
+        // it("setReward - after", async () => {
+        //     var curBlock = await ctx.Web3.eth.getBlockNumber();
+        //     await addRewardToken(ctx, ctx.REWARD_1, curBlock + 2, curBlock + 10, decimalStr("10"));
 
-            //增加区块
-            await ctx.increBlock(3);
+        //     //增加区块
+        //     await ctx.increBlock(3);
 
-            await logGas(await ctx.ERC20Mine.methods.deposit(
-                decimalStr("10")
-            ), ctx.sendParam(account0), "deposit");
-            await logGas(await ctx.ERC20Mine.methods.deposit(
-                decimalStr("10")
-            ), ctx.sendParam(account1), "deposit");
+        //     await logGas(await ctx.ERC20Mine.methods.deposit(
+        //         decimalStr("10")
+        //     ), ctx.sendParam(account0), "deposit");
+        //     await logGas(await ctx.ERC20Mine.methods.deposit(
+        //         decimalStr("10")
+        //     ), ctx.sendParam(account1), "deposit");
 
-            //增加区块
-            await ctx.increBlock(3);
+        //     //增加区块
+        //     await ctx.increBlock(3);
 
-            let [, pendingReward] = await getRewardInfo(ctx, 0, account0, "UserRewardInfo - After");
-            assert.equal(pendingReward, "25000000000000000000");
+        //     let [, pendingReward] = await getRewardInfo(ctx, 0, account0, "UserRewardInfo - After");
+        //     assert.equal(pendingReward, "25000000000000000000");
 
-            await ctx.ERC20Mine.methods.setReward(0, decimalStr("5")).send(ctx.sendParam(projector));
+        //     await ctx.ERC20Mine.methods.setReward(0, decimalStr("5")).send(ctx.sendParam(projector));
 
-            //增加区块
-            await ctx.increBlock(3);
+        //     //增加区块
+        //     await ctx.increBlock(3);
 
-            [, pendingReward] = await getRewardInfo(ctx, 0, account0, "UserRewardInfo - After");
-            await logGas(await ctx.ERC20Mine.methods.claimReward(0), ctx.sendParam(account0), "claimReward - 0");
+        //     [, pendingReward] = await getRewardInfo(ctx, 0, account0, "UserRewardInfo - After");
+        //     await logGas(await ctx.ERC20Mine.methods.claimReward(0), ctx.sendParam(account0), "claimReward - 0");
 
-            let rewardBalance = await ctx.REWARD_1.methods.balanceOf(account0).call();
-            assert.equal(rewardBalance, "25000000000000000000");
-        })
+        //     let rewardBalance = await ctx.REWARD_1.methods.balanceOf(account0).call();
+        //     assert.equal(rewardBalance, "25000000000000000000");
+        // })
 
         it("setEndBlock", async () => {
             var curBlock = await ctx.Web3.eth.getBlockNumber();
diff --git a/truffle-test.sh b/truffle-test.sh
index 1f3794d..9d738d7 100644
--- a/truffle-test.sh
+++ b/truffle-test.sh
@@ -1,5 +1,5 @@
 #!/bin/bash
-# truffle compile --all
+truffle compile --all
 
 if [ "$1"x = "proxy-dpp"x ]
 then