feat: WithdrawalRequests web3py extension

src/main.py

@@ -24,7 +24,8 @@
     KeysAPIClientModule,
     LidoValidatorsProvider,
     FallbackProviderModule,
-    LazyCSM
+    LazyCSM,
+    WithdrawalRequests
 )
 from src.web3py.middleware import metrics_collector
 from src.web3py.types import Web3
@@ -101,6 +102,7 @@ def main(module_name: OracleModule):
         'cc': lambda: cc,  # type: ignore[dict-item]
         'kac': lambda: kac,  # type: ignore[dict-item]
         'ipfs': lambda: ipfs,  # type: ignore[dict-item]
+        'withdrawal_requests': WithdrawalRequests,
     })
 
     logger.info({'msg': 'Add metrics middleware for ETH1 requests.'})

src/web3py/extensions/__init__.py

@@ -5,3 +5,4 @@
 from src.web3py.extensions.lido_validators import LidoValidatorsProvider
 from src.web3py.extensions.fallback import FallbackProviderModule
 from src.web3py.extensions.csm import CSM, LazyCSM
+from src.web3py.extensions.withdrawal_requests import WithdrawalRequests

src/web3py/extensions/withdrawal_requests.py

@@ -0,0 +1,30 @@
+import logging
+
+from web3 import Web3
+from web3.module import Module
+
+from src.providers.execution.exceptions import InconsistentData
+from src.types import BlockStamp
+
+logger = logging.getLogger(__name__)
+
+
+class WithdrawalRequests(Module):
+    """
+    Web3py extension to work with EIP-7002 withdrawal requests.
+    See https://eips.ethereum.org/EIPS/eip-7002 for details.
+    """
+
+    w3: Web3
+
+    ADDRESS = Web3.to_checksum_address("0x0c15F14308530b7CDB8460094BbB9cC28b9AaaAA")
+
+    QUEUE_HEAD_SLOT = 2
+    QUEUE_TAIL_SLOT = 3
+
+    def get_queue_len(self, blockstamp: BlockStamp):
+        head = self.w3.eth.get_storage_at(self.ADDRESS, self.QUEUE_HEAD_SLOT, block_identifier=blockstamp.block_hash)
+        tail = self.w3.eth.get_storage_at(self.ADDRESS, self.QUEUE_TAIL_SLOT, block_identifier=blockstamp.block_hash)
+        if head > tail:
+            raise InconsistentData("EIP-7002 queue's head is over the tail")
+        return int.from_bytes(tail) - int.from_bytes(head)

src/web3py/types.py

@@ -1,14 +1,14 @@
 from web3 import Web3 as _Web3
 
-
 from src.providers.ipfs import IPFSProvider
 from src.web3py.extensions import (
-    LidoContracts,
-    TransactionUtils,
+    CSM,
     ConsensusClientModule,
     KeysAPIClientModule,
+    LidoContracts,
     LidoValidatorsProvider,
-    CSM
+    TransactionUtils,
+    WithdrawalRequests,
 )
 
 
@@ -20,3 +20,4 @@ class Web3(_Web3):
     kac: KeysAPIClientModule
     csm: CSM
     ipfs: IPFSProvider
+    withdrawal_requests: WithdrawalRequests

tests/conftest.py

@@ -14,7 +14,7 @@
 from src.types import BlockNumber, EpochNumber, ReferenceBlockStamp, SlotNumber
 from src.variables import CONSENSUS_CLIENT_URI, EXECUTION_CLIENT_URI, KEYS_API_URI
 from src.web3py.contract_tweak import tweak_w3_contracts
-from src.web3py.extensions import LidoContracts, LidoValidatorsProvider, TransactionUtils
+from src.web3py.extensions import LidoContracts, LidoValidatorsProvider, TransactionUtils, WithdrawalRequests
 from src.web3py.types import Web3
 from tests.providers_utils import (
     ResponseFromFile,
@@ -182,6 +182,11 @@ def get_blockstamp_by_state(w3, state_id) -> ReferenceBlockStamp:
     )
 
 
+@pytest.fixture()
+def withdrawal_requests(web3):
+    web3.attach_modules({"withdrawal_requests": WithdrawalRequests})
+
+
 @dataclass
 class Account:
     address: ChecksumAddress

tests/web3_extentions/test_withdrawal_requests.py

@@ -0,0 +1,55 @@
+from unittest.mock import Mock
+
+import pytest
+from hexbytes import HexBytes
+
+from src.providers.execution.exceptions import InconsistentData
+from src.web3py.types import Web3
+from tests.factory.blockstamp import ReferenceBlockStampFactory
+
+blockstamp = ReferenceBlockStampFactory.build()
+
+
+@pytest.mark.unit
+@pytest.mark.usefixtures("withdrawal_requests")
+def test_queue_len(web3: Web3):
+    web3.eth.get_storage_at = Mock(
+        side_effect=[
+            HexBytes(bytes.fromhex("")),
+            HexBytes(bytes.fromhex("")),
+        ]
+    )
+    assert web3.withdrawal_requests.get_queue_len(blockstamp) == 0
+
+    web3.eth.get_storage_at = Mock(
+        side_effect=[
+            HexBytes(bytes.fromhex("0000000000000000000000000000000000000000000000000000000000000000")),
+            HexBytes(bytes.fromhex("0000000000000000000000000000000000000000000000000000000000000000")),
+        ]
+    )
+    assert web3.withdrawal_requests.get_queue_len(blockstamp) == 0
+
+    web3.eth.get_storage_at = Mock(
+        side_effect=[
+            HexBytes(bytes.fromhex("0000000000000000000000000000000000000000000000000000000000000000")),
+            HexBytes(bytes.fromhex("00000000000000000000000000000000000000000000000000000000000001c7")),
+        ]
+    )
+    assert web3.withdrawal_requests.get_queue_len(blockstamp) == 455
+
+    web3.eth.get_storage_at = Mock(
+        side_effect=[
+            HexBytes(bytes.fromhex("0000000000000000000000000000000000000000000000000000000000000020")),
+            HexBytes(bytes.fromhex("00000000000000000000000000000000000000000000000000000000000001c7")),
+        ]
+    )
+    assert web3.withdrawal_requests.get_queue_len(blockstamp) == 423
+
+    web3.eth.get_storage_at = Mock(
+        side_effect=[
+            HexBytes(bytes.fromhex("00000000000000000000000000000000000000000000000000000000000001c7")),
+            HexBytes(bytes.fromhex("0000000000000000000000000000000000000000000000000000000000000020")),
+        ]
+    )
+    with pytest.raises(InconsistentData):
+        web3.withdrawal_requests.get_queue_len(blockstamp)