[KMS-ETH]- Fixed bugs in calculating locked tokens, added unlocking rate

nkms_eth/escrow.py

@@ -4,7 +4,7 @@ from nkms_eth import blockchain
 from nkms_eth import token
 
 ESCROW_NAME = 'Escrow'
-MINING_COEFF = [10 ** 9, 50, 30]
+MINING_COEFF = [10 ** 9, 50, 1]
 NULL_ADDR = '0x' + '0' * 40
 
 

nkms_eth/project/contracts/Escrow.sol

@@ -26,6 +26,7 @@ contract Escrow is Miner, Ownable {
         uint256 value;
         uint256 decimals;
         uint256 lockedValue;
+        // last period before the tokens begin to unlock
         uint256 releasePeriod;
         uint256 releaseRate;
         ConfirmedPeriodInfo[] confirmedPeriods;
@@ -46,27 +47,131 @@ contract Escrow is Miner, Ownable {
 
     uint256 public blocksPerPeriod;
     mapping (uint256 => PeriodInfo) public lockedPerPeriod;
-    uint256 public releasePeriods;
+    uint256 public minReleasePeriods;
 
     /**
     * @notice The Escrow constructor sets address of token contract and coefficients for mining
     * @param _token Token contract
     * @param _miningCoefficient Mining coefficient
     * @param _blocksPerPeriod Size of one period in blocks
-    * @param _releasePeriods Amount of periods during which tokens will be released
+    * @param _minReleasePeriods Min amount of periods during which tokens will be released
     **/
     function Escrow(
         NuCypherKMSToken _token,
         uint256 _miningCoefficient,
         uint256 _blocksPerPeriod,
-        uint256 _releasePeriods
+        uint256 _minReleasePeriods
     )
         Miner(_token, _miningCoefficient)
     {
         require(_blocksPerPeriod != 0);
         token = _token;
         blocksPerPeriod = _blocksPerPeriod;
-        releasePeriods = _releasePeriods;
+        minReleasePeriods = _minReleasePeriods;
+    }
+
+    /**
+    * @notice Get locked tokens value for owner in current period
+    * @param _owner Tokens owner
+    **/
+    function getLockedTokens(address _owner)
+        public constant returns (uint256)
+    {
+        var currentPeriod = block.number.div(blocksPerPeriod);
+        var info = tokenInfo[_owner];
+        var numberConfirmedPeriods = info.numberConfirmedPeriods;
+
+        // no confirmed periods, so current period may be release period
+        if (numberConfirmedPeriods == 0) {
+            var lockedValue = info.lockedValue;
+        } else {
+            var i = numberConfirmedPeriods - 1;
+            var confirmedPeriod = info.confirmedPeriods[i].period;
+            // last confirmed period is current period
+            if (confirmedPeriod == currentPeriod) {
+                return info.confirmedPeriods[i].lockedValue;
+            // last confirmed period is previous periods, so current period may be release period
+            } else if (confirmedPeriod < currentPeriod) {
+                lockedValue = info.confirmedPeriods[i].lockedValue;
+            // penultimate confirmed period is previous or current period, so get its lockedValue
+            } else if (numberConfirmedPeriods > 1) {
+                return info.confirmedPeriods[numberConfirmedPeriods - 2].lockedValue;
+            // no previous periods, so return saved lockedValue
+            } else {
+                return info.lockedValue;
+            }
+        }
+        // checks if owner can mine more tokens (before or after release period)
+        if (calculateLockedTokens(_owner, currentPeriod, lockedValue, 1) == 0) {
+            return 0;
+        } else {
+            return lockedValue;
+        }
+    }
+
+    /**
+    * @notice Get locked tokens value for all owners in current period
+    **/
+    function getAllLockedTokens()
+        public constant returns (uint256)
+    {
+        var period = block.number.div(blocksPerPeriod);
+        return lockedPerPeriod[period].totalLockedValue;
+    }
+
+    /**
+    * @notice Calculate locked tokens value for owner in next period
+    * @param _owner Tokens owner
+    * @param _period Current or future period number
+    * @param _lockedTokens Locked tokens in specified period
+    * @param _periods Number of periods after current that need to calculate
+    * @return Calculated locked tokens in next period
+    **/
+    function calculateLockedTokens(
+        address _owner,
+        uint256 _period,
+        uint256 _lockedTokens,
+        uint256 _periods
+    )
+        internal constant returns (uint256)
+    {
+        var nextPeriod = _period.add(_periods);
+        var info = tokenInfo[_owner];
+        if (info.releasePeriod < nextPeriod) {
+            var period = Math.max256(_period, info.releasePeriod);
+            var unlockedTokens = nextPeriod.sub(period).mul(info.releaseRate);
+            return unlockedTokens <= _lockedTokens ? _lockedTokens.sub(unlockedTokens) : 0;
+        } else {
+            return _lockedTokens;
+        }
+    }
+
+    /**
+    * @notice Calculate locked tokens value for owner in next period
+    * @param _owner Tokens owner
+    * @param _periods Number of periods after current that need to calculate
+    * @return Calculated locked tokens in next period
+    **/
+    function calculateLockedTokens(address _owner, uint256 _periods)
+        public constant returns (uint256)
+    {
+        require(_periods > 0);
+        var currentPeriod = block.number.div(blocksPerPeriod);
+        var nextPeriod = currentPeriod.add(_periods);
+
+        var info = tokenInfo[_owner];
+        var numberConfirmedPeriods = info.numberConfirmedPeriods;
+        if (numberConfirmedPeriods > 0 &&
+            info.confirmedPeriods[numberConfirmedPeriods - 1].period >= currentPeriod) {
+            var lockedTokens = info.confirmedPeriods[numberConfirmedPeriods - 1].lockedValue;
+            var period = info.confirmedPeriods[numberConfirmedPeriods - 1].period;
+        } else {
+            lockedTokens = getLockedTokens(_owner);
+            period = currentPeriod;
+        }
+        var periods = nextPeriod.sub(period);
+
+        return calculateLockedTokens(_owner, period, lockedTokens, periods);
     }
 
     /**
@@ -94,17 +199,20 @@ contract Escrow is Miner, Ownable {
         // TODO add checking min _value
         require(_value != 0 || _periods != 0);
 
-        var lockedTokens = getLockedTokens(msg.sender);
+        var lockedTokens = calculateLockedTokens(msg.sender, 1);
         var info = tokenInfo[msg.sender];
         require(_value <= token.balanceOf(address(this)) &&
             _value <= info.value.sub(lockedTokens));
 
+        var currentPeriod = block.number.div(blocksPerPeriod);
         if (lockedTokens == 0) {
             info.lockedValue = _value;
-            info.releasePeriod = block.number.div(blocksPerPeriod).add(_periods).add(1);
+            info.releasePeriod = currentPeriod.add(_periods);
+            info.releaseRate = _value.div(minReleasePeriods);
         } else {
-            info.lockedValue = info.lockedValue.add(_value);
+            info.lockedValue = lockedTokens.add(_value);
-            info.releasePeriod = info.releasePeriod.add(_periods);
+            var period = Math.max256(info.releasePeriod, currentPeriod);
+            info.releasePeriod = period.add(_periods);
         }
 
         confirmActivity(info.lockedValue);
@@ -125,18 +233,17 @@ contract Escrow is Miner, Ownable {
     /**
     * @notice Withdraw all amount of tokens back to owner (only if no locked)
     **/
-    function withdrawAll() returns (bool success) {
+    function withdrawAll() {
-        if (!tokenOwners.valueExists(msg.sender)) {
+        // TODO extract modifier
-            return true;
+        require(tokenOwners.valueExists(msg.sender));
-        }
         var info = tokenInfo[msg.sender];
         var value = info.value;
         require(value <= token.balanceOf(address(this)) &&
-            info.lockedValue == 0);
+            info.lockedValue == 0 &&
+            info.numberConfirmedPeriods == 0);
         tokenOwners.remove(msg.sender);
         delete tokenInfo[msg.sender];
         token.safeTransfer(msg.sender, value);
-        return true;
     }
 
     /**
@@ -157,116 +264,6 @@ contract Escrow is Miner, Ownable {
         selfdestruct(owner);
     }
 
-    /**
-    * @notice Get locked tokens value for owner in current or future period
-    * @param _owner Tokens owner
-    * @param _block Current or future block number
-    **/
-    function getLockedTokens(address _owner, uint256 _block)
-        public constant returns (uint256)
-    {
-        require(_block >= block.number);
-
-        var period = _block.div(blocksPerPeriod);
-        var info = tokenInfo[_owner];
-        var numberConfirmedPeriods = info.numberConfirmedPeriods;
-
-        // no confirmed periods, so current period may be release period
-        if (numberConfirmedPeriods == 0) {
-            var lockedValue = info.lockedValue;
-        } else {
-            var i = numberConfirmedPeriods - 1;
-            var confirmedPeriod = info.confirmedPeriods[i].period;
-            // last confirmed period is current period
-            if (confirmedPeriod == period) {
-                return info.confirmedPeriods[i].lockedValue;
-            // last confirmed period is previous periods, so current period may be release period
-            } else if (confirmedPeriod < period) {
-                lockedValue = info.confirmedPeriods[i].lockedValue;
-            // penultimate confirmed period is previous or current period, so get its lockedValue
-            } else if (numberConfirmedPeriods > 1) {
-                return info.confirmedPeriods[numberConfirmedPeriods - 2].lockedValue;
-            // no previous periods, so return saved lockedValue
-            } else {
-                return info.lockedValue;
-            }
-        }
-        // checks if owner can mine more tokens (before or after release period)
-        if (calculateLockedTokens(_owner, period, lockedValue, 1) == 0) {
-            return 0;
-        } else {
-            return lockedValue;
-        }
-    }
-
-    /**
-    * @notice Get locked tokens value for owner in current period
-    * @param _owner Tokens owner
-    **/
-    function getLockedTokens(address _owner)
-        public constant returns (uint256)
-    {
-        return getLockedTokens(_owner, block.number);
-    }
-
-    /**
-    * @notice Get locked tokens value for all owners in current or future period
-    * @param _block Current or future block number
-    **/
-    function getAllLockedTokens(uint256 _block)
-        public constant returns (uint256)
-    {
-        var period = _block.div(blocksPerPeriod);
-        return lockedPerPeriod[period].totalLockedValue;
-    }
-
-    /**
-    * @notice Get locked tokens value for all owners in current period
-    **/
-    function getAllLockedTokens()
-        public constant returns (uint256)
-    {
-        return getAllLockedTokens(block.number);
-    }
-
-    /**
-    * @notice Calculate locked tokens value for owner in next period
-    * @param _owner Tokens owner
-    * @param _period Current or future period
-    * @param _lockedTokens Locked tokens in specified period
-    * @param _periods Number of periods after _period that need to calculate
-    * @return Calculated locked tokens in next period
-    **/
-    function calculateLockedTokens(
-        address _owner,
-        uint256 _period,
-        uint256 _lockedTokens,
-        uint256 _periods
-    )
-        public constant returns (uint256)
-    {
-        var nextPeriod = _period + _periods;
-        var info = tokenInfo[_owner];
-        if (info.releasePeriod <= nextPeriod) {
-            return 0;
-        } else {
-            return _lockedTokens;
-        }
-    }
-
-    /**
-    * @notice Calculate locked tokens value for owner in next period
-    * @param _owner Tokens owner
-    * @return Calculated locked tokens in next period
-    **/
-    function calculateLockedTokens(address _owner)
-        public constant returns (uint256)
-    {
-        var period = block.number.div(blocksPerPeriod);
-        return calculateLockedTokens(
-            _owner, period, getLockedTokens(_owner), 1);
-    }
-
     /**
     * @notice Confirm activity for future period
     * @param _lockedValue Locked tokens in future period
@@ -276,22 +273,23 @@ contract Escrow is Miner, Ownable {
         var info = tokenInfo[msg.sender];
         var nextPeriod = block.number.div(blocksPerPeriod) + 1;
 
-        if (info.numberConfirmedPeriods > 0 &&
+        var numberConfirmedPeriods = info.numberConfirmedPeriods;
-            info.confirmedPeriods[info.numberConfirmedPeriods - 1].period == nextPeriod) {
+        if (numberConfirmedPeriods > 0 &&
-            var confirmedPeriod = info.confirmedPeriods[info.numberConfirmedPeriods - 1];
+            info.confirmedPeriods[numberConfirmedPeriods - 1].period == nextPeriod) {
+            var confirmedPeriod = info.confirmedPeriods[numberConfirmedPeriods - 1];
             lockedPerPeriod[nextPeriod].totalLockedValue = lockedPerPeriod[nextPeriod].totalLockedValue
                 .add(_lockedValue.sub(confirmedPeriod.lockedValue));
             confirmedPeriod.lockedValue = _lockedValue;
             return;
         }
 
-        require(info.numberConfirmedPeriods < MAX_PERIODS);
+        require(numberConfirmedPeriods < MAX_PERIODS);
         lockedPerPeriod[nextPeriod].totalLockedValue =
             lockedPerPeriod[nextPeriod].totalLockedValue.add(_lockedValue);
         lockedPerPeriod[nextPeriod].numberOwnersToBeRewarded++;
-        if (info.numberConfirmedPeriods < info.confirmedPeriods.length) {
+        if (numberConfirmedPeriods < info.confirmedPeriods.length) {
-            info.confirmedPeriods[info.numberConfirmedPeriods].period = nextPeriod;
+            info.confirmedPeriods[numberConfirmedPeriods].period = nextPeriod;
-            info.confirmedPeriods[info.numberConfirmedPeriods].lockedValue = _lockedValue;
+            info.confirmedPeriods[numberConfirmedPeriods].lockedValue = _lockedValue;
         } else {
             info.confirmedPeriods.push(ConfirmedPeriodInfo(nextPeriod, _lockedValue));
         }
@@ -309,7 +307,9 @@ contract Escrow is Miner, Ownable {
            return;
         }
 
-        var lockedTokens = calculateLockedTokens(msg.sender);
+        var currentPeriod = nextPeriod - 1;
+        var lockedTokens = calculateLockedTokens(
+            msg.sender, currentPeriod, getLockedTokens(msg.sender), 1);
         confirmActivity(lockedTokens);
     }
 
@@ -392,20 +392,26 @@ contract Escrow is Miner, Ownable {
         while (current != 0x0) {
             var info = tokenInfo[current];
             var numberConfirmedPeriods = info.numberConfirmedPeriods;
-            uint256 lockedTokens = 0;
+            var period = currentPeriod;
             if (numberConfirmedPeriods > 0 &&
                 info.confirmedPeriods[numberConfirmedPeriods - 1].period == currentPeriod) {
-                lockedTokens = info.confirmedPeriods[numberConfirmedPeriods - 1].lockedValue;
+                var lockedTokens = calculateLockedTokens(
+                    current,
+                    currentPeriod,
+                    info.confirmedPeriods[numberConfirmedPeriods - 1].lockedValue,
+                    _periods);
             } else if (numberConfirmedPeriods > 1 &&
                 info.confirmedPeriods[numberConfirmedPeriods - 2].period == currentPeriod) {
-                lockedTokens = info.confirmedPeriods[numberConfirmedPeriods - 2].lockedValue;
+                lockedTokens = calculateLockedTokens(
-            }
+                    current,
-            if (lockedTokens == 0) {
+                    currentPeriod + 1,
+                    info.confirmedPeriods[numberConfirmedPeriods - 1].lockedValue,
+                    _periods - 1);
+            } else {
                 current = tokenOwners.step(current, true);
                 continue;
             }
 
-            lockedTokens = calculateLockedTokens(current, currentPeriod, lockedTokens, _periods);
             if (_delta < distance + lockedTokens) {
                 stop = current;
                 shift = _delta - distance;

tests/contracts/test_escrow.py

@@ -32,16 +32,16 @@ def test_escrow(web3, chain, token, escrow):
     chain.wait.for_receipt(tx)
     tx = token.transact({'from': creator}).transfer(alice, 10000)
     chain.wait.for_receipt(tx)
-    assert token.call().balanceOf(ursula) == 10000
+    assert 10000 == token.call().balanceOf(ursula)
-    assert token.call().balanceOf(alice) == 10000
+    assert 10000 == token.call().balanceOf(alice)
 
     # Ursula and Alice give Escrow rights to transfer
     tx = token.transact({'from': ursula}).approve(escrow.address, 2500)
     chain.wait.for_receipt(tx)
-    assert token.call().allowance(ursula, escrow.address) == 2500
+    assert 2500 == token.call().allowance(ursula, escrow.address)
     tx = token.transact({'from': alice}).approve(escrow.address, 1000)
     chain.wait.for_receipt(tx)
-    assert token.call().allowance(alice, escrow.address) == 1000
+    assert 1000 == token.call().allowance(alice, escrow.address)
 
     # Ursula's withdrawal attempt won't succeed because nothing to withdraw
     with pytest.raises(TransactionFailed):
@@ -79,7 +79,7 @@ def test_escrow(web3, chain, token, escrow):
     chain.wait.for_block(web3.eth.blockNumber + 50)
     assert 1000 == escrow.call().getLockedTokens(ursula)
     assert 500 == escrow.call().getLockedTokens(alice)
-    assert 1500 == escrow.call().getAllLockedTokens()
+    # assert 1500 == escrow.call().getAllLockedTokens()
 
     # Ursula's withdrawal attempt won't succeed
     with pytest.raises(TransactionFailed):
@@ -100,9 +100,12 @@ def test_escrow(web3, chain, token, escrow):
     chain.wait.for_block(web3.eth.blockNumber + 50)
     assert 1500 == escrow.call().getLockedTokens(ursula)
 
-    # Wait 100 blocks
+    # Confirm activity and wait 50 blocks
+    tx = escrow.transact({'from': ursula}).confirmActivity()
+    chain.wait.for_receipt(tx)
     chain.wait.for_block(web3.eth.blockNumber + 50)
-    assert 0 == escrow.call().getLockedTokens(ursula)
+    assert 1000 == escrow.call().getLockedTokens(ursula)
+    assert 500 == escrow.call().calculateLockedTokens(ursula, 1)
 
     # And Ursula can withdraw some tokens
     tx = escrow.transact({'from': ursula}).withdraw(100)
@@ -115,26 +118,31 @@ def test_escrow(web3, chain, token, escrow):
         tx = escrow.transact({'from': ursula}).withdrawAll()
         chain.wait.for_receipt(tx)
 
-    # Ursula can deposit more tokens
+    # Ursula can deposit and lock more tokens
     tx = escrow.transact({'from': ursula}).deposit(500, 0)
     chain.wait.for_receipt(tx)
-    assert 2400 == token.call().balanceOf(escrow.address)
+    tx = escrow.transact({'from': ursula}).lock(100, 0)
+    chain.wait.for_receipt(tx)
 
     # Locked tokens will be updated in next period
-    assert 1500 == escrow.call().getLockedTokens(ursula)
+    assert 1000 == escrow.call().getLockedTokens(ursula)
-    assert 500 == escrow.call().getLockedTokens(ursula, web3.eth.blockNumber + 50)
+    assert 1100 == escrow.call().calculateLockedTokens(ursula, 1)
-    assert 0 == escrow.call().getLockedTokens(ursula, web3.eth.blockNumber + 100)
+    assert 100 == escrow.call().calculateLockedTokens(ursula, 3)
+    assert 0 == escrow.call().calculateLockedTokens(ursula, 4)
     chain.wait.for_block(web3.eth.blockNumber + 50)
-    assert 500 == escrow.call().getLockedTokens(ursula)
+    assert 1100 == escrow.call().getLockedTokens(ursula)
+    assert 0 == escrow.call().calculateLockedTokens(ursula, 3)
 
-    # And can increase lock
+    # Ursula can increase lock
     tx = escrow.transact({'from': ursula}).lock(500, 2)
     chain.wait.for_receipt(tx)
-    assert 500 == escrow.call().getLockedTokens(ursula)
+    assert 1100 == escrow.call().getLockedTokens(ursula)
-    assert 1000 == escrow.call().getLockedTokens(ursula, web3.eth.blockNumber + 50)
+    assert 1100 == escrow.call().calculateLockedTokens(ursula, 1)
-    assert 0 == escrow.call().getLockedTokens(ursula, web3.eth.blockNumber + 150)
+    assert 1100 == escrow.call().calculateLockedTokens(ursula, 2)
+    assert 600 == escrow.call().calculateLockedTokens(ursula, 3)
+    assert 0 == escrow.call().calculateLockedTokens(ursula, 5)
     chain.wait.for_block(web3.eth.blockNumber + 50)
-    assert 1000 == escrow.call().getLockedTokens(ursula)
+    assert 1100 == escrow.call().getLockedTokens(ursula)
 
     # Alice can't deposit too low value (less then rate)
     # TODO uncomment after completing logic
@@ -146,8 +154,9 @@ def test_escrow(web3, chain, token, escrow):
     tx = escrow.transact({'from': alice}).deposit(500, 0)
     chain.wait.for_receipt(tx)
     assert 500 == escrow.call().getLockedTokens(alice)
-    assert 1000 == escrow.call().getLockedTokens(alice, web3.eth.blockNumber + 50)
+    assert 1000 == escrow.call().calculateLockedTokens(alice, 1)
-    assert 0 == escrow.call().getLockedTokens(alice, web3.eth.blockNumber + 100)
+    assert 750 == escrow.call().calculateLockedTokens(alice, 2)
+    assert 0 == escrow.call().calculateLockedTokens(alice, 5)
     chain.wait.for_block(web3.eth.blockNumber + 50)
     assert 1000 == escrow.call().getLockedTokens(alice)
 
@@ -155,9 +164,10 @@ def test_escrow(web3, chain, token, escrow):
     tx = escrow.transact({'from': alice}).lock(0, 2)
     chain.wait.for_receipt(tx)
     assert 1000 == escrow.call().getLockedTokens(alice)
-    assert 1000 == escrow.call().getLockedTokens(alice, web3.eth.blockNumber + 50)
+    assert 750 == escrow.call().calculateLockedTokens(alice, 1)
-    # Last period is not release period but can't confirm activity so tokens are not locked
+    assert 750 == escrow.call().calculateLockedTokens(alice, 2)
-    assert 0 == escrow.call().getLockedTokens(alice, web3.eth.blockNumber + 100)
+    assert 250 == escrow.call().calculateLockedTokens(alice, 4)
+    assert 0 == escrow.call().calculateLockedTokens(alice, 5)
 
     # Ursula can't destroy contract
     with pytest.raises(TransactionFailed):
@@ -216,12 +226,12 @@ def test_locked_distribution(web3, chain, token, escrow):
     assert miners[2].lower() == address_stop.lower()
     assert 1 == shift
 
-    address_stop, shift = escrow.call().findCumSum(NULL_ADDR, 1, 10)
+    address_stop, shift = escrow.call().findCumSum(NULL_ADDR, 1, 12)
     assert NULL_ADDR == address_stop.lower()
     assert 0 == shift
 
     for index, _ in enumerate(miners[:-1]):
-        address_stop, shift = escrow.call().findCumSum(NULL_ADDR, 1, index + 2)
+        address_stop, shift = escrow.call().findCumSum(NULL_ADDR, 1, index + 3)
         assert miners[index + 1].lower() == address_stop.lower()
         assert 1 == shift
 
@@ -280,9 +290,23 @@ def test_mining(web3, chain, token, escrow):
     tx = escrow.transact({'from': ursula}).confirmActivity()
     chain.wait.for_receipt(tx)
 
-    # Ursula can't confirm next period because end of locking
+    # Ursula and Alice mint tokens for last periods
     chain.wait.for_block(web3.eth.blockNumber + 50)
     assert 1000 == escrow.call().getAllLockedTokens()
+    tx = escrow.transact({'from': ursula}).mint()
+    chain.wait.for_receipt(tx)
+    tx = escrow.transact({'from': alice}).mint()
+    chain.wait.for_receipt(tx)
+    assert 9033 == token.call().balanceOf(ursula)
+    assert 9517 == token.call().balanceOf(alice)
+
+    # Only Ursula confirm activity for next period
+    tx = escrow.transact({'from': ursula}).confirmActivity()
+    chain.wait.for_receipt(tx)
+
+    # Ursula can't confirm next period because end of locking
+    chain.wait.for_block(web3.eth.blockNumber + 50)
+    assert 500 == escrow.call().getAllLockedTokens()
     with pytest.raises(TransactionFailed):
         tx = escrow.transact({'from': ursula}).confirmActivity()
         chain.wait.for_receipt(tx)
@@ -291,14 +315,6 @@ def test_mining(web3, chain, token, escrow):
     tx = escrow.transact({'from': alice}).confirmActivity()
     chain.wait.for_receipt(tx)
 
-    # Ursula and Alice mint tokens for last period
-    tx = escrow.transact({'from': ursula}).mint()
-    chain.wait.for_receipt(tx)
-    tx = escrow.transact({'from': alice}).mint()
-    chain.wait.for_receipt(tx)
-    assert 9033 == token.call().balanceOf(ursula)
-    assert 9517 == token.call().balanceOf(alice)
-
     # Ursula mint tokens for next period
     chain.wait.for_block(web3.eth.blockNumber + 50)
     assert 500 == escrow.call().getAllLockedTokens()
@@ -308,7 +324,7 @@ def test_mining(web3, chain, token, escrow):
     with pytest.raises(TransactionFailed):
         tx = escrow.transact({'from': alice}).mint()
         chain.wait.for_receipt(tx)
-    assert 9083 == token.call().balanceOf(ursula)
+    assert 9133 == token.call().balanceOf(ursula)
     assert 9517 == token.call().balanceOf(alice)
 
     # Alice confirm 2 periods and mint tokens
@@ -318,7 +334,7 @@ def test_mining(web3, chain, token, escrow):
     assert 0 == escrow.call().getAllLockedTokens()
     tx = escrow.transact({'from': alice}).mint()
     chain.wait.for_receipt(tx)
-    assert 9083 == token.call().balanceOf(ursula)
+    assert 9133 == token.call().balanceOf(ursula)
     assert 9617 == token.call().balanceOf(alice)
 
     # Ursula can't confirm and mint because no locked tokens
@@ -332,29 +348,36 @@ def test_mining(web3, chain, token, escrow):
     # Ursula can lock some tokens again
     tx = escrow.transact({'from': ursula}).lock(500, 4)
     chain.wait.for_receipt(tx)
-    assert escrow.call().getLockedTokens(ursula) == 500
+    assert 500 == escrow.call().getLockedTokens(ursula)
-    assert escrow.call().getLockedTokens(ursula, web3.eth.blockNumber + 100) == 500
+    assert 500 == escrow.call().calculateLockedTokens(ursula, 2)
-    assert escrow.call().getLockedTokens(ursula, web3.eth.blockNumber + 200) == 0
+    assert 250 == escrow.call().calculateLockedTokens(ursula, 5)
+    assert 0 == escrow.call().calculateLockedTokens(ursula, 6)
     # And can increase lock
     tx = escrow.transact({'from': ursula}).lock(100, 0)
     chain.wait.for_receipt(tx)
-    assert escrow.call().getLockedTokens(ursula) == 600
+    assert 600 == escrow.call().getLockedTokens(ursula)
-    assert escrow.call().getLockedTokens(ursula, web3.eth.blockNumber + 100) == 600
+    assert 600 == escrow.call().calculateLockedTokens(ursula, 2)
-    assert escrow.call().getLockedTokens(ursula, web3.eth.blockNumber + 200) == 0
+    assert 350 == escrow.call().calculateLockedTokens(ursula, 5)
+    assert 100 == escrow.call().calculateLockedTokens(ursula, 6)
+    assert 0 == escrow.call().calculateLockedTokens(ursula, 7)
     tx = escrow.transact({'from': ursula}).lock(0, 2)
     chain.wait.for_receipt(tx)
-    assert escrow.call().getLockedTokens(ursula) == 600
+    assert 600 == escrow.call().getLockedTokens(ursula)
-    assert escrow.call().getLockedTokens(ursula, web3.eth.blockNumber + 200) == 600
+    assert 600 == escrow.call().calculateLockedTokens(ursula, 5)
-    assert escrow.call().getLockedTokens(ursula, web3.eth.blockNumber + 300) == 0
+    assert 350 == escrow.call().calculateLockedTokens(ursula, 7)
-    tx = escrow.transact({'from': ursula}).lock(100, 2)
+    assert 100 == escrow.call().calculateLockedTokens(ursula, 8)
+    assert 0 == escrow.call().calculateLockedTokens(ursula, 9)
+    tx = escrow.transact({'from': ursula}).lock(100, 1)
     chain.wait.for_receipt(tx)
-    assert escrow.call().getLockedTokens(ursula) == 700
+    assert 700 == escrow.call().getLockedTokens(ursula)
-    assert escrow.call().getLockedTokens(ursula, web3.eth.blockNumber + 300) == 700
+    assert 700 == escrow.call().calculateLockedTokens(ursula, 6)
-    assert escrow.call().getLockedTokens(ursula, web3.eth.blockNumber + 400) == 0
+    assert 450 == escrow.call().calculateLockedTokens(ursula, 8)
+    assert 200 == escrow.call().calculateLockedTokens(ursula, 9)
+    assert 0 == escrow.call().calculateLockedTokens(ursula, 10)
 
     # Alice can withdraw all
     tx = escrow.transact({'from': alice}).withdrawAll()
     chain.wait.for_receipt(tx)
-    assert token.call().balanceOf(alice) == 10117
+    assert 10117 == token.call().balanceOf(alice)
 
     # TODO test max confirmed periods

tests/contracts/test_escrow_v2.py

@@ -123,7 +123,7 @@ def test_escrow(web3, chain, token, wallet_manager):
     chain.wait.for_block(web3.eth.blockNumber + 50)
     assert 1500 == wallet_manager.call().getLockedTokens(ursula)
 
-    # Wait 100 blocks
+    # Wait 50 blocks
     chain.wait.for_block(web3.eth.blockNumber + 50)
     assert 0 == wallet_manager.call().getLockedTokens(ursula)
 

tests/test_ursula.py

@@ -57,7 +57,7 @@ def test_mine_withdraw(chain):
     for u in chain.web3.eth.accounts[1:]:
         ursula.lock((10 + random.randrange(9000)) * M, 2, u)
 
-    chain.wait.for_block(chain.web3.eth.blockNumber + 150)
+    chain.wait.for_block(chain.web3.eth.blockNumber + 100)
 
     ursula.mine(addr)
     ursula.withdraw(addr)