diff --git a/tests/core/pyspec/eth2spec/test/helpers/fork_transition.py b/tests/core/pyspec/eth2spec/test/helpers/fork_transition.py index 80c54d9c1..93879512d 100644 --- a/tests/core/pyspec/eth2spec/test/helpers/fork_transition.py +++ b/tests/core/pyspec/eth2spec/test/helpers/fork_transition.py @@ -11,12 +11,8 @@ from eth2spec.test.helpers.block import ( ) from eth2spec.test.helpers.bls_to_execution_changes import get_signed_address_change from eth2spec.test.helpers.constants import ( - ALTAIR, - BELLATRIX, - CAPELLA, - DENEB, - EIP6110, - EIP7002, + PHASE0, + PREVIOUS_FORK_OF, ) from eth2spec.test.helpers.deposits import ( prepare_state_and_deposit, @@ -146,45 +142,34 @@ def state_transition_across_slots_with_ignoring_proposers(spec, next_slot(spec, state) +def get_upgrade_fn(spec, fork): + try: + # TODO: make all upgrade_to_* function names consistent? + fn = eval(f"spec.upgrade_to_{fork}") + return fn + except Exception: + raise ValueError(f"Unknown fork: {fork}") + + def do_fork(state, spec, post_spec, fork_epoch, with_block=True, sync_aggregate=None, operation_dict=None): spec.process_slots(state, state.slot + 1) assert state.slot % spec.SLOTS_PER_EPOCH == 0 assert spec.get_current_epoch(state) == fork_epoch - if post_spec.fork == ALTAIR: - state = post_spec.upgrade_to_altair(state) - elif post_spec.fork == BELLATRIX: - state = post_spec.upgrade_to_bellatrix(state) - elif post_spec.fork == CAPELLA: - state = post_spec.upgrade_to_capella(state) - elif post_spec.fork == DENEB: - state = post_spec.upgrade_to_deneb(state) - elif post_spec.fork == EIP6110: - state = post_spec.upgrade_to_eip6110(state) - elif post_spec.fork == EIP7002: - state = post_spec.upgrade_to_eip7002(state) + state = get_upgrade_fn(post_spec, post_spec.fork)(state) assert state.fork.epoch == fork_epoch - if post_spec.fork == ALTAIR: - assert state.fork.previous_version == post_spec.config.GENESIS_FORK_VERSION - assert state.fork.current_version == post_spec.config.ALTAIR_FORK_VERSION - elif post_spec.fork == BELLATRIX: - assert state.fork.previous_version == post_spec.config.ALTAIR_FORK_VERSION - assert state.fork.current_version == post_spec.config.BELLATRIX_FORK_VERSION - elif post_spec.fork == CAPELLA: - assert state.fork.previous_version == post_spec.config.BELLATRIX_FORK_VERSION - assert state.fork.current_version == post_spec.config.CAPELLA_FORK_VERSION - elif post_spec.fork == DENEB: - assert state.fork.previous_version == post_spec.config.CAPELLA_FORK_VERSION - assert state.fork.current_version == post_spec.config.DENEB_FORK_VERSION - elif post_spec.fork == EIP6110: - assert state.fork.previous_version == post_spec.config.DENEB_FORK_VERSION - assert state.fork.current_version == post_spec.config.EIP6110_FORK_VERSION - elif post_spec.fork == EIP7002: - assert state.fork.previous_version == post_spec.config.CAPELLA_FORK_VERSION - assert state.fork.current_version == post_spec.config.EIP7002_FORK_VERSION + previous_fork = PREVIOUS_FORK_OF[post_spec.fork] + if previous_fork == PHASE0: + previous_version = spec.config.GENESIS_FORK_VERSION + else: + previous_version = getattr(post_spec.config, f"{previous_fork.upper()}_FORK_VERSION") + current_version = getattr(post_spec.config, f"{post_spec.fork.upper()}_FORK_VERSION") + + assert state.fork.previous_version == previous_version + assert state.fork.current_version == current_version if with_block: return state, _state_transition_and_sign_block_at_slot(