refactor(plugins/execute): Refactor account refunds logic (#1204)

* refactor(execute): Move refund logic to `pre` fixture

* fix(plugins/execute): Overwrite tx instead of bumping nonce

* fix(rpc): handle null transaction response

* fix(execute): Return `--sender-funding-txs-gas-price`

Author: marioevz

src/cli/gentest/request_manager.py

@@ -36,6 +36,7 @@ def __init__(self):
     def eth_get_transaction_by_hash(self, transaction_hash: Hash) -> TransactionByHashResponse:
         """Get transaction data."""
         res = self.rpc.get_transaction_by_hash(transaction_hash)
+        assert res is not None, "Transaction not found"
         block_number = res.block_number
         assert block_number is not None, "Transaction does not seem to be included in any block"
 

src/ethereum_test_rpc/rpc.py

@@ -136,13 +136,13 @@ def get_transaction_count(
         block = hex(block_number) if isinstance(block_number, int) else block_number
         return int(self.post_request("getTransactionCount", f"{address}", block), 16)
 
-    def get_transaction_by_hash(self, transaction_hash: Hash) -> TransactionByHashResponse:
+    def get_transaction_by_hash(self, transaction_hash: Hash) -> TransactionByHashResponse | None:
         """`eth_getTransactionByHash`: Returns transaction details."""
         try:
-            resp = TransactionByHashResponse(
-                **self.post_request("getTransactionByHash", f"{transaction_hash}")
-            )
-            return resp
+            response = self.post_request("getTransactionByHash", f"{transaction_hash}")
+            if response is None:
+                return None
+            return TransactionByHashResponse(**response)
         except ValidationError as e:
             pprint(e.errors())
             raise e
@@ -200,7 +200,7 @@ def wait_for_transaction(self, transaction: Transaction) -> TransactionByHashRes
         start_time = time.time()
         while True:
             tx = self.get_transaction_by_hash(tx_hash)
-            if tx.block_number is not None:
+            if tx is not None and tx.block_number is not None:
                 return tx
             if (time.time() - start_time) > self.transaction_wait_timeout:
                 break
@@ -225,7 +225,7 @@ def wait_for_transactions(
             while i < len(tx_hashes):
                 tx_hash = tx_hashes[i]
                 tx = self.get_transaction_by_hash(tx_hash)
-                if tx.block_number is not None:
+                if tx is not None and tx.block_number is not None:
                     responses.append(tx)
                     tx_hashes.pop(i)
                 else:

src/pytest_plugins/execute/execute.py

@@ -7,11 +7,10 @@
 import pytest
 from pytest_metadata.plugin import metadata_key  # type: ignore
 
-from ethereum_test_base_types import Number
 from ethereum_test_execution import BaseExecute
 from ethereum_test_forks import Fork
 from ethereum_test_rpc import EthRPC
-from ethereum_test_tools import SPEC_TYPES, BaseTest, TestInfo, Transaction
+from ethereum_test_tools import SPEC_TYPES, BaseTest, TestInfo
 from ethereum_test_types import TransactionDefaults
 from pytest_plugins.spec_version_checker.spec_version_checker import EIPSpecTestItem
 
@@ -173,7 +172,9 @@ def pytest_html_report_title(report):
 @pytest.fixture(scope="session")
 def default_gas_price(request) -> int:
     """Return default gas price used for transactions."""
-    return request.config.getoption("default_gas_price")
+    gas_price = request.config.getoption("default_gas_price")
+    assert gas_price > 0, "Gas price must be greater than 0"
+    return gas_price
 
 
 @pytest.fixture(scope="session")
@@ -252,7 +253,6 @@ def base_test_parametrizer_func(
         eips: List[int],
         eth_rpc: EthRPC,
         collector: Collector,
-        default_gas_price: int,
     ):
         """
         Fixture used to instantiate an auto-fillable BaseTest object from within
@@ -281,13 +281,14 @@ def __init__(self, *args, **kwargs):
 
                 # wait for pre-requisite transactions to be included in blocks
                 pre.wait_for_transactions()
-                for deployed_contract, deployed_code in pre._deployed_contracts:
-                    if eth_rpc.get_code(deployed_contract) == deployed_code:
-                        pass
-                    else:
+                for deployed_contract, expected_code in pre._deployed_contracts:
+                    actual_code = eth_rpc.get_code(deployed_contract)
+                    if actual_code != expected_code:
                         raise Exception(
                             f"Deployed test contract didn't match expected code at address "
-                            f"{deployed_contract} (not enough gas_limit?)."
+                            f"{deployed_contract} (not enough gas_limit?).\n"
+                            f"Expected: {expected_code}\n"
+                            f"Actual: {actual_code}"
                         )
                 request.node.config.funded_accounts = ", ".join(
                     [str(eoa) for eoa in pre._funded_eoa]
@@ -297,33 +298,7 @@ def __init__(self, *args, **kwargs):
                 execute.execute(eth_rpc)
                 collector.collect(request.node.nodeid, execute)
 
-        sender_start_balance = eth_rpc.get_balance(pre._sender)
-
-        yield BaseTestWrapper
-
-        # Refund all EOAs (regardless of whether the test passed or failed)
-        refund_txs = []
-        for eoa in pre._funded_eoa:
-            remaining_balance = eth_rpc.get_balance(eoa)
-            eoa.nonce = Number(eth_rpc.get_transaction_count(eoa))
-            refund_gas_limit = 21_000
-            tx_cost = refund_gas_limit * default_gas_price
-            if remaining_balance < tx_cost:
-                continue
-            refund_txs.append(
-                Transaction(
-                    sender=eoa,
-                    to=pre._sender,
-                    gas_limit=21_000,
-                    gas_price=default_gas_price,
-                    value=remaining_balance - tx_cost,
-                ).with_signature_and_sender()
-            )
-        eth_rpc.send_wait_transactions(refund_txs)
-
-        sender_end_balance = eth_rpc.get_balance(pre._sender)
-        used_balance = sender_start_balance - sender_end_balance
-        print(f"Used balance={used_balance / 10**18:.18f}")
+        return BaseTestWrapper
 
     return base_test_parametrizer_func
 

src/pytest_plugins/execute/pre_alloc.py

@@ -2,12 +2,12 @@
 
 from itertools import count
 from random import randint
-from typing import Iterator, List, Literal, Tuple
+from typing import Generator, Iterator, List, Literal, Tuple
 
 import pytest
 from pydantic import PrivateAttr
 
-from ethereum_test_base_types import Number, StorageRootType, ZeroPaddedHexNumber
+from ethereum_test_base_types import Bytes, Number, StorageRootType, ZeroPaddedHexNumber
 from ethereum_test_base_types.conversions import (
     BytesConvertible,
     FixedSizeBytesConvertible,
@@ -94,7 +94,7 @@ class Alloc(BaseAlloc):
     _sender: EOA = PrivateAttr()
     _eth_rpc: EthRPC = PrivateAttr()
     _txs: List[Transaction] = PrivateAttr(default_factory=list)
-    _deployed_contracts: List[Tuple[Address, bytes]] = PrivateAttr(default_factory=list)
+    _deployed_contracts: List[Tuple[Address, Bytes]] = PrivateAttr(default_factory=list)
     _funded_eoa: List[EOA] = PrivateAttr(default_factory=list)
     _evm_code_type: EVMCodeType | None = PrivateAttr(None)
     _chain_id: int = PrivateAttr()
@@ -206,7 +206,7 @@ def deploy_contract(
         self._txs.append(deploy_tx)
 
         contract_address = deploy_tx.created_contract
-        self._deployed_contracts.append((contract_address, bytes(code)))
+        self._deployed_contracts.append((contract_address, Bytes(code)))
 
         assert Number(nonce) >= 1, "impossible to deploy contract with nonce lower than one"
 
@@ -373,9 +373,14 @@ def pre(
     evm_code_type: EVMCodeType,
     chain_id: int,
     eoa_fund_amount_default: int,
-) -> Alloc:
+    default_gas_price: int,
+) -> Generator[Alloc, None, None]:
     """Return default pre allocation for all tests (Empty alloc)."""
-    return Alloc(
+    # Record the starting balance of the sender
+    sender_test_starting_balance = eth_rpc.get_balance(sender_key)
+
+    # Prepare the pre-alloc
+    pre = Alloc(
         fork=fork,
         sender=sender_key,
         eth_rpc=eth_rpc,
@@ -384,3 +389,31 @@ def pre(
         chain_id=chain_id,
         eoa_fund_amount_default=eoa_fund_amount_default,
     )
+
+    # Yield the pre-alloc for usage during the test
+    yield pre
+
+    # Refund all EOAs (regardless of whether the test passed or failed)
+    refund_txs = []
+    for eoa in pre._funded_eoa:
+        remaining_balance = eth_rpc.get_balance(eoa)
+        eoa.nonce = Number(eth_rpc.get_transaction_count(eoa))
+        refund_gas_limit = 21_000
+        tx_cost = refund_gas_limit * default_gas_price
+        if remaining_balance < tx_cost:
+            continue
+        refund_txs.append(
+            Transaction(
+                sender=eoa,
+                to=sender_key,
+                gas_limit=21_000,
+                gas_price=default_gas_price,
+                value=remaining_balance - tx_cost,
+            ).with_signature_and_sender()
+        )
+    eth_rpc.send_wait_transactions(refund_txs)
+
+    # Record the ending balance of the sender
+    sender_test_ending_balance = eth_rpc.get_balance(sender_key)
+    used_balance = sender_test_starting_balance - sender_test_ending_balance
+    print(f"Used balance={used_balance / 10**18:.18f}")

src/pytest_plugins/execute/rpc/hive.py

@@ -721,6 +721,7 @@ def wait_for_transactions(
             while tx_id < len(tx_hashes):
                 tx_hash = tx_hashes[tx_id]
                 tx = self.get_transaction_by_hash(tx_hash)
+                assert tx is not None, f"Transaction {tx_hash} not found"
                 if tx.block_number is not None:
                     responses.append(tx)
                     tx_hashes.pop(tx_id)

src/pytest_plugins/execute/sender.py

@@ -34,7 +34,7 @@ def pytest_addoption(parser):
         action="store",
         dest="sender_funding_transactions_gas_price",
         type=Wei,
-        default=10**9,
+        default=None,
         help=("Gas price set for the funding transactions of each worker's sender key."),
     )
 
@@ -49,9 +49,15 @@ def pytest_addoption(parser):
 
 
 @pytest.fixture(scope="session")
-def sender_funding_transactions_gas_price(request: pytest.FixtureRequest) -> int:
+def sender_funding_transactions_gas_price(
+    request: pytest.FixtureRequest, default_gas_price: int
+) -> int:
     """Get the gas price for the funding transactions."""
-    return request.config.option.sender_funding_transactions_gas_price
+    gas_price: int | None = request.config.option.sender_funding_transactions_gas_price
+    if gas_price is None:
+        gas_price = default_gas_price
+    assert gas_price > 0, "Gas price must be greater than 0"
+    return gas_price
 
 
 @pytest.fixture(scope="session")
@@ -159,13 +165,16 @@ def sender_key(
 
     # refund seed sender
     remaining_balance = eth_rpc.get_balance(sender)
+    sender.nonce = Number(eth_rpc.get_transaction_count(sender))
     used_balance = sender_key_initial_balance - remaining_balance
     request.config.stash[metadata_key]["Senders"][str(sender)] = (
         f"Used balance={used_balance / 10**18:.18f}"
     )
 
     refund_gas_limit = sender_fund_refund_gas_limit
-    refund_gas_price = sender_funding_transactions_gas_price
+    # double the gas price to ensure the transaction is included and overwrites any other
+    # transaction that might have been sent by the sender.
+    refund_gas_price = sender_funding_transactions_gas_price * 2
     tx_cost = refund_gas_limit * refund_gas_price
 
     if (remaining_balance - 1) < tx_cost: