Fix and refactor tick_and_run_on_attestation

ethereum · Mar 12, 2021 · 09fc717 · 09fc717
1 parent eebe081
commit 09fc717
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/tests/core/pyspec/eth2spec/test/helpers/fork_choice.py b/tests/core/pyspec/eth2spec/test/helpers/fork_choice.py
@@ -33,7 +33,7 @@ def tick_and_run_on_block(spec, store, signed_block, test_steps=None):
     yield from run_on_block(spec, store, signed_block, test_steps)
 
 
-def add_attestation_to_store(spec, store, attestation, test_steps=None):
+def tick_and_run_on_attestation(spec, store, attestation, test_steps=None):
     if test_steps is None:
         test_steps = []
 
@@ -47,6 +47,7 @@ def add_attestation_to_store(spec, store, attestation, test_steps=None):
         test_steps.append({'tick': int(next_epoch_time)})
 
     spec.on_attestation(store, attestation)
+    yield get_attestation_file_name(attestation), attestation
     test_steps.append({'attestation': get_attestation_file_name(attestation)})
 
 

diff --git a/tests/core/pyspec/eth2spec/test/phase0/fork_choice/test_get_head.py b/tests/core/pyspec/eth2spec/test/phase0/fork_choice/test_get_head.py
@@ -10,7 +10,7 @@
 from eth2spec.test.helpers.attestations import get_valid_attestation, next_epoch_with_attestations
 from eth2spec.test.helpers.block import build_empty_block_for_next_slot
 from eth2spec.test.helpers.fork_choice import (
-    add_attestation_to_store,
+    tick_and_run_on_attestation,
     tick_and_run_on_block, get_anchor_root,
     get_genesis_forkchoice_store_and_block,
     get_attestation_file_name,
@@ -160,7 +160,7 @@ def test_shorter_chain_but_heavier_weight(spec, state):
     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)
+    yield from tick_and_run_on_attestation(spec, store, short_attestation, test_steps)
 
     assert spec.get_head(store) == spec.hash_tree_root(short_block)
     test_steps.append({
@@ -250,9 +250,7 @@ def test_filtered_block_tree(spec, state):
     yield from run_on_block(spec, store, signed_rogue_block, test_steps)
 
     for attestation in attestations:
-        spec.on_attestation(store, attestation)
-        test_steps.append({'attestation': get_attestation_file_name(attestation)})
-        yield get_attestation_file_name(attestation), attestation
+        yield from tick_and_run_on_attestation(spec, store, attestation, test_steps)
 
     # ensure that get_head still returns the head from the previous branch
     assert spec.get_head(store) == expected_head_root