Refactor test_sync_protocol.py a bit

tests/core/pyspec/eth2spec/test/altair/unittests/test_sync_protocol.py

@@ -5,38 +5,29 @@ from eth2spec.test.context import (
     with_presets,
     with_altair_and_later,
 )
-from eth2spec.test.helpers.attestations import next_epoch_with_attestations
+from eth2spec.test.helpers.attestations import (
+    next_epoch_with_attestations,
+)
 from eth2spec.test.helpers.block import (
     build_empty_block,
     build_empty_block_for_next_slot,
 )
 from eth2spec.test.helpers.constants import MINIMAL
+from eth2spec.test.helpers.light_client import (
+    get_sync_aggregate,
+    initialize_light_client_store,
+)
 from eth2spec.test.helpers.state import (
     next_slots,
     state_transition_and_sign_block,
 )
-from eth2spec.test.helpers.sync_committee import (
-    compute_aggregate_sync_committee_signature,
-)
 from eth2spec.test.helpers.merkle import build_proof
 
 
-def _initialize_light_client_store(spec, state):
-    return spec.LightClientStore(
-        finalized_header=spec.BeaconBlockHeader(),
-        current_sync_committee=state.current_sync_committee,
-        next_sync_committee=state.next_sync_committee,
-        best_valid_update=None,
-        optimistic_header=spec.BeaconBlockHeader(),
-        previous_max_active_participants=0,
-        current_max_active_participants=0,
-    )
-
-
 @with_altair_and_later
 @spec_state_test
 def test_process_light_client_update_not_timeout(spec, state):
-    store = _initialize_light_client_store(spec, state)
+    store = initialize_light_client_store(spec, state)
 
     # Block at slot 1 doesn't increase sync committee period, so it won't force update store.finalized_header
     block = build_empty_block_for_next_slot(spec, state)
@@ -49,19 +40,7 @@ def test_process_light_client_update_not_timeout(spec, state):
         body_root=signed_block.message.body.hash_tree_root(),
     )
     # Sync committee signing the header
-    all_pubkeys = [v.pubkey for v in state.validators]
-    committee = [all_pubkeys.index(pubkey) for pubkey in state.current_sync_committee.pubkeys]
-    sync_committee_bits = [True] * len(committee)
-    sync_committee_signature = compute_aggregate_sync_committee_signature(
-        spec,
-        state,
-        block_header.slot,
-        committee,
-    )
-    sync_aggregate = spec.SyncAggregate(
-        sync_committee_bits=sync_committee_bits,
-        sync_committee_signature=sync_committee_signature,
-    )
+    sync_aggregate = get_sync_aggregate(spec, state, block_header, block_root=None)
     next_sync_committee_branch = [spec.Bytes32() for _ in range(spec.floorlog2(spec.NEXT_SYNC_COMMITTEE_INDEX))]
 
     # Ensure that finality checkpoint is genesis
@@ -94,7 +73,7 @@ def test_process_light_client_update_not_timeout(spec, state):
 @spec_state_test
 @with_presets([MINIMAL], reason="too slow")
 def test_process_light_client_update_timeout(spec, state):
-    store = _initialize_light_client_store(spec, state)
+    store = initialize_light_client_store(spec, state)
 
     # Forward to next sync committee period
     next_slots(spec, state, spec.UPDATE_TIMEOUT)
@@ -113,20 +92,8 @@ def test_process_light_client_update_timeout(spec, state):
     )
 
     # Sync committee signing the finalized_block_header
-    all_pubkeys = [v.pubkey for v in state.validators]
-    committee = [all_pubkeys.index(pubkey) for pubkey in state.current_sync_committee.pubkeys]
-    sync_committee_bits = [True] * len(committee)
-    sync_committee_signature = compute_aggregate_sync_committee_signature(
-        spec,
-        state,
-        block_header.slot,
-        committee,
-        block_root=spec.Root(block_header.hash_tree_root()),
-    )
-    sync_aggregate = spec.SyncAggregate(
-        sync_committee_bits=sync_committee_bits,
-        sync_committee_signature=sync_committee_signature,
-    )
+    sync_aggregate = get_sync_aggregate(
+        spec, state, block_header, block_root=spec.Root(block_header.hash_tree_root()))
 
     # Sync committee is updated
     next_sync_committee_branch = build_proof(state.get_backing(), spec.NEXT_SYNC_COMMITTEE_INDEX)
@@ -158,7 +125,7 @@ def test_process_light_client_update_timeout(spec, state):
 @spec_state_test
 @with_presets([MINIMAL], reason="too slow")
 def test_process_light_client_update_finality_updated(spec, state):
-    store = _initialize_light_client_store(spec, state)
+    store = initialize_light_client_store(spec, state)
 
     # Change finality
     blocks = []
@@ -191,20 +158,8 @@ def test_process_light_client_update_finality_updated(spec, state):
     )
 
     # Sync committee signing the finalized_block_header
-    all_pubkeys = [v.pubkey for v in state.validators]
-    committee = [all_pubkeys.index(pubkey) for pubkey in state.current_sync_committee.pubkeys]
-    sync_committee_bits = [True] * len(committee)
-    sync_committee_signature = compute_aggregate_sync_committee_signature(
-        spec,
-        state,
-        block_header.slot,
-        committee,
-        block_root=spec.Root(block_header.hash_tree_root()),
-    )
-    sync_aggregate = spec.SyncAggregate(
-        sync_committee_bits=sync_committee_bits,
-        sync_committee_signature=sync_committee_signature,
-    )
+    sync_aggregate = get_sync_aggregate(
+        spec, state, block_header, block_root=spec.Root(block_header.hash_tree_root()))
 
     update = spec.LightClientUpdate(
         attested_header=block_header,

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

@@ -0,0 +1,35 @@
+from eth2spec.test.helpers.sync_committee import (
+    compute_aggregate_sync_committee_signature,
+)
+
+
+def initialize_light_client_store(spec, state):
+    return spec.LightClientStore(
+        finalized_header=spec.BeaconBlockHeader(),
+        current_sync_committee=state.current_sync_committee,
+        next_sync_committee=state.next_sync_committee,
+        best_valid_update=None,
+        optimistic_header=spec.BeaconBlockHeader(),
+        previous_max_active_participants=0,
+        current_max_active_participants=0,
+    )
+
+
+def get_sync_aggregate(spec, state, block_header, block_root=None, signature_slot=None):
+    if signature_slot is None:
+        signature_slot = block_header.slot
+
+    all_pubkeys = [v.pubkey for v in state.validators]
+    committee = [all_pubkeys.index(pubkey) for pubkey in state.current_sync_committee.pubkeys]
+    sync_committee_bits = [True] * len(committee)
+    sync_committee_signature = compute_aggregate_sync_committee_signature(
+        spec,
+        state,
+        block_header.slot,
+        committee,
+        block_root=block_root,
+    )
+    return spec.SyncAggregate(
+        sync_committee_bits=sync_committee_bits,
+        sync_committee_signature=sync_committee_signature,
+    )