From 8c5fedf3beb851cae61e6671cc5a1ca53d7fc3a2 Mon Sep 17 00:00:00 2001
From: trizin <25263018+trizin@users.noreply.github.com>
Date: Mon, 20 Nov 2023 13:07:45 +0300
Subject: [PATCH] #378 - Payout; auto convert wROSE to ROSE (#379)

* Add WrappedToken class

* Add docstring

* Update payout module

* Check zero balance

* Update tests

* Consistency

* Inherit from Token class
---
 pdr_backend/models/dfrewards.py            |  3 ++
 pdr_backend/models/wrapped_token.py        | 34 ++++++++++++++++++++++
 pdr_backend/predictoor/payout.py           | 17 +++++++++--
 pdr_backend/predictoor/test/test_payout.py | 16 ++++++++--
 4 files changed, 65 insertions(+), 5 deletions(-)
 create mode 100644 pdr_backend/models/wrapped_token.py

diff --git a/pdr_backend/models/dfrewards.py b/pdr_backend/models/dfrewards.py
index d6dd8ec3c..6d67b5dda 100644
--- a/pdr_backend/models/dfrewards.py
+++ b/pdr_backend/models/dfrewards.py
@@ -19,6 +19,9 @@ def claim_rewards(self, user_addr: str, token_addr: str, wait_for_receipt=True):
         return self.config.w3.eth.wait_for_transaction_receipt(tx)
 
     def get_claimable_rewards(self, user_addr: str, token_addr: str) -> float:
+        """
+        Returns the amount of claimable rewards in units of ETH
+        """
         claimable_wei = self.contract_instance.functions.claimable(
             user_addr, token_addr
         ).call()
diff --git a/pdr_backend/models/wrapped_token.py b/pdr_backend/models/wrapped_token.py
new file mode 100644
index 000000000..af548ef0e
--- /dev/null
+++ b/pdr_backend/models/wrapped_token.py
@@ -0,0 +1,34 @@
+from pdr_backend.models.token import Token
+from pdr_backend.util.web3_config import Web3Config
+
+
+class WrappedToken(Token):
+    def __init__(self, config: Web3Config, address: str):
+        super().__init__(config, address)
+        abi = [
+            {
+                "constant": False,
+                "inputs": [{"type": "uint256"}],
+                "name": "withdraw",
+                "outputs": [],
+                "payable": False,
+                "stateMutability": "external",
+                "type": "function",
+            },
+        ]
+        self.contract_instance_wrapped = config.w3.eth.contract(
+            address=self.contract_address, abi=abi
+        )
+
+    def withdraw(self, amount: int, wait_for_receipt=True):
+        """
+        Converts Wrapped Token to Token, amount is in wei.
+        """
+        gas_price = self.config.w3.eth.gas_price
+
+        tx = self.contract_instance_wrapped.functions.withdraw(amount).transact(
+            {"from": self.config.owner, "gasPrice": gas_price}
+        )
+        if not wait_for_receipt:
+            return tx
+        return self.config.w3.eth.wait_for_transaction_receipt(tx)
diff --git a/pdr_backend/predictoor/payout.py b/pdr_backend/predictoor/payout.py
index 8f9491adf..0cf1894f1 100644
--- a/pdr_backend/predictoor/payout.py
+++ b/pdr_backend/predictoor/payout.py
@@ -5,6 +5,7 @@
 from pdr_backend.models.base_config import BaseConfig
 
 from pdr_backend.models.predictoor_contract import PredictoorContract
+from pdr_backend.models.wrapped_token import WrappedToken
 from pdr_backend.util.subgraph import query_pending_payouts, wait_until_subgraph_syncs
 from pdr_backend.models.dfrewards import DFRewards
 
@@ -65,12 +66,22 @@ def do_rose_payout():
         raise Exception("Unsupported network")
     contract = DFRewards(config.web3_config, address)
     claimable_rewards = contract.get_claimable_rewards(owner, wrapped_rose)
-    print(f"Found {claimable_rewards} ROSE available to claim")
+    print(f"Found {claimable_rewards} wROSE available to claim")
 
     if claimable_rewards > 0:
-        print("Claiming ROSE rewards...")
+        print("Claiming wROSE rewards...")
         contract.claim_rewards(owner, wrapped_rose)
     else:
-        print("No rewards, exiting")
+        print("No rewards available to claim")
+
+    print("Converting wROSE to ROSE")
+    time.sleep(10)
+    wrose = WrappedToken(config.web3_config, wrapped_rose)
+    wrose_balance = wrose.balanceOf(config.web3_config.owner)
+    if wrose_balance == 0:
+        print("wROSE balance is 0")
+    else:
+        print(f"Found {wrose_balance/1e18} wROSE, converting to ROSE...")
+        wrose.withdraw(wrose_balance)
 
     print("ROSE reward claim done")
diff --git a/pdr_backend/predictoor/test/test_payout.py b/pdr_backend/predictoor/test/test_payout.py
index 80c154c4b..b9df3f6e3 100644
--- a/pdr_backend/predictoor/test/test_payout.py
+++ b/pdr_backend/predictoor/test/test_payout.py
@@ -4,6 +4,7 @@
 
 from pdr_backend.models.dfrewards import DFRewards
 from pdr_backend.models.predictoor_contract import PredictoorContract
+from pdr_backend.models.wrapped_token import WrappedToken
 from pdr_backend.predictoor.payout import (
     batchify,
     do_payout,
@@ -91,10 +92,21 @@ def test_do_rose_payout():
     mock_contract.get_claimable_rewards.return_value = 100
     mock_contract.claim_rewards = Mock()
 
-    with patch(
+    mock_wrose = Mock(spec=WrappedToken)
+    mock_wrose.balanceOf = Mock()
+    mock_wrose.balanceOf.return_value = 100
+    mock_wrose.withdraw = Mock()
+
+    with patch("pdr_backend.predictoor.payout.time"), patch(
         "pdr_backend.predictoor.payout.BaseConfig", return_value=mock_config
-    ), patch("pdr_backend.predictoor.payout.DFRewards", return_value=mock_contract):
+    ), patch(
+        "pdr_backend.predictoor.payout.WrappedToken", return_value=mock_wrose
+    ), patch(
+        "pdr_backend.predictoor.payout.DFRewards", return_value=mock_contract
+    ):
         do_rose_payout()
         mock_contract.claim_rewards.assert_called_with(
             "mock_owner", "0x8Bc2B030b299964eEfb5e1e0b36991352E56D2D3"
         )
+        mock_wrose.balanceOf.assert_called()
+        mock_wrose.withdraw.assert_called_with(100)