Apply with_config_overrides to all phases

tests/core/pyspec/eth2spec/test/altair/unittests/test_config_override.py

@@ -1,11 +1,15 @@
 from eth2spec.test.context import (
     spec_configured_state_test,
     spec_state_test_with_matching_config,
+    spec_test,
     with_all_phases,
+    with_config_overrides,
+    with_matching_spec_config,
     with_phases,
+    with_state,
 )
 from eth2spec.test.helpers.constants import (
-    PHASE0, ALTAIR,
+    PHASE0, ALTAIR, BELLATRIX,
     ALL_PHASES,
 )
 from eth2spec.test.helpers.forks import is_post_fork
@@ -30,7 +34,7 @@ def test_config_override(spec, state):
 
 @with_all_phases
 @spec_state_test_with_matching_config
-def test_override_config_fork_epoch(spec, state):
+def test_config_override_matching_fork_epochs(spec, state):
     # Fork schedule must be consistent with state fork
     epoch = spec.get_current_epoch(state)
     if is_post_fork(spec.fork, ALTAIR):
@@ -56,3 +60,27 @@ def test_override_config_fork_epoch(spec, state):
             continue
         fork_epoch_field = fork.upper() + '_FORK_EPOCH'
         assert getattr(spec.config, fork_epoch_field) <= epoch
+
+
+@with_phases(phases=[ALTAIR], other_phases=[BELLATRIX])
+@spec_test
+@with_config_overrides({
+    'ALTAIR_FORK_VERSION': '0x11111111',
+    'BELLATRIX_FORK_EPOCH': 4,
+}, emit=False)
+@with_state
+@with_matching_spec_config(emitted_fork=BELLATRIX)
+def test_config_override_across_phases(spec, phases, state):
+    assert state.fork.current_version == spec.config.ALTAIR_FORK_VERSION
+
+    assert spec.config.ALTAIR_FORK_VERSION == spec.Version('0x11111111')
+    assert spec.config.ALTAIR_FORK_EPOCH == 0
+    assert not hasattr(spec.config, 'BELLATRIX_FORK_EPOCH')
+
+    assert phases[ALTAIR].config.ALTAIR_FORK_VERSION == spec.Version('0x11111111')
+    assert phases[ALTAIR].config.ALTAIR_FORK_EPOCH == 0
+    assert not hasattr(phases[ALTAIR].config, 'BELLATRIX_FORK_EPOCH')
+
+    assert phases[ALTAIR].config.ALTAIR_FORK_VERSION == spec.Version('0x11111111')
+    assert phases[BELLATRIX].config.ALTAIR_FORK_EPOCH == 0
+    assert phases[BELLATRIX].config.BELLATRIX_FORK_EPOCH == 4

tests/core/pyspec/eth2spec/test/context.py

@@ -1,4 +1,5 @@
 import pytest
+from copy import deepcopy
 from dataclasses import dataclass
 import importlib
 
@@ -309,14 +310,18 @@ def config_fork_epoch_overrides(spec, state):
             return overrides
 
 
-def spec_state_test_with_matching_config(fn):
+def with_matching_spec_config(emitted_fork=None):
     def decorator(fn):
         def wrapper(*args, spec: Spec, **kw):
-            conf = config_fork_epoch_overrides(spec, kw['state'])
-            overrides = with_config_overrides(conf)
-            return overrides(fn)(*args, spec=spec, **kw)
+            overrides = config_fork_epoch_overrides(spec, kw['state'])
+            deco = with_config_overrides(overrides, emitted_fork)
+            return deco(fn)(*args, spec=spec, **kw)
         return wrapper
-    return spec_test(with_state(decorator(single_phase(fn))))
+    return decorator
+
+
+def spec_state_test_with_matching_config(fn):
+    return spec_test(with_state(with_matching_spec_config()(single_phase(fn))))
 
 
 def expect_assertion_error(fn):
@@ -557,10 +562,30 @@ def _get_copy_of_spec(spec):
     module_spec = importlib.util.find_spec(module_path)
     module = importlib.util.module_from_spec(module_spec)
     module_spec.loader.exec_module(module)
+
+    # Preserve existing config overrides
+    module.config = deepcopy(spec.config)
+
     return module
 
 
-def with_config_overrides(config_overrides):
+def spec_with_config_overrides(spec, config_overrides):
+    # apply our overrides to a copy of it, and apply it to the spec
+    config = spec.config._asdict()
+    config.update((k, config_overrides[k]) for k in config.keys() & config_overrides.keys())
+    config_types = spec.Configuration.__annotations__
+    modified_config = {k: config_types[k](v) for k, v in config.items()}
+
+    spec.config = spec.Configuration(**modified_config)
+
+    # To output the changed config in a format compatible with yaml test vectors,
+    # the dict SSZ objects have to be converted into Python built-in types.
+    output_config = _get_basic_dict(modified_config)
+
+    return spec, output_config
+
+
+def with_config_overrides(config_overrides, emitted_fork=None, emit=True):
     """
     WARNING: the spec_test decorator must wrap this, to ensure the decorated test actually runs.
     This decorator forces the test to yield, and pytest doesn't run generator tests, and instead silently passes it.
@@ -570,23 +595,26 @@ def with_config_overrides(config_overrides):
     """
     def decorator(fn):
         def wrapper(*args, spec: Spec, **kw):
-            spec = _get_copy_of_spec(spec)
-
-            # apply our overrides to a copy of it, and apply it to the spec
-            config = spec.config._asdict()
-            config.update(config_overrides)
-            config_types = spec.Configuration.__annotations__
-            modified_config = {k: config_types[k](v) for k, v in config.items()}
-
-            # To output the changed config to could be serialized with yaml test vectors,
-            # the dict SSZ objects have to be converted into Python built-in types.
-            output_config = _get_basic_dict(modified_config)
-            yield 'config', 'cfg', output_config
-
-            spec.config = spec.Configuration(**modified_config)
+            # Apply config overrides to spec
+            spec, output_config = spec_with_config_overrides(_get_copy_of_spec(spec), config_overrides)
+
+            # Apply config overrides to additional phases, if present
+            if 'phases' in kw:
+                phases = {}
+                for fork in kw['phases']:
+                    phases[fork], output = spec_with_config_overrides(
+                        _get_copy_of_spec(kw['phases'][fork]), config_overrides)
+                    if emitted_fork == fork:
+                        output_config = output
+                kw['phases'] = phases
+
+            # Emit requested spec (with overrides)
+            if emit:
+                yield 'config', 'cfg', output_config
 
             # Run the function
             out = fn(*args, spec=spec, **kw)
+
             # If it's not returning None like a normal test function,
             # it's generating things, and we need to complete it before setting back the config.
             if out is not None: