Update CLI ROSE and OCEAN payout to work with pred submitter manager

pdr_backend/cli/cli_module.py

@@ -1,5 +1,4 @@
 import logging
-import os
 import sys
 
 from enforce_typing import enforce_types
@@ -126,14 +125,6 @@ def do_claim_OCEAN(args, nested_args=None):
     )
     do_ocean_payout(ppss)
 
-    # check if there's a second pk
-    pk2 = os.getenv("PRIVATE_KEY2")
-    if pk2 is None:
-        return
-    web3_config = ppss.web3_pp.web3_config.copy_with_pk(pk2)
-    ppss.web3_pp.set_web3_config(web3_config)
-    do_ocean_payout(ppss)
-
 
 @enforce_types
 def do_claim_ROSE(args, nested_args=None):
@@ -144,14 +135,6 @@ def do_claim_ROSE(args, nested_args=None):
     )
     do_rose_payout(ppss)
 
-    # check if there's a second pk
-    pk2 = os.getenv("PRIVATE_KEY2")
-    if pk2 is None:
-        return
-    web3_config = ppss.web3_pp.web3_config.copy_with_pk(pk2)
-    ppss.web3_pp.set_web3_config(web3_config)
-    do_rose_payout(ppss)
-
 
 @enforce_types
 def do_multisim(args, nested_args=None):

pdr_backend/payout/payout.py

@@ -5,8 +5,10 @@
 from enforce_typing import enforce_types
 
 from pdr_backend.contract.dfrewards import DFRewards
+from pdr_backend.contract.pred_submitter_mgr import PredSubmitterMgr
 from pdr_backend.contract.feed_contract import FeedContract
 from pdr_backend.contract.wrapped_token import WrappedToken
+from pdr_backend.predictoor.util import find_shared_slots, to_checksum
 from pdr_backend.ppss.ppss import PPSS
 from pdr_backend.subgraph.subgraph_pending_payouts import query_pending_payouts
 from pdr_backend.subgraph.subgraph_sync import wait_until_subgraph_syncs
@@ -49,29 +51,41 @@ def request_payout_batches(
 
 
 @enforce_types
-def do_ocean_payout(ppss: PPSS, check_network: bool = True):
+def find_slots_and_payout_with_mgr(pred_submitter_mgr, ppss):
+    # we only need to query in one direction, since both predict on the same slots
+    up_addr = pred_submitter_mgr.pred_submitter_up_address()
     web3_config = ppss.web3_pp.web3_config
     subgraph_url: str = ppss.web3_pp.subgraph_url
+    logger.info("Starting payout")
+    wait_until_subgraph_syncs(web3_config, subgraph_url)
+    logger.info("Finding pending payouts")
+    pending_slots = query_pending_payouts(subgraph_url, up_addr)
+    shared_slots = find_shared_slots(pending_slots)
+    if not shared_slots:
+        logger.info("No payouts available")
+        return
+    logger.info("Found %d slots", len(shared_slots))
+    for slot_tuple in shared_slots:
+        contract_addrs, slots = slot_tuple
+        contract_addrs = to_checksum(ppss.web3_pp.w3, contract_addrs)
+        tx = pred_submitter_mgr.get_payout(slots, contract_addrs)
+        cur_index = shared_slots.index(slot_tuple)
+        progress = f"{cur_index + 1}/{len(shared_slots)}"
+        logger.info("Payout tx %s: %s", progress, tx["transactionHash"].hex())
+    logger.info("Payout done")
 
+
+@enforce_types
+def do_ocean_payout(ppss: PPSS, check_network: bool = True):
+    web3_config = ppss.web3_pp.web3_config
     if check_network:
         assert ppss.web3_pp.network == "sapphire-mainnet"
         assert web3_config.w3.eth.chain_id == SAPPHIRE_MAINNET_CHAINID
 
-    logger.info("Starting payout")
-    wait_until_subgraph_syncs(web3_config, subgraph_url)
-    logger.info("Finding pending payouts")
-    pending_payouts = query_pending_payouts(subgraph_url, web3_config.owner)
-    total_timestamps = sum(len(timestamps) for timestamps in pending_payouts.values())
-    logger.info("Found %d slots", total_timestamps)
-
-    for pdr_contract_addr in pending_payouts:
-        logger.info("Claiming payouts for %s", pdr_contract_addr)
-        pdr_contract = FeedContract(ppss.web3_pp, pdr_contract_addr)
-        request_payout_batches(
-            pdr_contract, ppss.payout_ss.batch_size, pending_payouts[pdr_contract_addr]
-        )
+    pred_submitter_mgr_addr = ppss.predictoor_ss.pred_submitter_mgr
+    pred_submitter_mgr = PredSubmitterMgr(ppss.web3_pp, pred_submitter_mgr_addr)
 
-    logger.info("Payout done")
+    find_slots_and_payout_with_mgr(pred_submitter_mgr, ppss)
 
 
 @enforce_types
@@ -82,24 +96,51 @@ def do_rose_payout(ppss: PPSS, check_network: bool = True):
         assert ppss.web3_pp.network == "sapphire-mainnet"
         assert web3_config.w3.eth.chain_id == SAPPHIRE_MAINNET_CHAINID
 
+    web3_config = ppss.web3_pp.web3_config
+    pred_submitter_mgr_addr = ppss.predictoor_ss.pred_submitter_mgr
+    pred_submitter_mgr = PredSubmitterMgr(ppss.web3_pp, pred_submitter_mgr_addr)
+    up_addr = pred_submitter_mgr.pred_submitter_up_address()
+    down_addr = pred_submitter_mgr.pred_submitter_down_address()
+
     dfrewards_addr = "0xc37F8341Ac6e4a94538302bCd4d49Cf0852D30C0"
     wROSE_addr = "0x8Bc2B030b299964eEfb5e1e0b36991352E56D2D3"
+    wROSE = WrappedToken(ppss.web3_pp, wROSE_addr)
 
     dfrewards_contract = DFRewards(ppss.web3_pp, dfrewards_addr)
-    claimable_rewards = dfrewards_contract.get_claimable_rewards(
-        web3_config.owner, wROSE_addr
+    claimable_rewards_up = dfrewards_contract.get_claimable_rewards(up_addr, wROSE_addr)
+    claimable_rewards_down = dfrewards_contract.get_claimable_rewards(
+        down_addr, wROSE_addr
     )
-    logger.info("Found %s wROSE available to claim", claimable_rewards.amt_eth)
-
-    if claimable_rewards > Eth(0):
-        logger.info("Claiming wROSE rewards...")
-        dfrewards_contract.claim_rewards(web3_config.owner, wROSE_addr)
+    total_claimable = claimable_rewards_up + claimable_rewards_down
+    logger.info("Found %s wROSE available to claim", total_claimable.to_eth().amt_eth)
+
+    if total_claimable > Eth(0):
+        logger.info("Claiming wROSE rewards from the manager contract...")
+        receipt = pred_submitter_mgr.claim_dfrewards(wROSE_addr, dfrewards_addr, True)
+        if receipt["status"] != 1:
+            logger.warning(
+                "Failed to claim wROSE rewards from the contract, tx: %s",
+                receipt["transactionHash"],
+            )
+            return
+        logger.info("Transfering wROSE to owner")
+        time.sleep(4)
+        balance = wROSE.balanceOf(pred_submitter_mgr.contract_address)
+        receipt = pred_submitter_mgr.transfer_erc20(
+            wROSE_addr, web3_config.owner, balance, True
+        )
+        if receipt["status"] != 1:
+            logger.warning(
+                "Failed to transfer wROSE tokens to the owner, tx: %s",
+                receipt["transactionHash"],
+            )
+            return
     else:
         logger.warning("No rewards available to claim")
+        return
 
-    logger.info("Sleeping")
-    time.sleep(10)
-    wROSE = WrappedToken(ppss.web3_pp, wROSE_addr)
+    logger.info("Converting wROSE to ROSE")
+    time.sleep(4)
     wROSE_bal = wROSE.balanceOf(web3_config.owner)
     if wROSE_bal == Wei(0):
         logger.warning("wROSE balance is 0")

pdr_backend/payout/test/test_payout.py

@@ -3,7 +3,7 @@
 import pytest
 from enforce_typing import enforce_types
 
-from pdr_backend.contract.dfrewards import DFRewards
+from pdr_backend.contract.pred_submitter_mgr import PredSubmitterMgr
 from pdr_backend.contract.feed_contract import FeedContract
 from pdr_backend.contract.wrapped_token import WrappedToken
 from pdr_backend.payout.payout import (
@@ -52,37 +52,50 @@ def test_do_ocean_payout(tmpdir):
     ppss = _ppss(tmpdir)
     ppss.payout_ss.set_batch_size(5)
 
+    mock_addr_1 = "0x0000000000000000000000000000000000000000"
+    mock_addr_2 = "0x0000000000000000000000000000000000000001"
     mock_pending_payouts = {
-        "address_1": [1, 2, 3],
-        "address_2": [4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
+        mock_addr_1: [1, 2, 3, 4, 5],
+        mock_addr_2: [4, 5, 6, 7, 8, 9, 10],
     }
 
-    mock_contract = Mock(spec=FeedContract)
-    mock_contract.payout_multiple = Mock()
+    mock_contract = Mock(spec=PredSubmitterMgr)
+    mock_contract.get_payout = Mock()
+    mock_contract.get_payout.return_value = {"transactionHash": b"0x1", "status": 1}
+    mock_contract.pred_submitter_up_address.return_value = "0x1"
+    mock_contract.pred_submitter_down_address.return_value = "0x1"
 
     with patch("pdr_backend.payout.payout.wait_until_subgraph_syncs"), patch(
         "pdr_backend.payout.payout.query_pending_payouts",
         return_value=mock_pending_payouts,
-    ), patch("pdr_backend.payout.payout.FeedContract", return_value=mock_contract):
+    ), patch("pdr_backend.payout.payout.PredSubmitterMgr", return_value=mock_contract):
         do_ocean_payout(ppss, check_network=False)
-        print(mock_contract.payout_multiple.call_args_list)
-        call_args = mock_contract.payout_multiple.call_args_list
-        assert call_args[0] == call([1, 2, 3], True)
-        assert call_args[1] == call([4, 5, 6, 7, 8], True)
-        assert call_args[2] == call([9, 10, 11, 12, 13], True)
-        assert call_args[3] == call([14, 15], True)
-        assert len(call_args) == 4
+        print(mock_contract.get_payout.call_args_list)
+        call_args = mock_contract.get_payout.call_args_list
+        assert call_args[0] == call([1, 2, 3], [mock_addr_1])
+        assert call_args[1] == call([4, 5], [mock_addr_1, mock_addr_2])
+        assert call_args[2] == call([6, 7, 8, 9, 10], [mock_addr_2])
+        assert len(call_args) == 3
 
 
 @enforce_types
 def test_do_rose_payout(tmpdir):
     ppss = _ppss(tmpdir)
     web3_config = ppss.web3_pp.web3_config
 
-    mock_contract = Mock(spec=DFRewards)
+    mock_contract = Mock(spec=PredSubmitterMgr)
     mock_contract.get_claimable_rewards = Mock()
     mock_contract.get_claimable_rewards.return_value = Eth(100)
-    mock_contract.claim_rewards = Mock()
+    mock_contract.contract_address = "0x0"
+    mock_contract.claim_dfrewards = Mock()
+    mock_contract.claim_dfrewards.return_value = {
+        "transactionHash": b"0x1",
+        "status": 1,
+    }
+    mock_contract.transfer_erc20 = Mock()
+    mock_contract.transfer_erc20.return_value = {"transactionHash": b"0x1", "status": 1}
+    mock_contract.pred_submitter_up_address.return_value = "0x1"
+    mock_contract.pred_submitter_down_address.return_value = "0x2"
 
     mock_wrose = Mock(spec=WrappedToken)
     mock_wrose.balanceOf = Mock()
@@ -91,10 +104,20 @@ def test_do_rose_payout(tmpdir):
 
     with patch("pdr_backend.payout.payout.time"), patch(
         "pdr_backend.payout.payout.WrappedToken", return_value=mock_wrose
-    ), patch("pdr_backend.payout.payout.DFRewards", return_value=mock_contract):
+    ), patch("pdr_backend.payout.payout.DFRewards", return_value=mock_contract), patch(
+        "pdr_backend.payout.payout.PredSubmitterMgr", return_value=mock_contract
+    ):
         do_rose_payout(ppss, check_network=False)
-        mock_contract.claim_rewards.assert_called_with(
-            web3_config.owner, "0x8Bc2B030b299964eEfb5e1e0b36991352E56D2D3"
+        mock_contract.claim_dfrewards.assert_called_with(
+            "0x8Bc2B030b299964eEfb5e1e0b36991352E56D2D3",
+            "0xc37F8341Ac6e4a94538302bCd4d49Cf0852D30C0",
+            True,
+        )
+        mock_contract.transfer_erc20.assert_called_with(
+            "0x8Bc2B030b299964eEfb5e1e0b36991352E56D2D3",
+            web3_config.owner,
+            Eth(100).to_wei(),
+            True,
         )
         mock_wrose.balanceOf.assert_called()
         mock_wrose.withdraw.assert_called_with(Eth(100).to_wei())

pdr_backend/predictoor/predictoor_agent.py

@@ -15,9 +15,9 @@
 from pdr_backend.ppss.ppss import PPSS
 from pdr_backend.predictoor.predictoor_logger import PredictoorAgentLogLine
 from pdr_backend.predictoor.stakes_per_slot import StakesPerSlot, StakeTup
-from pdr_backend.predictoor.util import find_shared_slots
+from pdr_backend.predictoor.util import to_checksum
+from pdr_backend.payout.payout import find_slots_and_payout_with_mgr
 from pdr_backend.subgraph.subgraph_feed import SubgraphFeed, print_feeds
-from pdr_backend.subgraph.subgraph_pending_payouts import query_pending_payouts
 from pdr_backend.util.currency_types import Eth, Wei
 from pdr_backend.util.logutil import logging_has_stdout
 from pdr_backend.util.time_types import UnixTimeS
@@ -226,8 +226,7 @@ def take_step(self):
     @enforce_types
     def _to_checksum(self, addrs: List[str]) -> List[str]:
         w3 = self.ppss.web3_pp.w3
-        checksummed_addrs = [w3.to_checksum_address(addr) for addr in addrs]
-        return checksummed_addrs
+        return to_checksum(w3, addrs)
 
     @property
     def cur_block(self):
@@ -474,23 +473,7 @@ def get_payout(self):
 
         # Update previous payouts history to avoid claiming for this epoch again
         self.prev_submit_payouts.append(self.cur_unique_epoch)
-
-        # we only need to query in one direction, since both predict on the same slots
-        up_addr = self.pred_submitter_mgr.pred_submitter_up_address()
-        pending_slots = query_pending_payouts(self.ppss.web3_pp.subgraph_url, up_addr)
-        shared_slots = find_shared_slots(pending_slots)
-        if not shared_slots:
-            logger.info("No payouts available")
-            return
-
-        for slot_tuple in shared_slots:
-            contract_addrs, slots = slot_tuple
-            contract_addrs = self._to_checksum(contract_addrs)
-            tx = self.pred_submitter_mgr.get_payout(slots, contract_addrs)
-
-            cur_index = shared_slots.index(slot_tuple)
-            progress = f"{cur_index + 1}/{len(shared_slots)}"
-            logger.info("Payout tx %s: %s", progress, tx["transactionHash"].hex())
+        find_slots_and_payout_with_mgr(self.pred_submitter_mgr, self.ppss)
 
 
 @enforce_types

pdr_backend/predictoor/util.py

@@ -38,3 +38,8 @@ def find_shared_slots(
         result.append(tup)
 
     return result
+
+
+def to_checksum(w3, addrs: List[str]) -> List[str]:
+    checksummed_addrs = [w3.to_checksum_address(addr) for addr in addrs]
+    return checksummed_addrs

system_tests/test_ocean_payout.py

@@ -3,7 +3,6 @@
 from unittest.mock import Mock, patch, MagicMock
 
 from pdr_backend.cli import cli_module
-from pdr_backend.contract.feed_contract import FeedContract
 from pdr_backend.ppss.web3_pp import Web3PP
 from pdr_backend.util.constants import SAPPHIRE_MAINNET_CHAINID
 from pdr_backend.util.web3_config import Web3Config
@@ -22,27 +21,36 @@ def test_ocean_payout_test(mock_wait_until_subgraph_syncs, caplog):
     mock_web3_config.owner = "0x00000000000000000000000000000000000c0ffe"
     mock_web3_config.w3 = Mock()
     mock_web3_config.w3.eth.chain_id = SAPPHIRE_MAINNET_CHAINID
+
+    def checksum_mock(_, y):
+        return y
+
+    mock_web3_config.w3.to_checksum_address = lambda x: x
     mock_web3_pp.web3_config = mock_web3_config
 
+    mock_contract = Mock()
+    mock_contract.get_payout = Mock()
+    mock_contract.get_payout.return_value = {"transactionHash": b"0x1", "status": 1}
+    mock_contract.pred_submitter_up_address.return_value = "0x1"
+    mock_contract.pred_submitter_down_address.return_value = "0x1"
+
     mock_token = MagicMock()
     mock_token.balanceOf.return_value = 100 * 1e18
 
     mock_query_pending_payouts = Mock()
     mock_query_pending_payouts.return_value = {"0x1": [1, 2, 3], "0x2": [5, 6, 7]}
-    number_of_slots = 6  # 3 + 3
-
-    mock_feed_contract = Mock(spec=FeedContract)
-    mock_feed_contract.payout_multiple.return_value = None
+    number_of_slots = 2  # 2 addresses
 
     with patch(
         "pdr_backend.payout.payout.query_pending_payouts", mock_query_pending_payouts
     ), patch("pdr_backend.ppss.ppss.Web3PP", return_value=mock_web3_pp), patch(
         "pdr_backend.contract.token.Token", return_value=mock_token
+    ), patch(
+        "pdr_backend.payout.payout.to_checksum", checksum_mock
     ), patch(
         "pdr_backend.payout.payout.WrappedToken", return_value=mock_token
     ), patch(
-        "pdr_backend.payout.payout.FeedContract",
-        return_value=mock_feed_contract,
+        "pdr_backend.payout.payout.PredSubmitterMgr", return_value=mock_contract
     ):
         # Mock sys.argv
         sys.argv = ["pdr", "claim_OCEAN", "ppss.yaml"]
@@ -56,14 +64,13 @@ def test_ocean_payout_test(mock_wait_until_subgraph_syncs, caplog):
         assert "Starting payout" in caplog.text
         assert "Finding pending payouts" in caplog.text
         assert f"Found {number_of_slots} slots" in caplog.text
-        assert "Claiming payouts for 0x1" in caplog.text
-        assert "Claiming payouts for 0x2" in caplog.text
+        assert "Payout tx 1/2: 307831" in caplog.text
+        assert "Payout tx 1/2: 307831" in caplog.text
         assert "Payout done" in caplog.text
 
         # Additional assertions
-        mock_query_pending_payouts.assert_called_with(
-            mock_web3_pp.subgraph_url, mock_web3_config.owner
-        )
-        mock_feed_contract.payout_multiple.assert_any_call([1, 2, 3], True)
-        mock_feed_contract.payout_multiple.assert_any_call([5, 6, 7], True)
-        assert mock_feed_contract.payout_multiple.call_count == 4
+        mock_query_pending_payouts.assert_called_with(mock_web3_pp.subgraph_url, "0x1")
+        print(mock_contract.get_payout.call_args_list)
+        mock_contract.get_payout.assert_any_call([1, 2, 3], ["0x1"])
+        mock_contract.get_payout.assert_any_call([5, 6, 7], ["0x2"])
+        assert mock_contract.get_payout.call_count == 2