Review fixes

Signed-off-by: cyc60 <avsysoev60@gmail.com>

Author: cyc60

src/commands/internal/process_redeemer.py

@@ -5,7 +5,7 @@
 from pathlib import Path
 
 import click
-from eth_typing import BlockNumber, ChecksumAddress, HexStr
+from eth_typing import BlockNumber, ChecksumAddress
 from hexbytes import HexBytes
 from multiproof import StandardMerkleTree
 from multiproof.standard import MultiProof
@@ -18,8 +18,6 @@
 from src.common.contracts import (
     MetaVaultContract,
     SubVaultsRegistryContract,
-    VaultContract,
-    multicall_contract,
     os_token_redeemer_contract,
 )
 from src.common.execution import (
@@ -29,6 +27,7 @@
 )
 from src.common.harvest import get_multiple_harvest_params
 from src.common.logging import LOG_LEVELS, setup_logging
+from src.common.typings import HarvestParams
 from src.common.utils import log_verbose
 from src.common.wallet import wallet
 from src.config.networks import AVAILABLE_NETWORKS, ZERO_CHECKSUM_ADDRESS
@@ -266,7 +265,7 @@ async def _redeem_os_token_positions(
 
     # Bring every involved vault up to date on-chain so subsequent reads
     # and redemption transactions run against fresh state — no harvest_params
-    # plumbing and no updateVaultState bundled in a multicall.
+    # plumbing through redeemOsTokenPositions.
     vaults = list({position.vault for position in os_token_positions})
     await update_vaults_state(vaults=vaults, block_number=block_number)
 
@@ -328,10 +327,11 @@ async def update_vaults_state(
     """Bring every vault in ``vaults`` up to date on-chain.
 
     Meta vaults that need a state update are harvested via process_meta_vault_tree.
-    Regular vaults with pending updates are batched into a single multicall
-    submitting ``updateState`` for each. After this call, ``get_withdrawable_assets``
-    reflects fresh state, so subsequent redemption decisions can be made without
-    threading harvest_params through every read.
+    Regular vaults with pending updates are batched into a single ``multicall`` on
+    the OsTokenRedeemer contract submitting ``updateVaultState`` for each, so the
+    transaction is attributed to the redeemer on Etherscan. After this call,
+    ``get_withdrawable_assets`` reflects fresh state, so subsequent redemption
+    decisions can be made without threading harvest_params through every read.
     """
     regular_vaults: list[ChecksumAddress] = []
     meta_vaults_map = await graph_get_vaults(
@@ -371,28 +371,28 @@ async def update_vaults_state(
     if not regular_vaults:
         return
 
-    vault_to_harvest_params = await get_multiple_harvest_params(regular_vaults, block_number)
-    calls: list[tuple[ChecksumAddress, HexStr]] = []
-    for vault in regular_vaults:
-        harvest_params = vault_to_harvest_params.get(vault)
-        if harvest_params is None:
-            continue
-        vault_contract = VaultContract(vault)
-        calls.append(
-            (vault_contract.contract_address, vault_contract.get_update_state_call(harvest_params))
-        )
-
-    if not calls:
+    vault_to_harvest_params: dict[ChecksumAddress, HarvestParams] = {
+        vault: harvest_params
+        for vault, harvest_params in (
+            await get_multiple_harvest_params(regular_vaults, block_number)
+        ).items()
+        if harvest_params is not None
+    }
+    if not vault_to_harvest_params:
         return
 
-    tx_hash = await multicall_contract.tx_aggregate(calls)
-    logger.info('Waiting for Update State multicall tx %s confirmation', tx_hash)
+    tx_hash = await os_token_redeemer_contract.update_vaults_state(vault_to_harvest_params)
+    logger.info(
+        'Waiting for OsTokenRedeemer updateVaultState multicall tx %s confirmation', tx_hash
+    )
     tx_receipt = await execution_client.eth.wait_for_transaction_receipt(
         HexBytes(Web3.to_bytes(hexstr=tx_hash)), timeout=settings.execution_transaction_timeout
     )
     if not tx_receipt['status']:
-        raise RuntimeError(f'Update State multicall tx failed. Tx Hash: {tx_hash}')
-    logger.info('Update State multicall confirmed. Tx Hash: %s', tx_hash)
+        raise RuntimeError(
+            f'OsTokenRedeemer updateVaultState multicall tx failed. Tx Hash: {tx_hash}'
+        )
+    logger.info('OsTokenRedeemer updateVaultState multicall confirmed. Tx Hash: %s', tx_hash)
 
 
 async def redeem_positions(

src/commands/tests/test_internal/test_process_redeemer.py

@@ -4,7 +4,7 @@
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
-from eth_typing import BlockNumber, ChecksumAddress, HexStr
+from eth_typing import BlockNumber, ChecksumAddress
 from hexbytes import HexBytes
 from sw_utils.tests import faker
 from web3 import Web3
@@ -518,16 +518,15 @@ async def test_meta_vault_claim_delay_logged_and_continues(self) -> None:
         exit_request = MagicMock()
         exit_request.vault = VAULT_1
         exit_request.position_ticket = 1234
+        params = make_harvest_params()
         with _mock_update_vaults_state(
             meta_vaults_map={VAULT_1: MagicMock()},
-            harvest_params={VAULT_2: make_harvest_params()},
+            harvest_params={VAULT_2: params},
             update_state_exception=ClaimDelayNotPassedException(exit_request),
         ) as mocks:
             await update_vaults_state(vaults=[VAULT_1, VAULT_2], block_number=BlockNumber(100))
         mocks['update_state'].assert_awaited_once()
-        mocks['multicall'].tx_aggregate.assert_awaited_once_with(
-            [(VAULT_2, ENCODED_UPDATE_STATE_CALL)]
-        )
+        mocks['redeemer'].update_vaults_state.assert_awaited_once_with({VAULT_2: params})
 
     @pytest.mark.parametrize(
         'has_params, expected_multicall_calls',
@@ -542,25 +541,25 @@ async def test_regular_vault_multicall_gated_by_harvest_params(
         with _mock_update_vaults_state(harvest_params={VAULT_1: params}) as mocks:
             await update_vaults_state(vaults=[VAULT_1], block_number=BlockNumber(100))
 
-        assert mocks['multicall'].tx_aggregate.await_count == expected_multicall_calls
+        assert mocks['redeemer'].update_vaults_state.await_count == expected_multicall_calls
         if has_params:
-            mocks['multicall'].tx_aggregate.assert_awaited_once_with(
-                [(VAULT_1, ENCODED_UPDATE_STATE_CALL)]
-            )
+            mocks['redeemer'].update_vaults_state.assert_awaited_once_with({VAULT_1: params})
 
     async def test_multicall_tx_failure_raises(self) -> None:
         """A failed multicall receipt aborts the round."""
         with _mock_update_vaults_state(
             harvest_params={VAULT_1: make_harvest_params()},
             multicall_tx_status=0,
         ):
-            with pytest.raises(RuntimeError, match='Update State multicall tx failed'):
+            with pytest.raises(
+                RuntimeError, match='OsTokenRedeemer updateVaultState multicall tx failed'
+            ):
                 await update_vaults_state(vaults=[VAULT_1], block_number=BlockNumber(100))
 
     async def test_mix_of_meta_and_regular_vaults(self) -> None:
         """Meta vault is harvested via meta_vault_tree_update_state; regular vaults are
-        batched into a single multicall. Harvest params are fetched only for the
-        regular vaults."""
+        batched into a single multicall on the OsTokenRedeemer. Harvest params are
+        fetched only for the regular vaults."""
         params = make_harvest_params()
         meta_vaults_map = {VAULT_1: MagicMock()}
         with _mock_update_vaults_state(
@@ -574,9 +573,7 @@ async def test_mix_of_meta_and_regular_vaults(self) -> None:
             mocks['update_state'].await_args.kwargs['root_meta_vault'] is meta_vaults_map[VAULT_1]
         )
         mocks['harvest_params'].assert_awaited_once_with([VAULT_2], BlockNumber(100))
-        mocks['multicall'].tx_aggregate.assert_awaited_once_with(
-            [(VAULT_2, ENCODED_UPDATE_STATE_CALL)]
-        )
+        mocks['redeemer'].update_vaults_state.assert_awaited_once_with({VAULT_2: params})
 
 
 class TestRedeemMetaVaultSubVaults:
@@ -987,9 +984,6 @@ def _mock_submit_redeem_position(
         }
 
 
-ENCODED_UPDATE_STATE_CALL = HexStr('0xdeadbeef')
-
-
 @contextmanager
 def _mock_update_vaults_state(
     needs_update: bool = True,
@@ -1004,7 +998,8 @@ def _mock_update_vaults_state(
     in this map are treated as meta vaults by update_vaults_state. ``harvest_params``
     is the dict returned by get_multiple_harvest_params; a None value for a vault
     skips it from the multicall (production behavior). ``update_state_exception``
-    makes meta_vault_tree_update_state raise.
+    makes meta_vault_tree_update_state raise. ``multicall_tx_status`` controls the
+    receipt status of the OsTokenRedeemer.update_vaults_state batched tx.
     """
     meta_vaults_map = {} if meta_vaults_map is None else meta_vaults_map
     harvest_params = {} if harvest_params is None else harvest_params
@@ -1014,11 +1009,10 @@ def _mock_update_vaults_state(
     else:
         update_state_mock = AsyncMock()
 
-    def vault_factory(addr: ChecksumAddress) -> MagicMock:
-        mock_vault = MagicMock()
-        mock_vault.contract_address = addr
-        mock_vault.get_update_state_call = MagicMock(return_value=ENCODED_UPDATE_STATE_CALL)
-        return mock_vault
+    mock_client = AsyncMock()
+    mock_client.eth.wait_for_transaction_receipt = AsyncMock(
+        return_value={'status': multicall_tx_status}
+    )
 
     with (
         patch(
@@ -1034,34 +1028,20 @@ def vault_factory(addr: ChecksumAddress) -> MagicMock:
             f'{MODULE}.get_multiple_harvest_params',
             new=AsyncMock(return_value=harvest_params),
         ) as mock_harvest_params,
-        patch(f'{MODULE}.VaultContract', side_effect=vault_factory) as mock_vault_cls,
-        _mock_multicall_tx(tx_status=multicall_tx_status) as multicall_mocks,
+        patch(f'{MODULE}.os_token_redeemer_contract') as mock_redeemer,
+        patch(f'{MODULE}.execution_client', new=mock_client),
+        patch(f'{MODULE}.settings') as mock_settings,
     ):
+        mock_settings.execution_transaction_timeout = 120
+        mock_redeemer.update_vaults_state = AsyncMock(return_value='0x' + '11' * 32)
         yield {
             'graph_get_vaults': mock_graph,
             'update_state': update_state_mock,
             'harvest_params': mock_harvest_params,
-            'multicall': multicall_mocks['mock_multicall'],
-            'vault_cls': mock_vault_cls,
+            'redeemer': mock_redeemer,
         }
 
 
-@contextmanager
-def _mock_multicall_tx(tx_status: int = 1) -> Iterator[dict[str, MagicMock]]:
-    """Mock multicall_contract.tx_aggregate with a fake tx hash + receipt status."""
-    mock_client = AsyncMock()
-    mock_client.eth.wait_for_transaction_receipt = AsyncMock(return_value={'status': tx_status})
-
-    with (
-        patch(f'{MODULE}.multicall_contract') as mock_multicall,
-        patch(f'{MODULE}.execution_client', new=mock_client),
-        patch(f'{MODULE}.settings') as mock_settings,
-    ):
-        mock_settings.execution_transaction_timeout = 120
-        mock_multicall.tx_aggregate = AsyncMock(return_value='0x' + '11' * 32)
-        yield {'mock_multicall': mock_multicall}
-
-
 @contextmanager
 def _mock_process(
     positions: list[OsTokenPosition] | None = None,

src/common/contracts.py

@@ -619,6 +619,34 @@ async def redeem_sub_vaults_assets(
             )
         return Web3.to_hex(tx_hash), tx_receipt['blockNumber']
 
+    async def update_vaults_state(
+        self, vault_to_harvest_params: dict[ChecksumAddress, HarvestParams]
+    ) -> HexStr:
+        """Batch ``updateVaultState`` calls into a single multicall on this contract."""
+        calls = [
+            self._encode_update_vault_state(vault, harvest_params)
+            for vault, harvest_params in vault_to_harvest_params.items()
+        ]
+        tx_function = self.contract.functions.multicall(calls)
+        tx_hash = await transaction_gas_wrapper(tx_function)
+        return Web3.to_hex(tx_hash)
+
+    def _encode_update_vault_state(
+        self, vault: ChecksumAddress, harvest_params: HarvestParams
+    ) -> HexStr:
+        return self.encode_abi(
+            'updateVaultState',
+            [
+                vault,
+                (
+                    harvest_params.rewards_root,
+                    harvest_params.reward,
+                    harvest_params.unlocked_mev_reward,
+                    harvest_params.proof,
+                ),
+            ],
+        )
+
 
 class ValidatorsCheckerContract(ContractWrapper):
     abi_path = 'abi/IValidatorsChecker.json'