Refactor do_fork with terrifying eval() and PREVIOUS_FORK_OF

Author: hwwhww

tests/core/pyspec/eth2spec/test/helpers/fork_transition.py

@@ -11,12 +11,8 @@
 )
 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(