From b9231f27d910aee52aa82e77f212feb64a3e174e Mon Sep 17 00:00:00 2001
From: pandadefi <87183122+pandadefi@users.noreply.github.com>
Date: Mon, 1 Nov 2021 16:46:17 +0200
Subject: [PATCH] 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 <banteeg@gmail.com>

* bump version

Co-authored-by: banteg <banteeg@gmail.com>
---
 contracts/VestingEscrowSimple.vy              |  6 ++++
 requirements-dev.txt                          |  4 +--
 scripts/deploy_empty.py                       |  3 +-
 tests/conftest.py                             |  4 +--
 .../VestingEscrow/test_collect_dust.py        | 32 +++++++++++++++++++
 5 files changed, 43 insertions(+), 6 deletions(-)
 create mode 100644 tests/functional/VestingEscrow/test_collect_dust.py

diff --git a/contracts/VestingEscrowSimple.vy b/contracts/VestingEscrowSimple.vy
index 78cf4db..af50352 100644
--- a/contracts/VestingEscrowSimple.vy
+++ b/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))
diff --git a/requirements-dev.txt b/requirements-dev.txt
index 127231b..e1b974b 100644
--- a/requirements-dev.txt
+++ b/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
diff --git a/scripts/deploy_empty.py b/scripts/deploy_empty.py
index 9982202..63e6bb4 100644
--- a/scripts/deploy_empty.py
+++ b/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})
diff --git a/tests/conftest.py b/tests/conftest.py
index 483daf7..e055e8c 100644
--- a/tests/conftest.py
+++ b/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")
diff --git a/tests/functional/VestingEscrow/test_collect_dust.py b/tests/functional/VestingEscrow/test_collect_dust.py
new file mode 100644
index 0000000..d316ed4
--- /dev/null
+++ b/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