From 95c712598af0bdb234c577e4da859cf2f8043ad7 Mon Sep 17 00:00:00 2001
From: Hsiao-Wei Wang <hsiaowei.eth@gmail.com>
Date: Wed, 9 Mar 2022 10:07:01 +0800
Subject: [PATCH] Fix `test_discard_equivocations` test vectors

---
 .../eth2spec/test/helpers/fork_choice.py      | 40 +++++++++++++++++++
 .../test/phase0/fork_choice/test_get_head.py  | 17 ++++----
 2 files changed, 50 insertions(+), 7 deletions(-)

diff --git a/tests/core/pyspec/eth2spec/test/helpers/fork_choice.py b/tests/core/pyspec/eth2spec/test/helpers/fork_choice.py
index f056b9acd..d524060a2 100644
--- a/tests/core/pyspec/eth2spec/test/helpers/fork_choice.py
+++ b/tests/core/pyspec/eth2spec/test/helpers/fork_choice.py
@@ -92,6 +92,10 @@ def get_attestation_file_name(attestation):
     return f"attestation_{encode_hex(attestation.hash_tree_root())}"
 
 
+def get_attester_slashing_file_name(attester_slashing):
+    return f"attester_slashing_{encode_hex(attester_slashing.hash_tree_root())}"
+
+
 def on_tick_and_append_step(spec, store, time, test_steps):
     spec.on_tick(store, time)
     test_steps.append({'tick': int(time)})
@@ -142,6 +146,10 @@ def add_block(spec,
     for attestation in signed_block.message.body.attestations:
         run_on_attestation(spec, store, attestation, is_from_block=True, valid=True)
 
+    # An on_block step implies receiving block's attester slashings
+    for attester_slashing in signed_block.message.body.attester_slashings:
+        run_on_attester_slashing(spec, store, attester_slashing, valid=True)
+
     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
@@ -168,6 +176,38 @@ def add_block(spec,
     return store.block_states[signed_block.message.hash_tree_root()]
 
 
+def run_on_attester_slashing(spec, store, attester_slashing, valid=True):
+    if not valid:
+        try:
+            spec.on_attester_slashing(store, attester_slashing)
+        except AssertionError:
+            return
+        else:
+            assert False
+
+    spec.on_attester_slashing(store, attester_slashing)
+
+
+def add_attester_slashing(spec, store, attester_slashing, test_steps, valid=True):
+    slashing_file_name = get_attester_slashing_file_name(attester_slashing)
+    yield get_attester_slashing_file_name(attester_slashing), attester_slashing
+
+    if not valid:
+        try:
+            run_on_attester_slashing(spec, store, attester_slashing)
+        except AssertionError:
+            test_steps.append({
+                'attester_slashing': slashing_file_name,
+                'valid': False,
+            })
+            return
+        else:
+            assert False
+
+    run_on_attester_slashing(spec, store, attester_slashing)
+    test_steps.append({'attester_slashing': slashing_file_name})
+
+
 def get_formatted_head_output(spec, store):
     head = spec.get_head(store)
     slot = store.blocks[head].slot
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
index 1f37625a8..587a1c6e8 100644
--- 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
@@ -8,24 +8,27 @@ from eth2spec.test.context import (
     with_presets,
 )
 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.block import (
+    apply_empty_block,
+    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_add_block,
+    add_attester_slashing,
+    add_block,
     get_anchor_root,
     get_genesis_forkchoice_store_and_block,
     get_formatted_head_output,
     on_tick_and_append_step,
-    add_block,
+    run_on_attestation,
+    tick_and_run_on_attestation,
+    tick_and_add_block,
 )
 from eth2spec.test.helpers.state import (
     next_slots,
     next_epoch,
     state_transition_and_sign_block,
 )
-from tests.core.pyspec.eth2spec.test.helpers.block import apply_empty_block
-from tests.core.pyspec.eth2spec.test.helpers.fork_choice import run_on_attestation
 
 
 rng = random.Random(1001)
@@ -411,7 +414,7 @@ def test_discard_equivocations(spec, state):
 
     # Process attester_slashing
     # The head should revert to block_2
-    spec.on_attester_slashing(store, attester_slashing)
+    yield from add_attester_slashing(spec, store, attester_slashing, test_steps)
     assert spec.get_head(store) == spec.hash_tree_root(block_2)
 
     test_steps.append({