Fix after rectoring + use more run_on_block

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

@@ -10,7 +10,17 @@ def get_anchor_root(spec, state):
     return spec.hash_tree_root(anchor_block_header)
 
 
-def add_block_to_store(spec, store, signed_block, test_steps=None):
+def add_block_to_store(spec, store, signed_block):
+    pre_state = store.block_states[signed_block.message.parent_root]
+    block_time = pre_state.genesis_time + signed_block.message.slot * spec.SECONDS_PER_SLOT
+
+    if store.time < block_time:
+        spec.on_tick(store, block_time)
+
+    spec.on_block(store, signed_block)
+
+
+def tick_and_run_on_block(spec, store, signed_block, test_steps=None):
     if test_steps is None:
         test_steps = []
 
@@ -20,8 +30,7 @@ def add_block_to_store(spec, store, signed_block, test_steps=None):
     if store.time < block_time:
         on_tick_and_append_step(spec, store, block_time, test_steps)
 
-    spec.on_block(store, signed_block)
-    test_steps.append({'block': get_block_file_name(signed_block)})
+    yield from run_on_block(spec, store, signed_block, test_steps)
 
 
 def add_attestation_to_store(spec, store, attestation, test_steps=None):
@@ -77,13 +86,13 @@ def run_on_block(spec, store, signed_block, test_steps, valid=True):
             assert False
 
     spec.on_block(store, signed_block)
+    yield get_block_file_name(signed_block), signed_block
+    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)
 
-    test_steps.append({'block': get_block_file_name(signed_block)})
-
     assert store.blocks[signed_block.message.hash_tree_root()] == signed_block.message
     test_steps.append({
         'checks': {

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

@@ -11,10 +11,9 @@
 from eth2spec.test.helpers.block import build_empty_block_for_next_slot
 from eth2spec.test.helpers.fork_choice import (
     add_attestation_to_store,
-    add_block_to_store, get_anchor_root,
+    tick_and_run_on_block, get_anchor_root,
     get_genesis_forkchoice_store_and_block,
     get_attestation_file_name,
-    get_block_file_name,
     get_formatted_head_output,
     on_tick_and_append_step,
     run_on_block,
@@ -69,14 +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)
-    add_block_to_store(spec, store, signed_block_1, test_steps)
-    yield get_block_file_name(signed_block_1), signed_block_1
+    yield from tick_and_run_on_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)
-    add_block_to_store(spec, store, signed_block_2, test_steps)
-    yield get_block_file_name(signed_block_2), signed_block_2
+    yield from tick_and_run_on_block(spec, store, signed_block_2, test_steps)
 
     assert spec.get_head(store) == spec.hash_tree_root(block_2)
     test_steps.append({
@@ -110,16 +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)
-    add_block_to_store(spec, store, signed_block_1, test_steps)
-    yield get_block_file_name(signed_block_1), signed_block_1
+    yield from tick_and_run_on_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)
-    add_block_to_store(spec, store, signed_block_2, test_steps)
-    yield get_block_file_name(signed_block_2), signed_block_2
+    yield from tick_and_run_on_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
@@ -155,16 +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)
-        add_block_to_store(spec, store, signed_long_block, test_steps)
-        yield get_block_file_name(signed_long_block), signed_long_block
+        yield from tick_and_run_on_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)
-    add_block_to_store(spec, store, signed_short_block, test_steps)
-    yield get_block_file_name(signed_short_block), signed_short_block
+    yield from tick_and_run_on_block(spec, store, signed_short_block, test_steps)
 
     short_attestation = get_valid_attestation(spec, short_state, short_block.slot, signed=True)
     add_attestation_to_store(spec, store, short_attestation, test_steps)
@@ -207,7 +200,7 @@ def test_filtered_block_tree(spec, state):
     current_time = state.slot * spec.SECONDS_PER_SLOT + store.genesis_time
     on_tick_and_append_step(spec, store, current_time, test_steps)
     for signed_block in signed_blocks:
-        run_on_block(spec, store, signed_block, test_steps)
+        yield from run_on_block(spec, store, signed_block, test_steps)
 
     assert store.justified_checkpoint == state.current_justified_checkpoint
 
@@ -254,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
-    run_on_block(spec, store, signed_rogue_block, test_steps)
+    yield from run_on_block(spec, store, signed_rogue_block, test_steps)
 
     for attestation in attestations:
         spec.on_attestation(store, attestation)