Fix test_aggregate_votes test and add verify_votes function

Author: hwwhww

eth/beacon/block_committees_info.py

@@ -1,6 +1,6 @@
 from typing import (
-    Iterable,
     NamedTuple,
+    Tuple,
     TYPE_CHECKING,
 )
 
@@ -15,6 +15,6 @@
         ('proposer_index_in_committee', int),
         ('proposer_shard_id', int),
         ('proposer_committee_size', int),
-        ('shards_and_committees', Iterable['ShardAndCommittee'])
+        ('shards_and_committees', Tuple['ShardAndCommittee'])
     )
 )

eth/beacon/helpers.py

@@ -439,11 +439,16 @@ def aggregate_attestation_record(last_justified_slot: int,
         last_justified_slot,
     )
 
+    # Verify
+    # TODO: Verification is very slow, needs to compute in parallel
+    voting_sigs, committee_indices = verify_votes(message, votes)
+
+    # Aggregate the votes
     bitfield, sigs = aggregate_votes(
-        message,
-        votes,
-        proposer_attestation.bitfield,
-        proposer_attestation.sigs,
+        bitfield=proposer_attestation.bitfield,
+        sigs=proposer_attestation.aggregate_sig,
+        voting_sigs=voting_sigs,
+        voting_committee_indices=committee_indices
     )
 
     return proposer_attestation.copy(
@@ -452,11 +457,12 @@ def aggregate_attestation_record(last_justified_slot: int,
     )
 
 
-def aggregate_votes(message: bytes,
-                    votes: Iterable[Tuple[int, bytes, int]],
-                    pre_bitfield: bytes,
-                    pre_sigs: Iterable[bytes]) -> Tuple[bytes, Iterable[int]]:
-    # Update the bitfield and append the signatures
+def verify_votes(
+        message: bytes,
+        votes: Iterable[Tuple[int, bytes, int]]) -> Tuple[Tuple[bytes, ...], Tuple[int, ...]]:
+    """
+    Verify the given votes
+    """
     sigs_with_committe_info = tuple(
         (sig, committee_index)
         for (committee_index, sig, public_key)
@@ -466,16 +472,26 @@ def aggregate_votes(message: bytes,
     try:
         sigs, committee_indices = zip(*sigs_with_committe_info)
     except ValueError:
-        sigs = ()
-        committee_indices = ()
+        sigs = tuple()
+        committee_indices = tuple()
+
+    return sigs, committee_indices
 
-    sigs = sigs + tuple(pre_sigs)
-    bitfield = pre_bitfield
+
+def aggregate_votes(bitfield: bytes,
+                    sigs: Iterable[bytes],
+                    voting_sigs: Iterable[bytes],
+                    voting_committee_indices: Iterable[int]) -> Tuple[bytes, Tuple[int]]:
+    """
+    Aggregate the votes
+    """
+    # Update the bitfield and append the signatures
+    sigs = tuple(sigs) + tuple(voting_sigs)
     bitfield = pipe(
         bitfield,
         *(
             set_voted(index=committee_index)
-            for committee_index in committee_indices
+            for committee_index in voting_committee_indices
         )
     )
 

tests/beacon/test_helpers.py

@@ -1,6 +1,9 @@
-import random
-
 import pytest
+from hypothesis import (
+    given,
+    settings,
+    strategies as st,
+)
 
 from eth.utils import bls
 from eth.utils.bitfield import (
@@ -20,6 +23,7 @@
     get_shards_and_committees_for_slot,
     get_signed_parent_hashes,
     get_block_committees_info,
+    verify_votes,
 )
 
 from tests.beacon.helpers import (
@@ -398,43 +402,53 @@ def mock_get_shards_and_committees_for_slot(parent_block,
         )
 
 
+@settings(max_examples=1)
+@given(random=st.randoms())
 @pytest.mark.parametrize(
     (
         'votes_count'
     ),
     [
         (0),
-        (1),
-        (5),
         (9),
     ],
 )
-def test_aggregate_votes(votes_count, privkeys, pubkeys):
+def test_aggregate_votes(votes_count, random, privkeys, pubkeys):
     bit_count = 10
     pre_bitfield = get_empty_bitfield(bit_count)
     pre_sigs = ()
 
     random_votes = random.sample([i for i in range(bit_count)], votes_count)
     message = b'hello'
 
-    # (committee_index, sig, public_key)
+    # Get votes: (committee_index, sig, public_key)
     votes = [
         (committee_index, bls.sign(message, privkeys[committee_index]), pubkeys[committee_index])
         for committee_index in random_votes
     ]
+
+    # Verify
+    sigs, committee_indices = verify_votes(message, votes)
+
+    # Aggregate the votes
     bitfield, sigs = aggregate_votes(
-        message,
-        votes,
-        pre_bitfield,
-        pre_sigs,
+        bitfield=pre_bitfield,
+        sigs=pre_sigs,
+        voting_sigs=sigs,
+        voting_committee_indices=committee_indices
     )
 
     try:
         _, _, pubs = zip(*votes)
     except ValueError:
         pubs = ()
 
+    voted_index = [
+        committee_index
+        for committee_index in random_votes
+        if has_voted(bitfield, committee_index)
+    ]
+    assert len(voted_index) == len(votes)
+
     aggregated_pubs = bls.aggregate_pubs(pubs)
-    for committee_index in random_votes:
-        if has_voted(bitfield, committee_index):
-            assert bls.verify(message, aggregated_pubs, sigs)
+    assert bls.verify(message, aggregated_pubs, sigs)