From 83760b15ac1dc7833a77e03559f666cc75945b19 Mon Sep 17 00:00:00 2001
From: Hsiao-Wei Wang <hwwang156@gmail.com>
Date: Mon, 10 Aug 2020 20:54:22 +0800
Subject: [PATCH] Use v-guide's `get_custody_secret`

---
 tests/core/pyspec/eth2spec/test/helpers/custody.py       | 9 ---------
 .../block_processing/test_process_custody_slashing.py    | 9 +++++++--
 .../pyspec/eth2spec/test/phase1/sanity/test_blocks.py    | 9 +++++++--
 3 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/tests/core/pyspec/eth2spec/test/helpers/custody.py b/tests/core/pyspec/eth2spec/test/helpers/custody.py
index 162ca3263..b3a8c0a95 100644
--- a/tests/core/pyspec/eth2spec/test/helpers/custody.py
+++ b/tests/core/pyspec/eth2spec/test/helpers/custody.py
@@ -179,15 +179,6 @@ def get_sample_shard_transition(spec, start_slot, block_lengths):
     return shard_transition
 
 
-def get_custody_secret(spec, state, validator_index, epoch=None):
-    period = spec.get_custody_period_for_validator(validator_index, epoch if epoch is not None
-                                                   else spec.get_current_epoch(state))
-    epoch_to_sign = spec.get_randao_epoch_for_custody_period(period, validator_index)
-    domain = spec.get_domain(state, spec.DOMAIN_RANDAO, epoch_to_sign)
-    signing_root = spec.compute_signing_root(spec.Epoch(epoch_to_sign), domain)
-    return bls.Sign(privkeys[validator_index], signing_root)
-
-
 def get_custody_slashable_test_vector(spec, custody_secret, length, slashable=True):
     test_vector = get_custody_test_vector(length)
     offset = 0
diff --git a/tests/core/pyspec/eth2spec/test/phase1/block_processing/test_process_custody_slashing.py b/tests/core/pyspec/eth2spec/test/phase1/block_processing/test_process_custody_slashing.py
index b08b65413..f8c92cb44 100644
--- a/tests/core/pyspec/eth2spec/test/phase1/block_processing/test_process_custody_slashing.py
+++ b/tests/core/pyspec/eth2spec/test/phase1/block_processing/test_process_custody_slashing.py
@@ -1,11 +1,11 @@
 from eth2spec.test.helpers.custody import (
     get_valid_custody_slashing,
-    get_custody_secret,
     get_custody_slashable_shard_transition,
 )
 from eth2spec.test.helpers.attestations import (
     get_valid_on_time_attestation,
 )
+from eth2spec.test.helpers.keys import privkeys
 from eth2spec.utils.ssz.ssz_typing import ByteList
 from eth2spec.test.helpers.state import get_balance, transition_to
 from eth2spec.test.context import (
@@ -77,7 +77,12 @@ def run_standard_custody_slashing_test(spec,
     if block_lengths is None:
         block_lengths = [2**15 // 3] * len(offset_slots)
 
-    custody_secret = get_custody_secret(spec, state, validator_index)
+    custody_secret = spec.get_custody_secret(
+        state,
+        validator_index,
+        privkeys[validator_index],
+        spec.get_current_epoch(state),
+    )
     shard_transition, slashable_test_vector = get_custody_slashable_shard_transition(
         spec,
         state.slot,
diff --git a/tests/core/pyspec/eth2spec/test/phase1/sanity/test_blocks.py b/tests/core/pyspec/eth2spec/test/phase1/sanity/test_blocks.py
index 5fe3394db..1ddbf1c6b 100644
--- a/tests/core/pyspec/eth2spec/test/phase1/sanity/test_blocks.py
+++ b/tests/core/pyspec/eth2spec/test/phase1/sanity/test_blocks.py
@@ -8,7 +8,6 @@ from eth2spec.test.context import (
 from eth2spec.test.helpers.attestations import get_valid_on_time_attestation
 from eth2spec.test.helpers.block import build_empty_block
 from eth2spec.test.helpers.custody import (
-    get_custody_secret,
     get_custody_slashable_test_vector,
     get_valid_chunk_challenge,
     get_valid_custody_chunk_response,
@@ -16,6 +15,7 @@ from eth2spec.test.helpers.custody import (
     get_valid_custody_slashing,
     get_valid_early_derived_secret_reveal,
 )
+from eth2spec.test.helpers.keys import privkeys
 from eth2spec.test.helpers.shard_block import (
     build_shard_block,
     get_committee_index_of_shard,
@@ -215,7 +215,12 @@ def test_custody_slashing(spec, state):
     committee_index = get_committee_index_of_shard(spec, state, state.slot, shard)
     # Create slashable shard block body
     validator_index = spec.get_beacon_committee(state, state.slot, committee_index)[0]
-    custody_secret = get_custody_secret(spec, state, validator_index)
+    custody_secret = spec.get_custody_secret(
+        state,
+        validator_index,
+        privkeys[validator_index],
+        spec.get_current_epoch(state),
+    )
     slashable_body = get_custody_slashable_test_vector(spec, custody_secret, length=100, slashable=True)
     shard_block = build_shard_block(spec, state, shard, body=slashable_body, slot=state.slot, signed=True)
     shard_block_dict: Dict[spec.Shard, Sequence[spec.SignedShardBlock]] = {shard: [shard_block]}