From 6af7813e6aa26cf9b192123ba240d642cf5c08d0 Mon Sep 17 00:00:00 2001
From: Kieran Prasch <kieran@nucypher.com>
Date: Mon, 16 Sep 2019 20:59:28 -0700
Subject: [PATCH] Restaking control via Staker Actor.

---
 nucypher/blockchain/eth/actors.py             | 21 ++++++++++++++
 nucypher/blockchain/eth/agents.py             |  4 +--
 .../eth/entities/actors/test_staker.py        | 28 +++++++++++++++++++
 .../agents/test_staking_escrow_agent.py       |  2 +-
 4 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/nucypher/blockchain/eth/actors.py b/nucypher/blockchain/eth/actors.py
index 52c6c823a..b1d386fe6 100644
--- a/nucypher/blockchain/eth/actors.py
+++ b/nucypher/blockchain/eth/actors.py
@@ -564,6 +564,27 @@ class Staker(NucypherTokenActor):
 
         return new_stake
 
+    def enable_restaking(self) -> dict:
+        receipt = self.staking_agent.set_restaking(staker_address=self.checksum_address, value=True)
+        return receipt
+
+    def enable_restaking_lock(self, release_period: int):
+        current_period = self.staking_agent.get_current_period()
+        if release_period < current_period:
+            raise ValueError(f"Terminal restaking period must be in the future.  "
+                             f"Current period is {current_period}, got '{release_period}'.")
+        receipt = self.staking_agent.lock_restaking(staker_address=self.checksum_address,
+                                                    release_period=release_period)
+        return receipt
+
+    @property
+    def restaking_lock_enabled(self) -> bool:
+        status = self.staking_agent.get_restaking_lock_status(staker_address=self.checksum_address)
+        return status
+
+    def disable_restaking(self) -> dict:
+        receipt = self.staking_agent.set_restaking(staker_address=self.checksum_address, value=False)
+        return receipt
     #
     # Reward and Collection
     #
diff --git a/nucypher/blockchain/eth/agents.py b/nucypher/blockchain/eth/agents.py
index 36b62254b..ab2a83564 100644
--- a/nucypher/blockchain/eth/agents.py
+++ b/nucypher/blockchain/eth/agents.py
@@ -377,8 +377,8 @@ class StakingEscrowAgent(EthereumContractAgent):
                                                    sender_address=staker_address)
         return receipt
 
-    def lock_restaking(self, staker_address: str, termination_period: int) -> dict:
-        contract_function = self.contract.functions.lockReStake(termination_period)
+    def lock_restaking(self, staker_address: str, release_period: int) -> dict:
+        contract_function = self.contract.functions.lockReStake(release_period)
         receipt = self.blockchain.send_transaction(contract_function=contract_function,
                                                    sender_address=staker_address)
         return receipt
diff --git a/tests/blockchain/eth/entities/actors/test_staker.py b/tests/blockchain/eth/entities/actors/test_staker.py
index cd606b351..faee6b528 100644
--- a/tests/blockchain/eth/entities/actors/test_staker.py
+++ b/tests/blockchain/eth/entities/actors/test_staker.py
@@ -17,8 +17,10 @@ along with nucypher.  If not, see <https://www.gnu.org/licenses/>.
 
 
 import pytest
+from eth_tester.exceptions import TransactionFailed
 
 from nucypher.blockchain.eth.actors import Staker
+from nucypher.blockchain.eth.agents import ContractAgency, StakingEscrowAgent
 from nucypher.blockchain.eth.token import NU, Stake
 from nucypher.crypto.powers import TransactingPower
 from nucypher.utilities.sandbox.blockchain import token_airdrop
@@ -90,6 +92,32 @@ def test_staker_divides_stake(staker, token_economics):
     assert expected_yet_another_stake == staker.stakes[stake_index + 3], 'Third stake values are invalid'
 
 
+def test_staker_manages_restaking(testerchain, test_registry, staker):
+
+    # Enable Restaking
+    receipt = staker.enable_restaking()
+    assert receipt['status'] == 1
+
+    # Enable Restaking Lock
+    staking_agent = ContractAgency.get_agent(StakingEscrowAgent, registry=test_registry)
+    current_period = staking_agent.get_current_period()
+    terminal_period = current_period + 2
+
+    assert not staker.restaking_lock_enabled
+    receipt = staker.enable_restaking_lock(release_period=terminal_period)
+    assert receipt['status'] == 1
+    assert staker.restaking_lock_enabled
+
+    with pytest.raises((TransactionFailed, ValueError)):
+      staker.disable_restaking()
+
+    # Wait until terminal period
+    testerchain.time_travel(periods=2)
+    receipt = staker.disable_restaking()
+    assert receipt['status'] == 1
+    assert not staker.restaking_lock_enabled
+
+
 @pytest.mark.slow()
 def test_staker_collects_staking_reward(testerchain,
                                         test_registry,
diff --git a/tests/blockchain/eth/entities/agents/test_staking_escrow_agent.py b/tests/blockchain/eth/entities/agents/test_staking_escrow_agent.py
index 20794a48b..7b35c86d8 100644
--- a/tests/blockchain/eth/entities/agents/test_staking_escrow_agent.py
+++ b/tests/blockchain/eth/entities/agents/test_staking_escrow_agent.py
@@ -238,7 +238,7 @@ def test_lock_restaking(agency, testerchain):
     staker_account, worker_account, *other = testerchain.unassigned_accounts
     current_period = staking_agent.get_current_period()
     terminal_period = current_period + 2
-    receipt = staking_agent.lock_restaking(staker_account, termination_period=terminal_period)
+    receipt = staking_agent.lock_restaking(staker_account, release_period=terminal_period)
     assert receipt['status'] == 1, "Transaction Rejected"