Make sure delegator stake balance is zero when delegating

keep-network · Jun 25, 2018 · c265be3 · c265be3
1 parent 06fe284
commit c265be3
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/contracts/solidity/contracts/StakeDelegatable.sol b/contracts/solidity/contracts/StakeDelegatable.sol
@@ -43,10 +43,19 @@ contract StakeDelegatable {
         notNull(_address)
         returns (uint256)
     {
+        // If provided address is a delegate return its delegator balance.
         address delegator = delegatorToDelegate[_address];
         if (delegator != address(0) && delegateToDelegator[delegator] == _address) {
             return stakeBalances[delegator];
         }
+
+        // If provided address is a delegator return zero balance since the 
+        // balance is delegated to a delegate.
+        address delegate = delegateToDelegator[_address];
+        if (delegate != address(0) && delegatorToDelegate[delegate] == _address) {
+            return 0;
+        }
+
         return stakeBalances[_address];
     }
 

diff --git a/contracts/solidity/test/TestStakeDelegate.js b/contracts/solidity/test/TestStakeDelegate.js
@@ -39,6 +39,7 @@ contract('TestStakeDelegate', function(accounts) {
     await stakingContract.requestDelegateFor(account_one, {from: account_three});
     await stakingContract.approveDelegateAt(account_three, {from: account_one});
     assert.equal(await stakingProxy.balanceOf(account_three), 200, "Delegate account should represent delegator's stake balance.");
+    assert.equal(await stakingProxy.balanceOf(account_one), 0, "Delegator account stake balance should become zero.");
   });
 
   it("should not be able to delegate stake to a delegate address that is not approved", async function() {
@@ -51,6 +52,7 @@ contract('TestStakeDelegate', function(accounts) {
     await stakingContract.requestDelegateFor(account_one, {from: account_three});
     await stakingContract.approveDelegateAt(account_three, {from: account_one});
     assert.equal(await stakingProxy.balanceOf(account_three), 200, "Delegate account should represent delegator's stake balance.");
+    assert.equal(await stakingProxy.balanceOf(account_one), 0, "Delegator account stake balance should become zero.");
 
     await stakingContract.approveDelegateAt(account_four, {from: account_one});
     await stakingContract.requestDelegateFor(account_one, {from: account_four});
@@ -64,6 +66,7 @@ contract('TestStakeDelegate', function(accounts) {
     assert.equal(await stakingProxy.balanceOf(account_three), 200, "Delegate account should represent delegator's stake balance.");
     await stakingContract.removeDelegate();
     assert.equal(await stakingProxy.balanceOf(account_three), 0, "Delegate account should stop representing delegator's stake balance.");
+    assert.equal(await stakingProxy.balanceOf(account_one), 200, "Delegator account should get its balance back.");
   });
 
   it("should be able to change stake and get delegate to reflect updated balance", async function() {
@@ -74,6 +77,7 @@ contract('TestStakeDelegate', function(accounts) {
     // Stake more tokens
     await token.approveAndCall(stakingContract.address, 100, "", {from: account_one});
     assert.equal(await stakingProxy.balanceOf(account_three), 300, "Delegate account should reflect delegator's updated stake balance.");
+    assert.equal(await stakingProxy.balanceOf(account_one), 0, "Delegator account stake balance should be zero.");
 
     // Unstake everything
     await stakingContract.initiateUnstake(300);