Refactored attestation creation in test utils

beacon_node/beacon_chain/src/test_utils.rs

@@ -8,11 +8,7 @@ use std::sync::Arc;
 use store::MemoryStore;
 use store::Store;
 use tree_hash::{SignedRoot, TreeHash};
-use types::{
+use types::{test_utils::TestingBeaconStateBuilder, AggregateSignature, Attestation, AttestationDataAndCustodyBit, BeaconBlock, BeaconState, Bitfield, ChainSpec, Domain, EthSpec, Hash256, Keypair, RelativeEpoch, SecretKey, Signature, Slot, CrosslinkCommittee};
-    test_utils::TestingBeaconStateBuilder, AggregateSignature, Attestation,
-    AttestationDataAndCustodyBit, BeaconBlock, BeaconState, Bitfield, ChainSpec, Domain, EthSpec,
-    Hash256, Keypair, RelativeEpoch, SecretKey, Signature, Slot,
-};
 
 pub use crate::persisted_beacon_chain::{PersistedBeaconChain, BEACON_CHAIN_DB_KEY};
 
@@ -171,7 +167,7 @@ where
             if let BlockProcessingOutcome::Processed { block_root } = outcome {
                 head_block_root = Some(block_root);
 
-                self.add_attestations_to_op_pool(
+                self.add_attestations_to_chain(
                     &attestation_strategy,
                     &new_state,
                     block_root,
@@ -256,18 +252,16 @@ where
         (block, state)
     }
 
-    /// Adds attestations to the `BeaconChain` operations pool to be included in future blocks.
+    /// Adds attestations to the `BeaconChain` operations pool and fork choice.
     ///
     /// The `attestation_strategy` dictates which validators should attest.
-    fn add_attestations_to_op_pool(
+    fn add_attestations_to_chain(
         &self,
         attestation_strategy: &AttestationStrategy,
         state: &BeaconState<E>,
         head_block_root: Hash256,
         head_block_slot: Slot,
     ) {
-        let spec = &self.spec;
-        let fork = &state.fork;
 
         let attesting_validators: Vec<usize> = match attestation_strategy {
             AttestationStrategy::AllValidators => (0..self.keypairs.len()).collect(),
@@ -279,16 +273,45 @@ where
             .expect("should get committees")
             .iter()
             .for_each(|cc| {
-                let committee_size = cc.committee.len();
-
                 for (i, validator_index) in cc.committee.iter().enumerate() {
                     // Note: searching this array is worst-case `O(n)`. A hashset could be a better
                     // alternative.
                     if attesting_validators.contains(validator_index) {
+                        let attestation = self.create_attestation(
+                            *validator_index,
+                            cc,
+                            head_block_root,
+                            head_block_slot,
+                            state,
+                            i
+                        );
+
+                        self.chain
+                            .process_attestation(attestation)
+                            .expect("should process attestation");
+                    }
+                }
+            });
+    }
+
+    /// Creates an attestation for a validator with the given data.
+    pub fn create_attestation(
+        &self,
+        validator_index: usize,
+        crosslink_committee: &CrosslinkCommittee,
+        head_block_root: Hash256,
+        head_block_slot: Slot,
+        state: &BeaconState<E>,
+        bitfield_index: usize
+    ) -> Attestation {
+        let committee_size = crosslink_committee.committee.len();
+        let spec = &self.spec;
+        let fork = &state.fork;
+
         let data = self
             .chain
             .produce_attestation_data_for_block(
-                                cc.shard,
+                crosslink_committee.shard,
                 head_block_root,
                 head_block_slot,
                 state,
@@ -296,7 +319,7 @@ where
             .expect("should produce attestation data");
 
         let mut aggregation_bitfield = Bitfield::new();
-                        aggregation_bitfield.set(i, true);
+        aggregation_bitfield.set(bitfield_index, true);
         aggregation_bitfield.set(committee_size, false);
 
         let mut custody_bitfield = Bitfield::new();
@@ -316,26 +339,20 @@ where
             agg_sig.add(&Signature::new(
                 &message,
                 domain,
-                                self.get_sk(*validator_index),
+                self.get_sk(validator_index),
             ));
 
             agg_sig
         };
 
-                        let attestation = Attestation {
+        Attestation {
             aggregation_bitfield,
             data,
             custody_bitfield,
             signature,
-                        };
+        }
+    }
 
-                        self.chain
-                            .process_attestation(attestation)
-                            .expect("should process attestation");
-                    }
-                }
-            });
-    }
 
     /// Returns the secret key for the given validator index.
     fn get_sk(&self, validator_index: usize) -> &SecretKey {