Apply suggestions from code review #1788

Co-Authored-By: Derek Pierre <derek.pierre@gmail.com>
pull/1788/head
vzotova 2020-03-27 12:01:40 +03:00
parent 69d535ef21
commit 48379354bc
2 changed files with 5 additions and 2 deletions

View File

@ -1127,7 +1127,7 @@ class Staker(NucypherTokenActor):
"""Collect rewarded ETH."""
if self.is_contract:
if collector_address and collector_address != self.beneficiary_address:
raise ValueError("Allowed to withdraw policy reward only to the beneficiary address")
raise ValueError("Policy rewards must be withdrawn to the beneficiary address")
self.preallocation_escrow_agent.collect_policy_reward() # TODO save receipt
receipt = self.preallocation_escrow_agent.withdraw_eth()
else:

View File

@ -184,8 +184,11 @@ def test_withdraw_eth(testerchain, preallocation_escrow):
{'from': testerchain.client.coinbase, 'to': preallocation_escrow.address, 'value': value})
testerchain.wait_for_receipt(tx)
tx = preallocation_escrow.functions.withdrawETH().transact({'from': owner})
balance = testerchain.client.get_balance(owner)
tx = preallocation_escrow.functions.withdrawETH().transact({'from': owner, 'gas_price': 0})
testerchain.wait_for_receipt(tx)
assert testerchain.client.get_balance(preallocation_escrow.address) == 0
assert testerchain.client.get_balance(owner) == balance + value
events = log.get_all_entries()
assert len(events) == 1