From a8de54483c34e1e7f3beb8809d4ecc306cefaa12 Mon Sep 17 00:00:00 2001 From: danceratopz Date: Mon, 19 Aug 2024 14:50:05 +0200 Subject: [PATCH] feat: add a minimal poc to enable custom fork definitions WIP This adds a fork registry to allow dynamic fork class definitions. This enables non-mainnet evm implementations to add/modify their fork definitions. --- pytest.ini | 1 + src/ethereum_test_forks/__init__.py | 6 +++ src/ethereum_test_forks/fork_registry.py | 33 +++++++++++++++ .../chain_selector/chain_selector.py | 41 +++++++++++++++++++ src/pytest_plugins/forks/forks.py | 7 +++- src/pytest_plugins/optimism.py | 30 ++++++++++++++ 6 files changed, 116 insertions(+), 2 deletions(-) create mode 100644 src/ethereum_test_forks/fork_registry.py create mode 100644 src/pytest_plugins/chain_selector/chain_selector.py create mode 100644 src/pytest_plugins/optimism.py diff --git a/pytest.ini b/pytest.ini index 309f013e2f..4e81f8fc9b 100644 --- a/pytest.ini +++ b/pytest.ini @@ -10,6 +10,7 @@ addopts = -p pytest_plugins.filler.pre_alloc -p pytest_plugins.filler.filler -p pytest_plugins.forks.forks + -p pytest_plugins.chain_selector.chain_selector -p pytest_plugins.spec_version_checker.spec_version_checker -p pytest_plugins.help.help -m "not eip_version_check" diff --git a/src/ethereum_test_forks/__init__.py b/src/ethereum_test_forks/__init__.py index 456679967b..c0dc7c8be6 100644 --- a/src/ethereum_test_forks/__init__.py +++ b/src/ethereum_test_forks/__init__.py @@ -20,11 +20,15 @@ Prague, Shanghai, ) + from .forks.transition import ( BerlinToLondonAt5, ParisToShanghaiAtTime15k, ShanghaiToCancunAtTime15k, ) + +from .fork_registry import update_fork_registry, get_fork_registry + from .helpers import ( InvalidForkError, forks_from, @@ -67,10 +71,12 @@ "forks_from_until", "get_deployed_forks", "get_development_forks", + "get_fork_registry", "get_forks", "get_forks_with_solc_support", "get_forks_without_solc_support", "get_closest_fork_with_solc_support", "transition_fork_from_to", "transition_fork_to", + "update_fork_registry", ] diff --git a/src/ethereum_test_forks/fork_registry.py b/src/ethereum_test_forks/fork_registry.py new file mode 100644 index 0000000000..33c974e735 --- /dev/null +++ b/src/ethereum_test_forks/fork_registry.py @@ -0,0 +1,33 @@ +""" +Define a global registry for forks that can be used in the Ethereum test +framework. + +This approach allows different chains to register their fork class +definitions. +""" + +from typing import Dict, Type +from .base_fork import BaseFork +from .helpers import get_forks + +fork_registry: Dict[str, Dict[str, Type[BaseFork]]] = {} + +fork_registry["ethereum-mainnet"] = {fork.name(): fork for fork in get_forks()} + + +def update_fork_registry(chain_name: str, forks: Dict[str, Type[BaseFork]]): + """ + Updates the global fork registry with forks from a specific chain. + If the chain already exists, it merges the new forks with the existing ones. + """ + if chain_name in fork_registry: + fork_registry[chain_name].update(forks) + else: + fork_registry[chain_name] = forks + + +def get_fork_registry() -> Dict[str, Dict[str, Type[BaseFork]]]: + """ + Returns the current fork registry. + """ + return fork_registry diff --git a/src/pytest_plugins/chain_selector/chain_selector.py b/src/pytest_plugins/chain_selector/chain_selector.py new file mode 100644 index 0000000000..ede0a9620f --- /dev/null +++ b/src/pytest_plugins/chain_selector/chain_selector.py @@ -0,0 +1,41 @@ +""" + +""" + +import pytest +from ethereum_test_forks.fork_registry import get_fork_registry +from ethereum_test_forks.helpers import get_forks + + +def pytest_addoption(parser): + parser.addoption( + "--chain", + action="store", + default="ethereum-mainnet", + help="Select the EVM chain for testing. Default: ethereum-mainnet.", + ) + + +@pytest.fixture(scope="session") +def chain(request): + return request.config.getoption("--chain") + + +@pytest.fixture(scope="session") +def forks(chain): + fork_registry = get_fork_registry() + if chain not in fork_registry: + raise ValueError(f"No forks found for chain: {chain}") + return fork_registry[chain].values() + + +# @pytest.hookimpl(tryfirst=True) +# def pytest_generate_tests(metafunc): +# if "fork" in metafunc.fixturenames: +# chain = metafunc.config.getoption("chain") +# fork_registry = get_fork_registry() +# if chain in fork_registry: +# forks = fork_registry[chain] +# metafunc.parametrize( +# "fork", forks.values(), ids=[fork.__name__ for fork in forks.values()] +# ) diff --git a/src/pytest_plugins/forks/forks.py b/src/pytest_plugins/forks/forks.py index e6082e26b9..ce143e41fc 100644 --- a/src/pytest_plugins/forks/forks.py +++ b/src/pytest_plugins/forks/forks.py @@ -16,6 +16,7 @@ ForkAttribute, get_deployed_forks, get_forks, + get_fork_registry, get_transition_forks, transition_fork_to, ) @@ -274,7 +275,9 @@ def pytest_configure(config: pytest.Config): for d in fork_covariant_descriptors: config.addinivalue_line("markers", f"{d.marker_name}: {d.description}") - forks = set([fork for fork in get_forks() if not fork.ignore()]) + chain = config.getoption("chain") + assert chain in get_fork_registry(), f"No forks found for chain: {chain}" + forks = set([fork for fork in get_fork_registry()[chain].values() if not fork.ignore()]) config.forks = forks # type: ignore config.fork_names = set([fork.name() for fork in sorted(list(forks))]) # type: ignore config.forks_by_name = {fork.name(): fork for fork in forks} # type: ignore @@ -305,7 +308,7 @@ def get_fork_option(config, option_name: str, parameter_name: str) -> Set[Fork]: resulting_forks = set() - for fork in get_forks(): + for fork in get_fork_registry()[chain].values(): if fork.name() in forks_str: resulting_forks.add(fork) diff --git a/src/pytest_plugins/optimism.py b/src/pytest_plugins/optimism.py new file mode 100644 index 0000000000..d6f960a1ff --- /dev/null +++ b/src/pytest_plugins/optimism.py @@ -0,0 +1,30 @@ +""" +Chain definition for the Optimism network. +""" +from typing import List + +import pytest + +from ethereum_test_forks import Cancun, get_fork_registry, update_fork_registry + + +class OptimismCancun(Cancun): + """ + Defines the Optimism Cancun fork, which adds the secp256r1 precompile. + """ + + @classmethod + def precompiles(cls, block_number: int = 0, timestamp: int = 0) -> List[int]: + """ + Optimism added the secp256r1 precompile to its Cancun fork. + """ + return [0x100] + super(OptimismCancun, cls).precompiles(block_number, timestamp) + + +@pytest.hookimpl(tryfirst=True) +def pytest_configure(config): + ethereum_mainnet_forks = get_fork_registry()["ethereum-mainnet"] + assert "Cancun" in ethereum_mainnet_forks, "Cancun fork not found in ethereum-mainnet forks." + optimism_mainnet_forks = ethereum_mainnet_forks.copy() + optimism_mainnet_forks["Cancun"] = OptimismCancun + update_fork_registry("optimism", optimism_mainnet_forks)