StakingEscrow: enable wind down for claim from worklock

nucypher/blockchain/eth/sol/source/contracts/StakingEscrow.sol

@@ -41,7 +41,7 @@ interface WorkLockInterface {
 /**
 * @notice Contract holds and locks stakers tokens.
 * Each staker that locks their tokens will receive some compensation
-* @dev |v5.4.1|
+* @dev |v5.4.2|
 */
 contract StakingEscrow is Issuer, IERC900History {
 
@@ -543,6 +543,11 @@ contract StakingEscrow is Issuer, IERC900History {
         external
     {
         require(msg.sender == address(workLock));
+        StakerInfo storage info = stakerInfo[_staker];
+        if (!info.flags.bitSet(WIND_DOWN_INDEX) && info.subStakes.length == 0) {
+            info.flags = info.flags.toggleBit(WIND_DOWN_INDEX);
+            emit WindDownSet(_staker, true);
+        }
         deposit(_staker, msg.sender, MAX_SUB_STAKES, _value, _periods);
     }
 
@@ -557,12 +562,10 @@ contract StakingEscrow is Issuer, IERC900History {
             return;
         }
         info.flags = info.flags.toggleBit(WIND_DOWN_INDEX);
-
-        uint16 currentPeriod = getCurrentPeriod();
-        uint16 nextPeriod = currentPeriod + 1;
         emit WindDownSet(msg.sender, _windDown);
 
         // duration adjustment if next period is committed
+        uint16 nextPeriod = getCurrentPeriod() + 1;
         if (info.nextCommittedPeriod != nextPeriod) {
            return;
         }

tests/contracts/integration/test_intercontract_integration.py

@@ -487,6 +487,8 @@ def test_worklock_phases(testerchain,
     tx = worklock.functions.claim().transact({'from': staker2, 'gas_price': 0})
     testerchain.wait_for_receipt(tx)
     assert worklock.functions.workInfo(staker2).call()[2]
+    wind_down, _re_stake, _measure_work, _snapshots = escrow.functions.getFlags(staker2).call()
+    assert wind_down
 
     assert token.functions.balanceOf(staker2).call() == 0
     assert escrow.functions.getLockedTokens(staker2, 0).call() == 0
@@ -516,6 +518,8 @@ def test_worklock_phases(testerchain,
     pytest.staker1_tokens += staker1_claims
     assert escrow.functions.getLockedTokens(staker1, 1).call() == pytest.staker1_tokens
     pytest.escrow_supply += staker1_claims
+    wind_down, _re_stake, _measure_work, _snapshots = escrow.functions.getFlags(staker1).call()
+    assert not wind_down
 
     # Staker prolongs lock duration
     tx = escrow.functions.prolongStake(0, 3).transact({'from': staker2, 'gas_price': 0})

tests/contracts/main/staking_escrow/test_staking_escrow.py

@@ -1548,9 +1548,10 @@ def test_staking_from_worklock(testerchain, token, escrow_contract, token_econom
 
     maximum_allowed_locked = 1500
     escrow = escrow_contract(maximum_allowed_locked, disable_reward=True)
-    creator, staker1, staker2, staker3, staker4 = testerchain.client.accounts[0:5]
+    creator, staker1, staker2, staker3 = testerchain.client.accounts[0:4]
     deposit_log = escrow.events.Deposited.createFilter(fromBlock='latest')
     lock_log = escrow.events.Locked.createFilter(fromBlock='latest')
+    wind_down_log = escrow.events.WindDownSet.createFilter(fromBlock='latest')
 
     # Deploy WorkLock mock
     worklock, _ = deploy_contract('WorkLockForStakingEscrowMock', token.address, escrow.address)
@@ -1572,6 +1573,8 @@ def test_staking_from_worklock(testerchain, token, escrow_contract, token_econom
     assert token.functions.balanceOf(escrow.address).call() == 0
 
     # Deposit tokens from WorkLock
+    wind_down, _re_stake, _measure_work, _snapshots = escrow.functions.getFlags(staker1).call()
+    assert not wind_down
     current_period = escrow.functions.getCurrentPeriod().call()
     tx = worklock.functions.depositFromWorkLock(staker1, value, duration).transact()
     testerchain.wait_for_receipt(tx)
@@ -1580,6 +1583,8 @@ def test_staking_from_worklock(testerchain, token, escrow_contract, token_econom
     assert escrow.functions.getLockedTokens(staker1, 1).call() == value
     assert escrow.functions.getLockedTokens(staker1, duration).call() == value
     assert escrow.functions.getLockedTokens(staker1, duration + 1).call() == 0
+    wind_down, _re_stake, _measure_work, _snapshots = escrow.functions.getFlags(staker1).call()
+    assert wind_down
 
     # Check that all events are emitted
     events = deposit_log.get_all_entries()
@@ -1596,3 +1601,94 @@ def test_staking_from_worklock(testerchain, token, escrow_contract, token_econom
     assert event_args['value'] == value
     assert event_args['firstPeriod'] == current_period + 1
     assert event_args['periods'] == duration
+
+    events = wind_down_log.get_all_entries()
+    assert len(events) == 1
+    event_args = events[-1]['args']
+    assert event_args['staker'] == staker1
+    assert event_args['windDown']
+
+    # Deposit directly and then through WorkLock
+    tx = token.functions.transfer(staker2, maximum_allowed_locked).transact({'from': creator})
+    testerchain.wait_for_receipt(tx)
+    tx = token.functions.approve(escrow.address, maximum_allowed_locked).transact({'from': staker2})
+    testerchain.wait_for_receipt(tx)
+    tx = escrow.functions.deposit(staker2, value, duration).transact({'from': staker2})
+    testerchain.wait_for_receipt(tx)
+
+    wind_down, _re_stake, _measure_work, _snapshots = escrow.functions.getFlags(staker2).call()
+    assert not wind_down
+    current_period = escrow.functions.getCurrentPeriod().call()
+    tx = worklock.functions.depositFromWorkLock(staker2, value, duration).transact()
+    testerchain.wait_for_receipt(tx)
+    assert token.functions.balanceOf(escrow.address).call() == 3 * value
+    assert escrow.functions.getLockedTokens(staker2, 0).call() == 0
+    assert escrow.functions.getLockedTokens(staker2, 1).call() == 2 * value
+    assert escrow.functions.getLockedTokens(staker2, duration).call() == 2 * value
+    assert escrow.functions.getLockedTokens(staker2, duration + 1).call() == 0
+    wind_down, _re_stake, _measure_work, _snapshots = escrow.functions.getFlags(staker2).call()
+    assert not wind_down
+
+    # Check that all events are emitted
+    events = deposit_log.get_all_entries()
+    assert len(events) == 3
+    event_args = events[-1]['args']
+    assert event_args['staker'] == staker2
+    assert event_args['value'] == value
+    assert event_args['periods'] == duration
+
+    events = lock_log.get_all_entries()
+    assert len(events) == 3
+    event_args = events[-1]['args']
+    assert event_args['staker'] == staker2
+    assert event_args['value'] == value
+    assert event_args['firstPeriod'] == current_period + 1
+    assert event_args['periods'] == duration
+
+    events = wind_down_log.get_all_entries()
+    assert len(events) == 1
+
+    # Enable wind down before deposit from WorkLock
+    tx = token.functions.transfer(staker3, maximum_allowed_locked).transact({'from': creator})
+    testerchain.wait_for_receipt(tx)
+    tx = token.functions.approve(escrow.address, maximum_allowed_locked).transact({'from': staker3})
+    testerchain.wait_for_receipt(tx)
+    tx = escrow.functions.deposit(staker3, value, duration).transact({'from': staker2})
+    testerchain.wait_for_receipt(tx)
+
+    tx = escrow.functions.setWindDown(True).transact({'from': staker3})
+    testerchain.wait_for_receipt(tx)
+    wind_down, _re_stake, _measure_work, _snapshots = escrow.functions.getFlags(staker3).call()
+    assert wind_down
+    events = wind_down_log.get_all_entries()
+    assert len(events) == 2
+
+    current_period = escrow.functions.getCurrentPeriod().call()
+    tx = worklock.functions.depositFromWorkLock(staker3, value, duration).transact()
+    testerchain.wait_for_receipt(tx)
+    assert token.functions.balanceOf(escrow.address).call() == 5 * value
+    assert escrow.functions.getLockedTokens(staker3, 0).call() == 0
+    assert escrow.functions.getLockedTokens(staker3, 1).call() == 2 * value
+    assert escrow.functions.getLockedTokens(staker3, duration).call() == 2 * value
+    assert escrow.functions.getLockedTokens(staker3, duration + 1).call() == 0
+    wind_down, _re_stake, _measure_work, _snapshots = escrow.functions.getFlags(staker3).call()
+    assert wind_down
+
+    # Check that all events are emitted
+    events = deposit_log.get_all_entries()
+    assert len(events) == 5
+    event_args = events[-1]['args']
+    assert event_args['staker'] == staker3
+    assert event_args['value'] == value
+    assert event_args['periods'] == duration
+
+    events = lock_log.get_all_entries()
+    assert len(events) == 5
+    event_args = events[-1]['args']
+    assert event_args['staker'] == staker3
+    assert event_args['value'] == value
+    assert event_args['firstPeriod'] == current_period + 1
+    assert event_args['periods'] == duration
+
+    events = wind_down_log.get_all_entries()
+    assert len(events) == 2