From 35d444b9d225d9a18a1d3fc6e87c2294bd74ab05 Mon Sep 17 00:00:00 2001
From: Hsiao-Wei Wang <hsiaowei.eth@gmail.com>
Date: Wed, 25 Oct 2023 02:10:55 +0800
Subject: [PATCH] Update test format and output
 `should_override_forkchoice_update` result

---
 .../test/bellatrix/fork_choice/test_reorg.py  | 21 ++++++++++++++++---
 tests/formats/fork_choice/README.md           |  1 +
 tests/generators/fork_choice/main.py          |  2 ++
 3 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/tests/core/pyspec/eth2spec/test/bellatrix/fork_choice/test_reorg.py b/tests/core/pyspec/eth2spec/test/bellatrix/fork_choice/test_reorg.py
index cb7adac1c..617e0909c 100644
--- a/tests/core/pyspec/eth2spec/test/bellatrix/fork_choice/test_reorg.py
+++ b/tests/core/pyspec/eth2spec/test/bellatrix/fork_choice/test_reorg.py
@@ -17,6 +17,7 @@ from eth2spec.test.helpers.fork_choice import (
     apply_next_slots_with_attestations,
     get_genesis_forkchoice_store_and_block,
     on_tick_and_append_step,
+    output_store_checks,
     tick_and_add_block,
     tick_and_run_on_attestation,
 )
@@ -56,7 +57,15 @@ def test_should_override_forkchoice_update__false(spec, state):
     current_time = slot * spec.config.SECONDS_PER_SLOT + store.genesis_time
     on_tick_and_append_step(spec, store, current_time, test_steps)
 
-    assert not spec.should_override_forkchoice_update(store, head_root)
+    should_override = spec.should_override_forkchoice_update(store, head_root)
+    assert not should_override
+
+    output_store_checks(spec, store, test_steps)
+    test_steps.append({
+        'checks': {
+            'should_override_forkchoice_update': should_override,
+        }
+    })
 
     yield 'steps', test_steps
 
@@ -155,8 +164,14 @@ def test_should_override_forkchoice_update__true(spec, state):
     assert spec.is_head_weak(store, head_root)
     assert spec.is_parent_strong(store, parent_root)
 
-    assert spec.should_override_forkchoice_update(store, head_root)
+    should_override = spec.should_override_forkchoice_update(store, head_root)
+    assert should_override
 
-    # TODO: export the `should_override_forkchoice_update` result to test vectors?
+    output_store_checks(spec, store, test_steps)
+    test_steps.append({
+        'checks': {
+            'should_override_forkchoice_update': should_override,
+        }
+    })
 
     yield 'steps', test_steps
diff --git a/tests/formats/fork_choice/README.md b/tests/formats/fork_choice/README.md
index d23de865b..3bd773f52 100644
--- a/tests/formats/fork_choice/README.md
+++ b/tests/formats/fork_choice/README.md
@@ -177,6 +177,7 @@ finalized_checkpoint: {
     root: string,             -- Encoded 32-byte value from store.finalized_checkpoint.root
 }
 proposer_boost_root: string   -- Encoded 32-byte value from store.proposer_boost_root
+should_override_forkchoice_update: bool  -- [New in Bellatrix] the result of `should_override_forkchoice_update(store, head_root)`
 ```
 
 For example:
diff --git a/tests/generators/fork_choice/main.py b/tests/generators/fork_choice/main.py
index 7ff028cd8..e8d19eb09 100644
--- a/tests/generators/fork_choice/main.py
+++ b/tests/generators/fork_choice/main.py
@@ -10,12 +10,14 @@ if __name__ == "__main__":
         'ex_ante',
         'reorg',
         'withholding',
+        'get_proposer_head',
     ]}
 
     # For merge `on_merge_block` test kind added with `pow_block_N.ssz` files with several
     # PowBlock's which should be resolved by `get_pow_block(hash: Hash32) -> PowBlock` function
     _new_bellatrix_mods = {key: 'eth2spec.test.bellatrix.fork_choice.test_' + key for key in [
         'on_merge_block',
+        'reorg',
     ]}
     bellatrix_mods = combine_mods(_new_bellatrix_mods, altair_mods)
     capella_mods = bellatrix_mods  # No additional Capella specific fork choice tests