From 49d96f92ef6f6750e55b0cf089d173de5a643852 Mon Sep 17 00:00:00 2001
From: Hsiao-Wei Wang <hsiaowei.eth@gmail.com>
Date: Fri, 12 Nov 2021 17:06:10 +0800
Subject: [PATCH] Add a sample altair-to-merge-only transition

---
 tests/core/pyspec/eth2spec/test/context.py    |  6 +++-
 .../pyspec/eth2spec/test/helpers/constants.py |  3 +-
 .../test/merge/transition/__init__.py         |  0
 .../test/merge/transition/test_transition.py  | 35 +++++++++++++++++++
 tests/generators/transition/main.py           | 11 ++++--
 5 files changed, 51 insertions(+), 4 deletions(-)
 create mode 100644 tests/core/pyspec/eth2spec/test/merge/transition/__init__.py
 create mode 100644 tests/core/pyspec/eth2spec/test/merge/transition/test_transition.py

diff --git a/tests/core/pyspec/eth2spec/test/context.py b/tests/core/pyspec/eth2spec/test/context.py
index 04da8e653..7f7c6bf28 100644
--- a/tests/core/pyspec/eth2spec/test/context.py
+++ b/tests/core/pyspec/eth2spec/test/context.py
@@ -581,7 +581,11 @@ def yield_fork_meta(fork_metas: Sequence[ForkMeta]):
         def wrapper(*args, **kw):
             phases = kw.pop('phases')
             spec = kw["spec"]
-            fork_meta = next(filter(lambda m: m.pre_fork_name == spec.fork, fork_metas))
+            try:
+                fork_meta = next(filter(lambda m: m.pre_fork_name == spec.fork, fork_metas))
+            except StopIteration:
+                dump_skipping_message(f"doesn't support this fork: {spec.fork}")
+
             post_spec = phases[fork_meta.post_fork_name]
 
             # Reset counter
diff --git a/tests/core/pyspec/eth2spec/test/helpers/constants.py b/tests/core/pyspec/eth2spec/test/helpers/constants.py
index 4f0f8bd6d..bb8f49cbc 100644
--- a/tests/core/pyspec/eth2spec/test/helpers/constants.py
+++ b/tests/core/pyspec/eth2spec/test/helpers/constants.py
@@ -27,7 +27,8 @@ ALL_FORK_UPGRADES = {
     ALTAIR: MERGE,
 }
 ALL_PRE_POST_FORKS = ALL_FORK_UPGRADES.items()
-
+AFTER_MERGE_UPGRADES = {key: value for key, value in ALL_FORK_UPGRADES.items() if key not in FORKS_BEFORE_ALTAIR}
+AFTER_MERGE_PRE_POST_FORKS = AFTER_MERGE_UPGRADES.items()
 
 #
 # Config
diff --git a/tests/core/pyspec/eth2spec/test/merge/transition/__init__.py b/tests/core/pyspec/eth2spec/test/merge/transition/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/core/pyspec/eth2spec/test/merge/transition/test_transition.py b/tests/core/pyspec/eth2spec/test/merge/transition/test_transition.py
new file mode 100644
index 000000000..d488d81dc
--- /dev/null
+++ b/tests/core/pyspec/eth2spec/test/merge/transition/test_transition.py
@@ -0,0 +1,35 @@
+from eth2spec.test.context import (
+    ForkMeta,
+    with_fork_metas,
+)
+from eth2spec.test.helpers.constants import (
+    AFTER_MERGE_PRE_POST_FORKS,
+)
+from eth2spec.test.helpers.fork_transition import (
+    do_fork,
+    transition_to_next_epoch_and_append_blocks,
+    transition_until_fork,
+)
+
+
+@with_fork_metas([
+    ForkMeta(pre_fork_name=pre, post_fork_name=post, fork_epoch=2) for pre, post in AFTER_MERGE_PRE_POST_FORKS
+])
+def test_sample_transition(state, fork_epoch, spec, post_spec, pre_tag, post_tag):
+    transition_until_fork(spec, state, fork_epoch)
+
+    # check pre state
+    assert spec.get_current_epoch(state) < fork_epoch
+
+    yield "pre", state
+
+    # irregular state transition to handle fork:
+    blocks = []
+    state, block = do_fork(state, spec, post_spec, fork_epoch)
+    blocks.append(post_tag(block))
+
+    # continue regular state transition with new spec into next epoch
+    transition_to_next_epoch_and_append_blocks(post_spec, state, post_tag, blocks, only_last_block=True)
+
+    yield "blocks", blocks
+    yield "post", state
diff --git a/tests/generators/transition/main.py b/tests/generators/transition/main.py
index de00e774a..d6195de68 100644
--- a/tests/generators/transition/main.py
+++ b/tests/generators/transition/main.py
@@ -3,7 +3,7 @@ from typing import Iterable
 from eth2spec.test.helpers.constants import (
     MINIMAL,
     MAINNET,
-    ALL_FORK_UPGRADES,
+    ALL_PRE_POST_FORKS,
 )
 from eth2spec.gen_helpers.gen_base import gen_runner, gen_typing
 from eth2spec.gen_helpers.gen_from_tests.gen import (
@@ -16,6 +16,9 @@ from eth2spec.test.altair.transition import (
     test_slashing as test_altair_slashing,
     test_operations as test_altair_operations,
 )
+from eth2spec.test.merge.transition import (
+    test_transition as test_merge_transition,
+)
 
 
 def create_provider(tests_src, preset_name: str, pre_fork_name: str, post_fork_name: str) -> gen_typing.TestProvider:
@@ -44,7 +47,11 @@ if __name__ == "__main__":
         test_altair_slashing,
         test_altair_operations,
     )
-    for transition_test_module in altair_tests:
+    merge_tests = (
+        test_merge_transition,
+    )
+    all_tests = altair_tests + merge_tests
+    for transition_test_module in all_tests:
         for pre_fork, post_fork in ALL_PRE_POST_FORKS:
             gen_runner.run_generator("transition", [
                 create_provider(transition_test_module, MINIMAL, pre_fork, post_fork),