Rework `on_block` unit tests

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

@@ -1,4 +1,5 @@
 from eth_utils import encode_hex
+from eth2spec.test.helpers.attestations import next_epoch_with_attestations
 
 
 def get_anchor_root(spec, state):
@@ -18,7 +19,7 @@ def add_block_to_store(spec, store, signed_block):
     spec.on_block(store, signed_block)
 
 
-def tick_and_run_on_block(spec, store, signed_block, test_steps=None):
+def tick_and_add_block(spec, store, signed_block, test_steps=None, valid=True):
     if test_steps is None:
         test_steps = []
 
@@ -28,7 +29,7 @@ def tick_and_run_on_block(spec, store, signed_block, test_steps=None):
     if store.time < block_time:
         on_tick_and_append_step(spec, store, block_time, test_steps)
 
-    yield from run_on_block(spec, store, signed_block, test_steps)
+    yield from add_block(spec, store, signed_block, test_steps, valid=valid)
 
 
 def tick_and_run_on_attestation(spec, store, attestation, test_steps=None):
@@ -73,28 +74,47 @@ def on_tick_and_append_step(spec, store, time, test_steps):
     test_steps.append({'tick': int(time)})
 
 
-def run_on_block(spec, store, signed_block, test_steps, valid=True):
+def run_on_block(spec, store, signed_block, valid=True):
-    yield get_block_file_name(signed_block), signed_block
     if not valid:
         try:
             spec.on_block(store, signed_block)
         except AssertionError:
-            test_steps.append({
-                'block': get_block_file_name(signed_block),
-                'valid': True,
-            })
             return
         else:
             assert False
 
     spec.on_block(store, signed_block)
+    assert store.blocks[signed_block.message.hash_tree_root()] == signed_block.message
+
+
+def add_block(spec, store, signed_block, test_steps=None, valid=True):
+    if test_steps is None:
+        test_steps = []
+
+    yield get_block_file_name(signed_block), signed_block
+
+    if not valid:
+        try:
+            run_on_block(spec, store, signed_block, valid=True)
+        except AssertionError:
+            test_steps.append({
+                'block': get_block_file_name(signed_block),
+                'valid': False,
+            })
+            return
+        else:
+            assert False
+
+    run_on_block(spec, store, signed_block, valid=True)
     test_steps.append({'block': get_block_file_name(signed_block)})
 
     # An on_block step implies receiving block's attestations
     for attestation in signed_block.message.body.attestations:
         spec.on_attestation(store, attestation)
 
-    assert store.blocks[signed_block.message.hash_tree_root()] == signed_block.message
+    block_root = signed_block.message.hash_tree_root()
+    assert store.blocks[block_root] == signed_block.message
+    assert store.block_states[block_root].hash_tree_root() == signed_block.message.state_root
     test_steps.append({
         'checks': {
             'time': int(store.time),
@@ -113,3 +133,20 @@ def get_formatted_head_output(spec, store):
         'slot': int(slot),
         'root': encode_hex(head),
     }
+
+
+def apply_next_epoch_with_attestations(spec, state, store, test_steps=None):
+    if test_steps is None:
+        test_steps = []
+
+    _, new_signed_blocks, post_state = next_epoch_with_attestations(spec, state, True, False)
+    for signed_block in new_signed_blocks:
+        block = signed_block.message
+        yield from tick_and_add_block(spec, store, signed_block, test_steps)
+        block_root = block.hash_tree_root()
+        assert store.blocks[block_root] == block
+        last_signed_block = signed_block
+
+    assert store.block_states[block_root].hash_tree_root() == post_state.hash_tree_root()
+
+    return post_state, store, last_signed_block

tests/core/pyspec/eth2spec/test/phase0/fork_choice/test_get_head.py

@@ -11,12 +11,12 @@ from eth2spec.test.helpers.block import build_empty_block_for_next_slot
 from eth2spec.test.helpers.constants import MINIMAL
 from eth2spec.test.helpers.fork_choice import (
     tick_and_run_on_attestation,
-    tick_and_run_on_block,
+    tick_and_add_block,
     get_anchor_root,
     get_genesis_forkchoice_store_and_block,
     get_formatted_head_output,
     on_tick_and_append_step,
-    run_on_block,
+    add_block,
 )
 from eth2spec.test.helpers.state import (
     next_epoch,
@@ -68,12 +68,12 @@ def test_chain_no_attestations(spec, state):
     # On receiving a block of `GENESIS_SLOT + 1` slot
     block_1 = build_empty_block_for_next_slot(spec, state)
     signed_block_1 = state_transition_and_sign_block(spec, state, block_1)
-    yield from tick_and_run_on_block(spec, store, signed_block_1, test_steps)
+    yield from tick_and_add_block(spec, store, signed_block_1, test_steps)
 
     # On receiving a block of next epoch
     block_2 = build_empty_block_for_next_slot(spec, state)
     signed_block_2 = state_transition_and_sign_block(spec, state, block_2)
-    yield from tick_and_run_on_block(spec, store, signed_block_2, test_steps)
+    yield from tick_and_add_block(spec, store, signed_block_2, test_steps)
 
     assert spec.get_head(store) == spec.hash_tree_root(block_2)
     test_steps.append({
@@ -107,14 +107,14 @@ def test_split_tie_breaker_no_attestations(spec, state):
     block_1_state = genesis_state.copy()
     block_1 = build_empty_block_for_next_slot(spec, block_1_state)
     signed_block_1 = state_transition_and_sign_block(spec, block_1_state, block_1)
-    yield from tick_and_run_on_block(spec, store, signed_block_1, test_steps)
+    yield from tick_and_add_block(spec, store, signed_block_1, test_steps)
 
     # additional block at slot 1
     block_2_state = genesis_state.copy()
     block_2 = build_empty_block_for_next_slot(spec, block_2_state)
     block_2.body.graffiti = b'\x42' * 32
     signed_block_2 = state_transition_and_sign_block(spec, block_2_state, block_2)
-    yield from tick_and_run_on_block(spec, store, signed_block_2, test_steps)
+    yield from tick_and_add_block(spec, store, signed_block_2, test_steps)
 
     highest_root = max(spec.hash_tree_root(block_1), spec.hash_tree_root(block_2))
     assert spec.get_head(store) == highest_root
@@ -150,14 +150,14 @@ def test_shorter_chain_but_heavier_weight(spec, state):
     for _ in range(3):
         long_block = build_empty_block_for_next_slot(spec, long_state)
         signed_long_block = state_transition_and_sign_block(spec, long_state, long_block)
-        yield from tick_and_run_on_block(spec, store, signed_long_block, test_steps)
+        yield from tick_and_add_block(spec, store, signed_long_block, test_steps)
 
     # build short tree
     short_state = genesis_state.copy()
     short_block = build_empty_block_for_next_slot(spec, short_state)
     short_block.body.graffiti = b'\x42' * 32
     signed_short_block = state_transition_and_sign_block(spec, short_state, short_block)
-    yield from tick_and_run_on_block(spec, store, signed_short_block, test_steps)
+    yield from tick_and_add_block(spec, store, signed_short_block, test_steps)
 
     short_attestation = get_valid_attestation(spec, short_state, short_block.slot, signed=True)
     yield from tick_and_run_on_attestation(spec, store, short_attestation, test_steps)
@@ -200,7 +200,7 @@ def test_filtered_block_tree(spec, state):
     current_time = state.slot * spec.config.SECONDS_PER_SLOT + store.genesis_time
     on_tick_and_append_step(spec, store, current_time, test_steps)
     for signed_block in signed_blocks:
-        yield from run_on_block(spec, store, signed_block, test_steps)
+        yield from add_block(spec, store, signed_block, test_steps)
 
     assert store.justified_checkpoint == state.current_justified_checkpoint
 
@@ -247,7 +247,7 @@ def test_filtered_block_tree(spec, state):
     on_tick_and_append_step(spec, store, current_time, test_steps)
 
     # include rogue block and associated attestations in the store
-    yield from run_on_block(spec, store, signed_rogue_block, test_steps)
+    yield from add_block(spec, store, signed_rogue_block, test_steps)
 
     for attestation in attestations:
         yield from tick_and_run_on_attestation(spec, store, attestation, test_steps)

tests/core/pyspec/eth2spec/test/phase0/fork_choice/test_on_block.py

@@ -1,38 +1,30 @@
 from eth2spec.utils.ssz.ssz_impl import hash_tree_root
 
 from eth2spec.test.context import MINIMAL, spec_state_test, with_all_phases, with_presets
-from eth2spec.test.helpers.attestations import next_epoch_with_attestations
+from eth2spec.test.helpers.block import (
-from eth2spec.test.helpers.block import build_empty_block_for_next_slot, build_empty_block
+    build_empty_block_for_next_slot,
+    build_empty_block,
+    transition_unsigned_block,
+    sign_block,
+)
 from eth2spec.test.helpers.fork_choice import (
     get_genesis_forkchoice_store_and_block,
     on_tick_and_append_step,
-    run_on_block,
+    add_block,
-    tick_and_run_on_block,
+    tick_and_add_block,
+    apply_next_epoch_with_attestations,
+)
+from eth2spec.test.helpers.state import (
+    next_epoch,
+    state_transition_and_sign_block,
 )
-from eth2spec.test.helpers.state import next_epoch, state_transition_and_sign_block
-
-
-def apply_next_epoch_with_attestations(spec, state, store, test_steps=None):
-    if test_steps is None:
-        test_steps = []
-
-    _, new_signed_blocks, post_state = next_epoch_with_attestations(spec, state, True, False)
-    for signed_block in new_signed_blocks:
-        block = signed_block.message
-        block_root = hash_tree_root(block)
-        store.blocks[block_root] = block
-        store.block_states[block_root] = post_state
-        yield from tick_and_run_on_block(spec, store, signed_block, test_steps)
-        last_signed_block = signed_block
-
-    return post_state, store, last_signed_block
 
 
 @with_all_phases
 @spec_state_test
 def test_basic(spec, state):
-    # Initialization
     test_steps = []
+    # Initialization
     store, anchor_block = get_genesis_forkchoice_store_and_block(spec, state)
     yield 'anchor_state', state
     yield 'anchor_block', anchor_block
@@ -43,14 +35,14 @@ def test_basic(spec, state):
     # On receiving a block of `GENESIS_SLOT + 1` slot
     block = build_empty_block_for_next_slot(spec, state)
     signed_block = state_transition_and_sign_block(spec, state, block)
-    yield from tick_and_run_on_block(spec, store, signed_block, test_steps)
+    yield from tick_and_add_block(spec, store, signed_block, test_steps)
     assert spec.get_head(store) == signed_block.message.hash_tree_root()
 
     # On receiving a block of next epoch
     store.time = current_time + spec.config.SECONDS_PER_SLOT * spec.SLOTS_PER_EPOCH
     block = build_empty_block(spec, state, state.slot + spec.SLOTS_PER_EPOCH)
     signed_block = state_transition_and_sign_block(spec, state, block)
-    yield from tick_and_run_on_block(spec, store, signed_block, test_steps)
+    yield from tick_and_add_block(spec, store, signed_block, test_steps)
     assert spec.get_head(store) == signed_block.message.hash_tree_root()
 
     yield 'steps', test_steps
@@ -74,6 +66,7 @@ def test_on_block_checkpoints(spec, state):
     # Run for 1 epoch with full attestations
     next_epoch(spec, state)
     on_tick_and_append_step(spec, store, store.time + state.slot * spec.config.SECONDS_PER_SLOT, test_steps)
+
     state, store, last_signed_block = yield from apply_next_epoch_with_attestations(spec, state, store, test_steps)
     last_block_root = hash_tree_root(last_signed_block.message)
     assert spec.get_head(store) == last_block_root
@@ -90,7 +83,7 @@ def test_on_block_checkpoints(spec, state):
 
     block = build_empty_block_for_next_slot(spec, fin_state)
     signed_block = state_transition_and_sign_block(spec, fin_state.copy(), block)
-    yield from tick_and_run_on_block(spec, store, signed_block, test_steps)
+    yield from tick_and_add_block(spec, store, signed_block, test_steps)
     assert spec.get_head(store) == signed_block.message.hash_tree_root()
     yield 'steps', test_steps
 
@@ -107,233 +100,36 @@ def test_on_block_future_block(spec, state):
     on_tick_and_append_step(spec, store, current_time, test_steps)
     assert store.time == current_time
 
-    # do not tick time
+    # Do NOT tick time to `GENESIS_SLOT + 1` slot
-
     # Fail receiving block of `GENESIS_SLOT + 1` slot
     block = build_empty_block_for_next_slot(spec, state)
     signed_block = state_transition_and_sign_block(spec, state, block)
-    run_on_block(spec, store, signed_block, test_steps, valid=False)
+    yield from add_block(spec, store, signed_block, test_steps, valid=False)
 
     yield 'steps', test_steps
 
 
-# @with_all_phases
+@with_all_phases
-# @spec_state_test
+@spec_state_test
-# def test_on_block_bad_parent_root(spec, state):
+def test_on_block_bad_parent_root(spec, state):
-#     test_steps = []
+    test_steps = []
-#     # Initialization
+    # Initialization
-#     store = get_genesis_forkchoice_store(spec, state)
+    store, anchor_block = get_genesis_forkchoice_store_and_block(spec, state)
-#     time = 100
+    yield 'anchor_state', state
-#     on_tick_and_append_step(spec, store, time, test_steps)
+    yield 'anchor_block', anchor_block
+    current_time = state.slot * spec.config.SECONDS_PER_SLOT + store.genesis_time
+    on_tick_and_append_step(spec, store, current_time, test_steps)
+    assert store.time == current_time
 
-#     # Fail receiving block of `GENESIS_SLOT + 1` slot
+    # Fail receiving block of `GENESIS_SLOT + 1` slot
-#     block = build_empty_block_for_next_slot(spec, state)
+    block = build_empty_block_for_next_slot(spec, state)
-#     transition_unsigned_block(spec, state, block)
+    transition_unsigned_block(spec, state, block)
-#     block.state_root = state.hash_tree_root()
+    block.state_root = state.hash_tree_root()
 
-#     block.parent_root = b'\x45' * 32
+    block.parent_root = b'\x45' * 32
 
-#     signed_block = sign_block(spec, state, block)
+    signed_block = sign_block(spec, state, block)
 
-#     run_on_block(spec, store, signed_block, test_steps, valid=False)
+    yield from add_block(spec, store, signed_block, test_steps, valid=False)
 
-
+    yield 'steps', test_steps
-# @with_all_phases
-# @spec_state_test
-# def test_on_block_before_finalized(spec, state):
-#     test_steps = []
-#     # Initialization
-#     store = get_genesis_forkchoice_store(spec, state)
-#     time = 100
-#     on_tick_and_append_step(spec, store, time, test_steps)
-
-#     store.finalized_checkpoint = spec.Checkpoint(
-#         epoch=store.finalized_checkpoint.epoch + 2,
-#         root=store.finalized_checkpoint.root
-#     )
-
-#     # Fail receiving block of `GENESIS_SLOT + 1` slot
-#     block = build_empty_block_for_next_slot(spec, state)
-#     signed_block = state_transition_and_sign_block(spec, state, block)
-#     run_on_block(spec, store, signed_block, test_steps, valid=False)
-
-
-# @with_all_phases
-# @spec_state_test
-# def test_on_block_finalized_skip_slots(spec, state):
-#     test_steps = []
-#     # Initialization
-#     store = get_genesis_forkchoice_store(spec, state)
-#     time = 100
-#     on_tick_and_append_step(spec, store, time, test_steps)
-
-#     store.finalized_checkpoint = spec.Checkpoint(
-#         epoch=store.finalized_checkpoint.epoch + 2,
-#         root=store.finalized_checkpoint.root
-#     )
-
-#     # Build block that includes the skipped slots up to finality in chain
-#     block = build_empty_block(spec, state, spec.compute_start_slot_at_epoch(store.finalized_checkpoint.epoch) + 2)
-#     signed_block = state_transition_and_sign_block(spec, state, block)
-#     on_tick_and_append_step(spec, store, store.time + state.slot * spec.config.SECONDS_PER_SLOT, test_steps)
-#     run_on_block(spec, store, signed_block, test_steps)
-
-
-# @with_all_phases
-# @spec_state_test
-# def test_on_block_finalized_skip_slots_not_in_skip_chain(spec, state):
-#     test_steps = []
-#     # Initialization
-#     transition_to(spec, state, state.slot + spec.SLOTS_PER_EPOCH - 1)
-#     block = build_empty_block_for_next_slot(spec, state)
-#     transition_unsigned_block(spec, state, block)
-#     block.state_root = state.hash_tree_root()
-#     store = spec.get_forkchoice_store(state, block)
-#     store.finalized_checkpoint = spec.Checkpoint(
-#         epoch=store.finalized_checkpoint.epoch + 2,
-#         root=store.finalized_checkpoint.root
-#     )
-
-#     # First transition through the epoch to ensure no skipped slots
-#     state, store, _ = apply_next_epoch_with_attestations(spec, state, store)
-
-#     # Now build a block at later slot than finalized epoch
-#     # Includes finalized block in chain, but not at appropriate skip slot
-#     block = build_empty_block(spec, state, spec.compute_start_slot_at_epoch(store.finalized_checkpoint.epoch) + 2)
-#     signed_block = state_transition_and_sign_block(spec, state, block)
-#     on_tick_and_append_step(spec, store, store.time + state.slot * spec.config.SECONDS_PER_SLOT, test_steps)
-#     run_on_block(spec, store, signed_block, test_steps, valid=False)
-
-
-# @with_all_phases
-# @spec_state_test
-# def test_on_block_update_justified_checkpoint_within_safe_slots(spec, state):
-#     test_steps = []
-#     # Initialization
-#     store = get_genesis_forkchoice_store(spec, state)
-#     time = 0
-#     on_tick_and_append_step(spec, store, time, test_steps)
-
-#     next_epoch(spec, state)
-#     on_tick_and_append_step(spec, store, store.time + state.slot * spec.config.SECONDS_PER_SLOT, test_steps)
-#     state, store, last_signed_block = apply_next_epoch_with_attestations(spec, state, store)
-#     next_epoch(spec, state)
-#     on_tick_and_append_step(spec, store, store.time + state.slot * spec.config.SECONDS_PER_SLOT, test_steps)
-#     last_block_root = hash_tree_root(last_signed_block.message)
-
-#     # Mock the justified checkpoint
-#     just_state = store.block_states[last_block_root]
-#     new_justified = spec.Checkpoint(
-#         epoch=just_state.current_justified_checkpoint.epoch + 1,
-#         root=b'\x77' * 32,
-#     )
-#     just_state.current_justified_checkpoint = new_justified
-
-#     block = build_empty_block_for_next_slot(spec, just_state)
-#     signed_block = state_transition_and_sign_block(spec, deepcopy(just_state), block)
-#     assert spec.get_current_slot(store) % spec.SLOTS_PER_EPOCH < spec.SAFE_SLOTS_TO_UPDATE_JUSTIFIED
-#     run_on_block(spec, store, signed_block, test_steps)
-
-#     assert store.justified_checkpoint == new_justified
-
-
-# @with_all_phases
-# @spec_state_test
-# def test_on_block_outside_safe_slots_and_multiple_better_justified(spec, state):
-#     test_steps = []
-#     # Initialization
-#     store = get_genesis_forkchoice_store(spec, state)
-#     time = 0
-#     on_tick_and_append_step(spec, store, time, test_steps)
-
-#     next_epoch(spec, state)
-#     on_tick_and_append_step(spec, store, store.time + state.slot * spec.config.SECONDS_PER_SLOT, test_steps)
-#     state, store, last_signed_block = apply_next_epoch_with_attestations(spec, state, store)
-#     next_epoch(spec, state)
-#     on_tick_and_append_step(spec, store, store.time + state.slot * spec.config.SECONDS_PER_SLOT, test_steps)
-#     last_block_root = hash_tree_root(last_signed_block.message)
-
-#     # Mock justified block in store
-#     just_block = build_empty_block_for_next_slot(spec, state)
-#     # Slot is same as justified checkpoint so does not trigger an override in the store
-#     just_block.slot = spec.compute_start_slot_at_epoch(store.justified_checkpoint.epoch)
-#     store.blocks[just_block.hash_tree_root()] = just_block
-
-#     # Step time past safe slots
-#     time = store.time + spec.SAFE_SLOTS_TO_UPDATE_JUSTIFIED * spec.config.SECONDS_PER_SLOT
-#     on_tick_and_append_step(spec, store, time, test_steps)
-#     assert spec.get_current_slot(store) % spec.SLOTS_PER_EPOCH >= spec.SAFE_SLOTS_TO_UPDATE_JUSTIFIED
-
-#     previously_justified = store.justified_checkpoint
-
-#     # Add a series of new blocks with "better" justifications
-#     best_justified_checkpoint = spec.Checkpoint(epoch=0)
-#     for i in range(3, 0, -1):
-#         just_state = store.block_states[last_block_root]
-#         new_justified = spec.Checkpoint(
-#             epoch=previously_justified.epoch + i,
-#             root=just_block.hash_tree_root(),
-#         )
-#         if new_justified.epoch > best_justified_checkpoint.epoch:
-#             best_justified_checkpoint = new_justified
-
-#         just_state.current_justified_checkpoint = new_justified
-
-#         block = build_empty_block_for_next_slot(spec, just_state)
-#         signed_block = state_transition_and_sign_block(spec, deepcopy(just_state), block)
-
-#         run_on_block(spec, store, signed_block, test_steps)
-
-#     assert store.justified_checkpoint == previously_justified
-#     # ensure the best from the series was stored
-#     assert store.best_justified_checkpoint == best_justified_checkpoint
-
-
-# @with_all_phases
-# @spec_state_test
-# def test_on_block_outside_safe_slots_but_finality(spec, state):
-#     test_steps = []
-#     # Initialization
-#     store = get_genesis_forkchoice_store(spec, state)
-#     time = 100
-#     on_tick_and_append_step(spec, store, time, test_steps)
-
-#     next_epoch(spec, state)
-#     on_tick_and_append_step(spec, store, store.time + state.slot * spec.config.SECONDS_PER_SLOT, test_steps)
-#     state, store, last_signed_block = apply_next_epoch_with_attestations(spec, state, store)
-#     next_epoch(spec, state)
-#     on_tick_and_append_step(spec, store, store.time + state.slot * spec.config.SECONDS_PER_SLOT, test_steps)
-#     last_block_root = hash_tree_root(last_signed_block.message)
-
-#     # Mock justified block in store
-#     just_block = build_empty_block_for_next_slot(spec, state)
-#     # Slot is same as justified checkpoint so does not trigger an override in the store
-#     just_block.slot = spec.compute_start_slot_at_epoch(store.justified_checkpoint.epoch)
-#     store.blocks[just_block.hash_tree_root()] = just_block
-
-#     # Step time past safe slots
-#     time = store.time + spec.SAFE_SLOTS_TO_UPDATE_JUSTIFIED * spec.config.SECONDS_PER_SLOT
-#     on_tick_and_append_step(spec, store, time, test_steps)
-#     assert spec.get_current_slot(store) % spec.SLOTS_PER_EPOCH >= spec.SAFE_SLOTS_TO_UPDATE_JUSTIFIED
-
-#     # Mock justified and finalized update in state
-#     just_fin_state = store.block_states[last_block_root]
-#     new_justified = spec.Checkpoint(
-#         epoch=store.justified_checkpoint.epoch + 1,
-#         root=just_block.hash_tree_root(),
-#     )
-#     new_finalized = spec.Checkpoint(
-#         epoch=store.finalized_checkpoint.epoch + 1,
-#         root=just_block.parent_root,
-#     )
-#     just_fin_state.current_justified_checkpoint = new_justified
-#     just_fin_state.finalized_checkpoint = new_finalized
-
-#     # Build and add block that includes the new justified/finalized info
-#     block = build_empty_block_for_next_slot(spec, just_fin_state)
-#     signed_block = state_transition_and_sign_block(spec, deepcopy(just_fin_state), block)
-
-#     run_on_block(spec, store, signed_block, test_steps)
-
-#     assert store.finalized_checkpoint == new_finalized
-#     assert store.justified_checkpoint == new_justified

tests/core/pyspec/eth2spec/test/phase0/unittests/fork_choice/test_on_block.py

@@ -1,123 +1,17 @@
 from copy import deepcopy
 from eth2spec.utils.ssz.ssz_impl import hash_tree_root
 
-from eth2spec.test.context import with_all_phases, spec_state_test
+from eth2spec.test.context import with_all_phases, spec_state_test, with_phases, PHASE0
-from eth2spec.test.helpers.attestations import next_epoch_with_attestations
+from eth2spec.test.helpers.block import build_empty_block_for_next_slot, transition_unsigned_block, \
-from eth2spec.test.helpers.block import build_empty_block_for_next_slot, sign_block, transition_unsigned_block, \
+    build_empty_block, sign_block
-    build_empty_block
+from eth2spec.test.helpers.fork_choice import (
-from eth2spec.test.helpers.fork_choice import get_genesis_forkchoice_store
+    get_genesis_forkchoice_store,
+    run_on_block,
+    apply_next_epoch_with_attestations,
+)
 from eth2spec.test.helpers.state import next_epoch, state_transition_and_sign_block, transition_to
 
 
-def run_on_block(spec, store, signed_block, valid=True):
-    if not valid:
-        try:
-            spec.on_block(store, signed_block)
-        except AssertionError:
-            return
-        else:
-            assert False
-
-    spec.on_block(store, signed_block)
-    assert store.blocks[hash_tree_root(signed_block.message)] == signed_block.message
-
-
-def apply_next_epoch_with_attestations(spec, state, store):
-    _, new_signed_blocks, post_state = next_epoch_with_attestations(spec, state, True, False)
-    for signed_block in new_signed_blocks:
-        block = signed_block.message
-        block_root = hash_tree_root(block)
-        store.blocks[block_root] = block
-        store.block_states[block_root] = post_state
-        last_signed_block = signed_block
-    spec.on_tick(store, store.time + state.slot * spec.config.SECONDS_PER_SLOT)
-    return post_state, store, last_signed_block
-
-
-@with_all_phases
-@spec_state_test
-def test_basic(spec, state):
-    # Initialization
-    store = get_genesis_forkchoice_store(spec, state)
-    time = 100
-    spec.on_tick(store, time)
-    assert store.time == time
-
-    # On receiving a block of `GENESIS_SLOT + 1` slot
-    block = build_empty_block_for_next_slot(spec, state)
-    signed_block = state_transition_and_sign_block(spec, state, block)
-    run_on_block(spec, store, signed_block)
-
-    # On receiving a block of next epoch
-    store.time = time + spec.config.SECONDS_PER_SLOT * spec.SLOTS_PER_EPOCH
-    block = build_empty_block(spec, state, state.slot + spec.SLOTS_PER_EPOCH)
-    signed_block = state_transition_and_sign_block(spec, state, block)
-
-    run_on_block(spec, store, signed_block)
-
-    # TODO: add tests for justified_root and finalized_root
-
-
-@with_all_phases
-@spec_state_test
-def test_on_block_checkpoints(spec, state):
-    # Initialization
-    store = get_genesis_forkchoice_store(spec, state)
-    time = 100
-    spec.on_tick(store, time)
-
-    next_epoch(spec, state)
-    spec.on_tick(store, store.time + state.slot * spec.config.SECONDS_PER_SLOT)
-    state, store, last_signed_block = apply_next_epoch_with_attestations(spec, state, store)
-    next_epoch(spec, state)
-    spec.on_tick(store, store.time + state.slot * spec.config.SECONDS_PER_SLOT)
-    last_block_root = hash_tree_root(last_signed_block.message)
-
-    # Mock the finalized_checkpoint
-    fin_state = store.block_states[last_block_root]
-    fin_state.finalized_checkpoint = (
-        store.block_states[last_block_root].current_justified_checkpoint
-    )
-
-    block = build_empty_block_for_next_slot(spec, fin_state)
-    signed_block = state_transition_and_sign_block(spec, deepcopy(fin_state), block)
-    run_on_block(spec, store, signed_block)
-
-
-@with_all_phases
-@spec_state_test
-def test_on_block_future_block(spec, state):
-    # Initialization
-    store = get_genesis_forkchoice_store(spec, state)
-
-    # do not tick time
-
-    # Fail receiving block of `GENESIS_SLOT + 1` slot
-    block = build_empty_block_for_next_slot(spec, state)
-    signed_block = state_transition_and_sign_block(spec, state, block)
-    run_on_block(spec, store, signed_block, False)
-
-
-@with_all_phases
-@spec_state_test
-def test_on_block_bad_parent_root(spec, state):
-    # Initialization
-    store = get_genesis_forkchoice_store(spec, state)
-    time = 100
-    spec.on_tick(store, time)
-
-    # Fail receiving block of `GENESIS_SLOT + 1` slot
-    block = build_empty_block_for_next_slot(spec, state)
-    transition_unsigned_block(spec, state, block)
-    block.state_root = state.hash_tree_root()
-
-    block.parent_root = b'\x45' * 32
-
-    signed_block = sign_block(spec, state, block)
-
-    run_on_block(spec, store, signed_block, False)
-
-
 @with_all_phases
 @spec_state_test
 def test_on_block_before_finalized(spec, state):
@@ -134,7 +28,7 @@ def test_on_block_before_finalized(spec, state):
     # Fail receiving block of `GENESIS_SLOT + 1` slot
     block = build_empty_block_for_next_slot(spec, state)
     signed_block = state_transition_and_sign_block(spec, state, block)
-    run_on_block(spec, store, signed_block, False)
+    run_on_block(spec, store, signed_block, valid=False)
 
 
 @with_all_phases
@@ -166,23 +60,31 @@ def test_on_block_finalized_skip_slots_not_in_skip_chain(spec, state):
     transition_unsigned_block(spec, state, block)
     block.state_root = state.hash_tree_root()
     store = spec.get_forkchoice_store(state, block)
-    store.finalized_checkpoint = spec.Checkpoint(
-        epoch=store.finalized_checkpoint.epoch + 2,
-        root=store.finalized_checkpoint.root
-    )
 
-    # First transition through the epoch to ensure no skipped slots
+    current_time = state.slot * spec.config.SECONDS_PER_SLOT + store.genesis_time
-    state, store, _ = apply_next_epoch_with_attestations(spec, state, store)
+    spec.on_tick(store, current_time)
+    assert store.time == current_time
+
+    pre_finalized_checkpoint_epoch = store.finalized_checkpoint.epoch
+
+    # Finalized
+    for _ in range(3):
+        state, store, _ = yield from apply_next_epoch_with_attestations(spec, state, store)
+    assert store.finalized_checkpoint.epoch == pre_finalized_checkpoint_epoch + 1
 
     # Now build a block at later slot than finalized epoch
     # Includes finalized block in chain, but not at appropriate skip slot
-    block = build_empty_block(spec, state, spec.compute_start_slot_at_epoch(store.finalized_checkpoint.epoch) + 2)
+    pre_state = store.block_states[block.hash_tree_root()]
-    signed_block = state_transition_and_sign_block(spec, state, block)
+    block = build_empty_block(spec,
+                              state=pre_state,
+                              slot=spec.compute_start_slot_at_epoch(store.finalized_checkpoint.epoch) + 2)
+    signed_block = sign_block(spec, pre_state, block)
+
     spec.on_tick(store, store.time + state.slot * spec.config.SECONDS_PER_SLOT)
-    run_on_block(spec, store, signed_block, False)
+    run_on_block(spec, store, signed_block, valid=False)
 
 
-@with_all_phases
+@with_phases([PHASE0])
 @spec_state_test
 def test_on_block_update_justified_checkpoint_within_safe_slots(spec, state):
     # Initialization
@@ -192,21 +94,33 @@ def test_on_block_update_justified_checkpoint_within_safe_slots(spec, state):
 
     next_epoch(spec, state)
     spec.on_tick(store, store.time + state.slot * spec.config.SECONDS_PER_SLOT)
-    state, store, last_signed_block = apply_next_epoch_with_attestations(spec, state, store)
+    state, store, last_signed_block = yield from apply_next_epoch_with_attestations(spec, state, store)
     next_epoch(spec, state)
     spec.on_tick(store, store.time + state.slot * spec.config.SECONDS_PER_SLOT)
-    last_block_root = hash_tree_root(last_signed_block.message)
+    last_block = last_signed_block.message
+    last_block_root = last_block.hash_tree_root()
 
-    # Mock the justified checkpoint
+    # NOTE: Mock the justified checkpoint
     just_state = store.block_states[last_block_root]
     new_justified = spec.Checkpoint(
         epoch=just_state.current_justified_checkpoint.epoch + 1,
         root=b'\x77' * 32,
     )
-    just_state.current_justified_checkpoint = new_justified
+    just_state.current_justified_checkpoint = new_justified  # Mutate `store`
+
+    assert store.block_states[last_block_root].hash_tree_root() == just_state.hash_tree_root()
 
     block = build_empty_block_for_next_slot(spec, just_state)
-    signed_block = state_transition_and_sign_block(spec, deepcopy(just_state), block)
+
+    # NOTE: Mock store so that the modified state could be accessed
+    parent_block = last_signed_block.message.copy()
+    parent_block.state_root = just_state.hash_tree_root()
+    store.blocks[block.parent_root] = parent_block
+    store.block_states[block.parent_root] = just_state.copy()
+    assert block.parent_root in store.blocks.keys()
+    assert block.parent_root in store.block_states.keys()
+
+    signed_block = state_transition_and_sign_block(spec, just_state.copy(), block)
     assert spec.get_current_slot(store) % spec.SLOTS_PER_EPOCH < spec.SAFE_SLOTS_TO_UPDATE_JUSTIFIED
     run_on_block(spec, store, signed_block)
 
@@ -223,10 +137,10 @@ def test_on_block_outside_safe_slots_and_multiple_better_justified(spec, state):
 
     next_epoch(spec, state)
     spec.on_tick(store, store.time + state.slot * spec.config.SECONDS_PER_SLOT)
-    state, store, last_signed_block = apply_next_epoch_with_attestations(spec, state, store)
+    state, store, last_signed_block = yield from apply_next_epoch_with_attestations(spec, state, store)
     last_block_root = hash_tree_root(last_signed_block.message)
 
-    # Mock fictitious justified checkpoint in store
+    # NOTE: Mock fictitious justified checkpoint in store
     store.justified_checkpoint = spec.Checkpoint(
         epoch=spec.compute_epoch_at_slot(last_signed_block.message.slot),
         root=spec.Root("0x4a55535449464945440000000000000000000000000000000000000000000000")
@@ -248,6 +162,7 @@ def test_on_block_outside_safe_slots_and_multiple_better_justified(spec, state):
     # Add a series of new blocks with "better" justifications
     best_justified_checkpoint = spec.Checkpoint(epoch=0)
     for i in range(3, 0, -1):
+        # Mutate store
         just_state = store.block_states[last_block_root]
         new_justified = spec.Checkpoint(
             epoch=previously_justified.epoch + i,
@@ -261,6 +176,14 @@ def test_on_block_outside_safe_slots_and_multiple_better_justified(spec, state):
         block = build_empty_block_for_next_slot(spec, just_state)
         signed_block = state_transition_and_sign_block(spec, deepcopy(just_state), block)
 
+        # NOTE: Mock store so that the modified state could be accessed
+        parent_block = store.blocks[last_block_root].copy()
+        parent_block.state_root = just_state.hash_tree_root()
+        store.blocks[block.parent_root] = parent_block
+        store.block_states[block.parent_root] = just_state.copy()
+        assert block.parent_root in store.blocks.keys()
+        assert block.parent_root in store.block_states.keys()
+
         run_on_block(spec, store, signed_block)
 
     assert store.justified_checkpoint == previously_justified
@@ -273,15 +196,15 @@ def test_on_block_outside_safe_slots_and_multiple_better_justified(spec, state):
 def test_on_block_outside_safe_slots_but_finality(spec, state):
     # Initialization
     store = get_genesis_forkchoice_store(spec, state)
-    time = 100
+    time = 0
     spec.on_tick(store, time)
 
     next_epoch(spec, state)
     spec.on_tick(store, store.time + state.slot * spec.config.SECONDS_PER_SLOT)
-    state, store, last_signed_block = apply_next_epoch_with_attestations(spec, state, store)
+    state, store, last_signed_block = yield from apply_next_epoch_with_attestations(spec, state, store)
     last_block_root = hash_tree_root(last_signed_block.message)
 
-    # Mock fictitious justified checkpoint in store
+    # NOTE: Mock fictitious justified checkpoint in store
     store.justified_checkpoint = spec.Checkpoint(
         epoch=spec.compute_epoch_at_slot(last_signed_block.message.slot),
         root=spec.Root("0x4a55535449464945440000000000000000000000000000000000000000000000")
@@ -290,7 +213,7 @@ def test_on_block_outside_safe_slots_but_finality(spec, state):
     next_epoch(spec, state)
     spec.on_tick(store, store.time + state.slot * spec.config.SECONDS_PER_SLOT)
 
-    # Create new higher justified checkpoint not in branch of store's justified checkpoint
+    # NOTE: Mock a new higher justified checkpoint not in branch of store's justified checkpoint
     just_block = build_empty_block_for_next_slot(spec, state)
     store.blocks[just_block.hash_tree_root()] = just_block
 
@@ -298,7 +221,7 @@ def test_on_block_outside_safe_slots_but_finality(spec, state):
     spec.on_tick(store, store.time + spec.SAFE_SLOTS_TO_UPDATE_JUSTIFIED * spec.config.SECONDS_PER_SLOT)
     assert spec.get_current_slot(store) % spec.SLOTS_PER_EPOCH >= spec.SAFE_SLOTS_TO_UPDATE_JUSTIFIED
 
-    # Mock justified and finalized update in state
+    # NOTE: Mock justified and finalized update in state
     just_fin_state = store.block_states[last_block_root]
     new_justified = spec.Checkpoint(
         epoch=spec.compute_epoch_at_slot(just_block.slot) + 1,
@@ -317,6 +240,14 @@ def test_on_block_outside_safe_slots_but_finality(spec, state):
     block = build_empty_block_for_next_slot(spec, just_fin_state)
     signed_block = state_transition_and_sign_block(spec, deepcopy(just_fin_state), block)
 
+    # NOTE: Mock store so that the modified state could be accessed
+    parent_block = last_signed_block.message.copy()
+    parent_block.state_root = just_fin_state.hash_tree_root()
+    store.blocks[block.parent_root] = parent_block
+    store.block_states[block.parent_root] = just_fin_state.copy()
+    assert block.parent_root in store.blocks.keys()
+    assert block.parent_root in store.block_states.keys()
+
     run_on_block(spec, store, signed_block)
 
     assert store.finalized_checkpoint == new_finalized

tests/generators/fork_choice/main.py

@@ -4,7 +4,7 @@ from eth2spec.test.helpers.constants import PHASE0, ALTAIR, MERGE
 
 if __name__ == "__main__":
     phase_0_mods = {key: 'eth2spec.test.phase0.fork_choice.test_' + key for key in [
-        # 'get_head',
+        'get_head',
         'on_block',
     ]}
     # No additional Altair specific finality tests, yet.