Fix get_shard_committees_at_slot

Author: hwwhww

eth/beacon/helpers.py

@@ -170,30 +170,42 @@ def get_new_recent_block_hashes(old_block_hashes: Sequence[Hash32],
 # Get shards_committees or indices
 #
 @to_tuple
-def get_shards_committees_for_slot(
-        crystallized_state: 'CrystallizedState',
+def _get_shard_committees_at_slot(
+        latest_state_recalculation_slot: int,
+        shard_committees_at_slots: Sequence[Sequence[ShardCommittee]],
         slot: int,
         epoch_length: int) -> Iterable[ShardCommittee]:
-    """
-    FIXME
-    """
-    if len(crystallized_state.shard_committee_for_slots) != epoch_length * 2:
+    if len(shard_committees_at_slots) != epoch_length * 2:
         raise ValueError(
-            "Length of shard_committee_for_slots != epoch_length * 2"
+            "Length of shard_committees_at_slots != epoch_length * 2"
             "\texpected: %s, found: %s" % (
-                epoch_length * 2, len(crystallized_state.shard_committee_for_slots)
+                epoch_length * 2, len(shard_committees_at_slots)
             )
         )
 
-    slot_relative_position = crystallized_state.last_state_recalc - epoch_length
+    slot_relative_position = latest_state_recalculation_slot - epoch_length
 
     yield from _get_element_from_recent_list(
-        crystallized_state.shard_committee_for_slots,
+        shard_committees_at_slots,
         slot,
         slot_relative_position,
     )
 
 
+def get_shard_committees_at_slot(state: 'BeaconState',
+                                 slot: int,
+                                 epoch_length: int) -> Tuple[ShardCommittee]:
+    """
+    Returns the ``ShardCommittee`` for the ``slot``.
+    """
+    return _get_shard_committees_at_slot(
+        latest_state_recalculation_slot=state.latest_state_recalculation_slot,
+        shard_committees_at_slots=state.shard_committees_at_slots,
+        slot=slot,
+        epoch_length=epoch_length,
+    )
+
+
 @to_tuple
 def get_attestation_indices(crystallized_state: 'CrystallizedState',
                             attestation: 'AttestationRecord',
@@ -204,7 +216,7 @@ def get_attestation_indices(crystallized_state: 'CrystallizedState',
     """
     shard_id = attestation.shard_id
 
-    shards_committees_for_slot = get_shards_committees_for_slot(
+    shards_committees_for_slot = get_shard_committees_at_slot(
         crystallized_state,
         attestation.slot,
         epoch_length,
@@ -325,7 +337,7 @@ def get_new_shuffling(*,
 def get_block_committees_info(parent_block: 'BaseBeaconBlock',
                               crystallized_state: 'CrystallizedState',
                               epoch_length: int) -> BlockCommitteesInfo:
-    shards_committees = get_shards_committees_for_slot(
+    shards_committees = get_shard_committees_at_slot(
         crystallized_state,
         parent_block.slot_number,
         epoch_length,

tests/beacon/test_helpers.py

@@ -14,7 +14,7 @@
     get_hashes_from_recent_block_hashes,
     get_hashes_to_sign,
     get_new_shuffling,
-    get_shards_committees_for_slot,
+    _get_shard_committees_at_slot,
     get_signed_parent_hashes,
     get_block_committees_info,
 )
@@ -24,6 +24,21 @@
 )
 
 
+def get_sample_shard_committees_at_slots(num_slot,
+                                         num_shard_committee_per_slot,
+                                         sample_shard_committee_params):
+
+    return tuple(
+        [
+            [
+                ShardCommittee(**sample_shard_committee_params)
+                for _ in range(num_shard_committee_per_slot)
+            ]
+            for _ in range(num_slot)
+        ]
+    )
+
+
 def generate_mock_recent_block_hashes(
         genesis_block,
         current_block_number,
@@ -202,39 +217,100 @@ def test_get_new_recent_block_hashes(genesis_block,
 #
 # Get shards_committees or indices
 #
-@pytest.mark.xfail(reason="Need to be fixed")
 @pytest.mark.parametrize(
     (
-        'num_validators,slot,success'
+        'num_validators,'
+        'cycle_length,'
+        'latest_state_recalculation_slot,'
+        'num_slot,'
+        'num_shard_committee_per_slot,'
+        'slot,'
+        'success'
     ),
     [
-        (100, 0, True),
-        (100, 63, True),
-        (100, 64, False),
+        (
+            100,
+            64,
+            0,
+            128,
+            10,
+            0,
+            True,
+        ),
+        (
+            100,
+            64,
+            64,
+            128,
+            10,
+            64,
+            True,
+        ),
+        # The length of shard_committees_at_slots != epoch_length * 2
+        (
+            100,
+            64,
+            64,
+            127,
+            10,
+            0,
+            False,
+        ),
+        # slot is too small
+        (
+            100,
+            64,
+            128,
+            128,
+            10,
+            0,
+            False,
+        ),
+        # slot is too large
+        (
+            100,
+            64,
+            0,
+            128,
+            10,
+            64,
+            False,
+        ),
     ],
 )
 def test_get_shard_committee_for_slot(
-        genesis_crystallized_state,
         num_validators,
+        cycle_length,
+        latest_state_recalculation_slot,
+        num_slot,
+        num_shard_committee_per_slot,
         slot,
         success,
-        epoch_length):
-    crystallized_state = genesis_crystallized_state
+        epoch_length,
+        sample_shard_committee_params):
+
+    shard_committees_at_slots = get_sample_shard_committees_at_slots(
+        num_slot,
+        num_shard_committee_per_slot,
+        sample_shard_committee_params
+    )
 
     if success:
-        shards_committees_for_slot = get_shards_committees_for_slot(
-            crystallized_state,
-            slot,
-            epoch_length,
+        shard_committees = _get_shard_committees_at_slot(
+            latest_state_recalculation_slot=latest_state_recalculation_slot,
+            shard_committees_at_slots=shard_committees_at_slots,
+            slot=slot,
+            epoch_length=epoch_length,
         )
-        assert len(shards_committees_for_slot) > 0
-        assert len(shards_committees_for_slot[0].committee) > 0
+        assert len(shard_committees) > 0
+        assert len(shard_committees[0].committee) > 0
     else:
         with pytest.raises(ValueError):
-            get_shards_committees_for_slot(
-                crystallized_state,
-                slot,
-                epoch_length,
+            _get_shard_committees_at_slot(
+                latest_state_recalculation_slot=latest_state_recalculation_slot,
+                shard_committees_at_slots=shard_committees_at_slots,
+                slot=slot,
+                epoch_length=epoch_length,
             )
 
 
@@ -363,17 +439,17 @@ def test_get_block_committees_info(monkeypatch,
                                    epoch_length):
     from eth.beacon import helpers
 
-    def mock_get_shards_committees_for_slot(parent_block,
-                                            crystallized_state,
-                                            epoch_length):
+    def mock_get_shard_committees_at_slot(parent_block,
+                                          crystallized_state,
+                                          epoch_length):
         return [
             ShardCommittee(shard_id=1, committee=committee),
         ]
 
     monkeypatch.setattr(
         helpers,
-        'get_shards_committees_for_slot',
-        mock_get_shards_committees_for_slot
+        '_get_shard_committees_at_slot',
+        mock__get_shard_committees_at_slot
     )
 
     parent_block = genesis_block