feat: allow token recipient to collect token dust (#9)

* feat: allow token recipient to collect others token dust

* Apply suggestions from code review

Co-authored-by: banteg <[email protected]>

* bump version

Co-authored-by: banteg <[email protected]>

contracts/VestingEscrowSimple.vy

@@ -195,3 +195,9 @@ def renounce_ownership():
     self.future_admin = ZERO_ADDRESS
     self.admin = ZERO_ADDRESS
     log ApplyOwnership(ZERO_ADDRESS)
+
+@external
+def collect_dust(token: address):
+    assert msg.sender == self.recipient  # dev: recipient only
+    assert (token != self.token.address or block.timestamp > self.disabled_at)
+    assert ERC20(token).transfer(self.recipient, ERC20(token).balanceOf(self))

requirements-dev.txt

@@ -1,2 +1,2 @@
-black==19.10b0
-eth-brownie>=1.13.0,<2.0.0
+black==21.9b0
+eth-brownie>=1.17.0,<2.0.0

scripts/deploy_empty.py

@@ -1,6 +1,7 @@
 from brownie import VestingEscrowSimple, VestingEscrowFactory, accounts
 
+
 def main():
-    admin = accounts.load('deployer')
+    admin = accounts.load("deployer")
     template = VestingEscrowSimple.deploy({"from": admin})
     factory = VestingEscrowFactory.deploy(template, {"from": admin})

tests/conftest.py

@@ -51,9 +51,7 @@ def vesting_target(VestingEscrowSimple, accounts):
 
 @pytest.fixture(scope="module")
 def vesting_factory(VestingEscrowFactory, accounts, vesting_target):
-    yield VestingEscrowFactory.deploy(
-        vesting_target, {"from": accounts[0]}
-    )
+    yield VestingEscrowFactory.deploy(vesting_target, {"from": accounts[0]})
 
 
 @pytest.fixture(scope="module")

tests/functional/VestingEscrow/test_collect_dust.py

@@ -0,0 +1,32 @@
+import brownie
+import pytest
+from brownie import ZERO_ADDRESS
+
+
+@pytest.fixture
+def token2(ERC20, accounts):
+    yield ERC20.deploy("XYZ", "XYZ", 18, {"from": accounts[0]})
+
+
+def test_claim_non_vested_token(vesting, token, token2, accounts, chain, end_time):
+    token2._mint_for_testing(10 ** 20, {"from": accounts[0]})
+    token2.transfer(vesting, 10 ** 20)
+
+    vesting.collect_dust(token2, {"from": accounts[1]})
+    assert token2.balanceOf(accounts[1]) == 10 ** 20
+
+
+def test_do_not_allow_claim_of_vested_token(
+    vesting, token, token2, accounts, chain, end_time
+):
+    with brownie.reverts():
+        vesting.collect_dust(token, {"from": accounts[1]})
+
+
+def test_allow_vested_token_dust_to_be_claim_at_end(
+    vesting, token, accounts, chain, end_time
+):
+    chain.sleep(end_time - chain.time() + 1)
+    chain.mine()
+    vesting.collect_dust(token, {"from": accounts[1]})
+    assert token.balanceOf(accounts[1]) == 10 ** 20