diff --git a/packages/testing/src/consensus_testing/forks/__init__.py b/packages/testing/src/consensus_testing/forks/__init__.py index e49ba2db..9bfa6161 100644 --- a/packages/testing/src/consensus_testing/forks/__init__.py +++ b/packages/testing/src/consensus_testing/forks/__init__.py @@ -2,27 +2,19 @@ from typing import Type -from framework.forks import BaseFork, BaseForkMeta +from framework.forks import BaseFork, BaseForkMeta, ForkRegistry +from . import forks as _forks_module from .forks import Devnet -from .helpers import ( - ALL_FORKS, - get_fork_by_name, - get_forks, - get_forks_with_no_parents, - get_from_until_fork_set, -) Fork = Type[BaseFork] +registry = ForkRegistry(_forks_module) + __all__ = [ - "ALL_FORKS", "BaseFork", "BaseForkMeta", "Devnet", "Fork", - "get_fork_by_name", - "get_forks", - "get_forks_with_no_parents", - "get_from_until_fork_set", + "registry", ] diff --git a/packages/testing/src/consensus_testing/forks/helpers.py b/packages/testing/src/consensus_testing/forks/helpers.py deleted file mode 100644 index dcf278f8..00000000 --- a/packages/testing/src/consensus_testing/forks/helpers.py +++ /dev/null @@ -1,55 +0,0 @@ -"""Consensus layer fork discovery and helpers.""" - -from typing import FrozenSet, Set, Type - -from framework.forks import BaseFork -from framework.forks.helpers import ( - get_all_forks, - get_forks_with_no_parents, - get_from_until_fork_set, -) -from framework.forks.helpers import ( - get_fork_by_name as _get_fork_by_name, -) -from framework.forks.helpers import ( - get_forks as _get_forks, -) - -from . import forks - -# Discover all consensus forks at module import time -ALL_FORKS: FrozenSet[Type[BaseFork]] = get_all_forks(forks) -"""All available consensus forks, excluding ignored forks.""" - - -def get_forks() -> Set[Type[BaseFork]]: - """ - Return the set of all available consensus forks. - - Returns: - Set of all non-ignored consensus fork classes. - """ - return _get_forks(ALL_FORKS) - - -def get_fork_by_name(fork_name: str) -> Type[BaseFork] | None: - """ - Get a consensus fork class by its name. - - Args: - fork_name: Name of the fork (case-insensitive). - - Returns: - The fork class, or None if not found. - """ - return _get_fork_by_name(ALL_FORKS, fork_name) - - -# Re-export the generic helpers for convenience -__all__ = [ - "ALL_FORKS", - "get_forks", - "get_fork_by_name", - "get_forks_with_no_parents", - "get_from_until_fork_set", -] diff --git a/packages/testing/src/framework/forks/__init__.py b/packages/testing/src/framework/forks/__init__.py index 4f725630..dbc78c6e 100644 --- a/packages/testing/src/framework/forks/__init__.py +++ b/packages/testing/src/framework/forks/__init__.py @@ -1,22 +1,10 @@ """Base fork infrastructure for Ethereum testing.""" from framework.forks.base import BaseFork, BaseForkMeta -from framework.forks.helpers import ( - discover_forks, - get_all_forks, - get_fork_by_name, - get_forks, - get_forks_with_no_parents, - get_from_until_fork_set, -) +from framework.forks.registry import ForkRegistry __all__ = [ "BaseFork", "BaseForkMeta", - "discover_forks", - "get_all_forks", - "get_fork_by_name", - "get_forks", - "get_forks_with_no_parents", - "get_from_until_fork_set", + "ForkRegistry", ] diff --git a/packages/testing/src/framework/forks/helpers.py b/packages/testing/src/framework/forks/helpers.py deleted file mode 100644 index e0aa0047..00000000 --- a/packages/testing/src/framework/forks/helpers.py +++ /dev/null @@ -1,117 +0,0 @@ -"""Generic fork helper functions for any Ethereum layer.""" - -from types import ModuleType -from typing import FrozenSet, List, Set, Type - -from framework.forks import BaseFork - - -def discover_forks(forks_module: ModuleType) -> List[Type[BaseFork]]: - """ - Discover all fork classes by scanning a forks module. - - Args: - forks_module: The module containing fork definitions (e.g., consensus_testing.forks.forks). - - Returns: - List of all BaseFork subclasses found in the module. - """ - discovered: List[Type[BaseFork]] = [] - for name in dir(forks_module): - obj = getattr(forks_module, name) - # Check if it's a type (class) and subclass of BaseFork (but not BaseFork itself) - if isinstance(obj, type) and issubclass(obj, BaseFork) and obj is not BaseFork: - discovered.append(obj) - return discovered - - -def get_all_forks(forks_module: ModuleType) -> FrozenSet[Type[BaseFork]]: - """ - Get all available forks from a forks module, excluding ignored forks. - - Args: - forks_module: The module containing fork definitions. - - Returns: - Frozen set of all non-ignored fork classes. - """ - all_forks = discover_forks(forks_module) - return frozenset(fork for fork in all_forks if not fork.ignore()) - - -def get_forks(all_forks: FrozenSet[Type[BaseFork]]) -> Set[Type[BaseFork]]: - """ - Convert a frozen set of forks to a regular set. - - Args: - all_forks: Frozen set of fork classes. - - Returns: - Set of fork classes. - """ - return set(all_forks) - - -def get_fork_by_name(all_forks: FrozenSet[Type[BaseFork]], fork_name: str) -> Type[BaseFork] | None: - """ - Get a fork class by its name. - - Args: - all_forks: Set of available forks to search. - fork_name: Name of the fork (case-insensitive). - - Returns: - The fork class, or None if not found. - """ - for fork in all_forks: - if fork.name().lower() == fork_name.lower(): - return fork - return None - - -def get_forks_with_no_parents(forks: Set[Type[BaseFork]]) -> Set[Type[BaseFork]]: - """ - Get all forks that have no parent forks in the given set. - - Args: - forks: Set of forks to search. - - Returns: - Set of forks with no parents (root forks). - """ - result: Set[Type[BaseFork]] = set() - for fork in forks: - has_parent = False - for other_fork in forks - {fork}: - if other_fork < fork: # other_fork is older than fork - has_parent = True - break - if not has_parent: - result.add(fork) - return result - - -def get_from_until_fork_set( - forks: Set[Type[BaseFork]], - forks_from: Set[Type[BaseFork]], - forks_until: Set[Type[BaseFork]], -) -> Set[Type[BaseFork]]: - """ - Get all forks in the range from forks_from to forks_until (inclusive). - - Args: - forks: The complete set of forks to filter. - forks_from: Start of the range (inclusive). - forks_until: End of the range (inclusive). - - Returns: - Set of forks in the specified range. - """ - result: Set[Type[BaseFork]] = set() - for fork_from in forks_from: - for fork_until in forks_until: - for fork in forks: - # Fork must be >= fork_from and <= fork_until - if fork >= fork_from and fork <= fork_until: - result.add(fork) - return result diff --git a/packages/testing/src/framework/forks/registry.py b/packages/testing/src/framework/forks/registry.py new file mode 100644 index 00000000..7009118b --- /dev/null +++ b/packages/testing/src/framework/forks/registry.py @@ -0,0 +1,32 @@ +"""Fork registry for discovering and looking up forks by name.""" + +from types import ModuleType +from typing import Type + +from framework.forks.base import BaseFork + + +class ForkRegistry: + """Discovers forks from a module and provides O(1) name-based lookup.""" + + def __init__(self, forks_module: ModuleType) -> None: + """Scan a forks module and build the name index.""" + discovered: set[type[BaseFork]] = set() + for name in dir(forks_module): + obj = getattr(forks_module, name) + if isinstance(obj, type) and issubclass(obj, BaseFork) and obj is not BaseFork: + discovered.add(obj) + + self._forks = frozenset(fork for fork in discovered if not fork.ignore()) + self._by_name: dict[str, type[BaseFork]] = { + fork.name().lower(): fork for fork in self._forks + } + + @property + def forks(self) -> frozenset[type[BaseFork]]: + """All available non-ignored forks.""" + return self._forks + + def get_fork_by_name(self, fork_name: str) -> Type[BaseFork] | None: + """Case-insensitive fork lookup.""" + return self._by_name.get(fork_name.lower()) diff --git a/packages/testing/src/framework/pytest_plugins/filler.py b/packages/testing/src/framework/pytest_plugins/filler.py index b4c4aff2..b782df0e 100644 --- a/packages/testing/src/framework/pytest_plugins/filler.py +++ b/packages/testing/src/framework/pytest_plugins/filler.py @@ -233,11 +233,8 @@ def pytest_configure(config: pytest.Config) -> None: clean = config.getoption("--clean") # Get available forks from layer-specific module - get_forks = layer_module.forks.get_forks - get_fork_by_name = layer_module.forks.get_fork_by_name - - available_forks = get_forks() - available_fork_names = sorted(fork.name() for fork in available_forks) + registry = layer_module.forks.registry + available_fork_names = sorted(fork.name() for fork in registry.forks) # Validate fork if not fork_name: @@ -248,7 +245,7 @@ def pytest_configure(config: pytest.Config) -> None: ) pytest.exit("Missing required --fork option.", returncode=pytest.ExitCode.USAGE_ERROR) - fork_class = get_fork_by_name(fork_name) + fork_class = registry.get_fork_by_name(fork_name) if fork_class is None: print( f"Error: Unsupported fork for {layer} layer: {fork_name}\n", @@ -289,13 +286,13 @@ def pytest_collection_modifyitems(config: pytest.Config, items: List[pytest.Item fork_class = config.test_fork_class layer_module = config.layer_module # type: ignore[attr-defined] - get_fork_by_name = layer_module.forks.get_fork_by_name + registry = layer_module.forks.registry verbose = config.getoption("verbose") deselected = [] selected = [] for item in items: - if not _is_test_item_valid_for_fork(item, fork_class, get_fork_by_name): + if not _is_test_item_valid_for_fork(item, fork_class, registry.get_fork_by_name): if verbose < 2: deselected.append(item) else: @@ -500,9 +497,9 @@ def pytest_generate_tests(metafunc: pytest.Metafunc) -> None: fork_class = metafunc.config.test_fork_class # type: ignore[attr-defined] layer_module = metafunc.config.layer_module # type: ignore[attr-defined] - get_fork_by_name = layer_module.forks.get_fork_by_name + registry = layer_module.forks.registry - if not _is_test_valid_for_fork(metafunc, fork_class, get_fork_by_name): + if not _is_test_valid_for_fork(metafunc, fork_class, registry.get_fork_by_name): verbose = metafunc.config.getoption("verbose") if verbose >= 2: metafunc.parametrize(