Enhance setStakingContract function (#274)

* Enhance setStakingContract function * Update poa20_test.js
omni · Aug 17, 2019 · 1841968 · 1841968
1 parent de8effe
commit 1841968
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/contracts/ERC677BridgeTokenRewardable.sol b/contracts/ERC677BridgeTokenRewardable.sol
@@ -17,6 +17,7 @@ contract ERC677BridgeTokenRewardable is ERC677BridgeToken {
 
     function setStakingContract(address _stakingContract) external onlyOwner {
         require(AddressUtils.isContract(_stakingContract));
+        require(balanceOf(_stakingContract) == 0);
         stakingContract = _stakingContract;
     }
 

diff --git a/test/poa20_test.js b/test/poa20_test.js
@@ -139,6 +139,17 @@ async function testERC677BridgeToken(accounts, rewardable) {
         await token.setStakingContract(ZERO_ADDRESS).should.be.rejectedWith(ERROR_MSG)
         ;(await token.stakingContract()).should.be.equal(ZERO_ADDRESS)
       })
+
+      it('fail to set Staking contract address with non-zero balance', async () => {
+        const stakingContract = await StakingTest.new()
+        ;(await token.stakingContract()).should.be.equal(ZERO_ADDRESS)
+
+        await token.mint(user, 1, { from: owner }).should.be.fulfilled
+        await token.transfer(stakingContract.address, 1, { from: user }).should.be.fulfilled
+
+        await token.setStakingContract(stakingContract.address).should.be.rejectedWith(ERROR_MSG)
+        ;(await token.stakingContract()).should.be.equal(ZERO_ADDRESS)
+      })
     })
 
     describe('#mintReward', async () => {