Add attestation strategy to chain harness

beacon_node/beacon_chain/src/test_utils.rs

@@ -20,6 +20,11 @@ pub enum BuildStrategy {
     ForkCanonicalChainAt(Slot),
 }
 
+#[derive(Clone, Copy, Debug)]
+pub enum AttestationStrategy {
+    AllValidators,
+}
+
 pub struct CommonTypes<L, E>
 where
     L: LmdGhost<MemoryStore, E>,
@@ -95,7 +100,12 @@ where
         self.chain.catchup_state().expect("should catchup state");
     }
 
-    pub fn extend_chain(&self, build_strategy: BuildStrategy, blocks: usize) {
+    pub fn extend_chain(
+        &self,
+        build_strategy: BuildStrategy,
+        blocks: usize,
+        attestation_strategy: AttestationStrategy,
+    ) {
         // Get an initial state to build the block upon, based on the build strategy.
         let mut state = match build_strategy {
             BuildStrategy::OnCanonicalHead => self.chain.current_state().clone(),
@@ -134,7 +144,12 @@ where
                 .expect("should not error during block processing");
 
             if let BlockProcessingOutcome::Processed { block_root } = outcome {
-                self.add_attestations_to_op_pool(&new_state, block_root, slot);
+                self.add_attestations_to_op_pool(
+                    attestation_strategy,
+                    &new_state,
+                    block_root,
+                    slot,
+                );
             } else {
                 panic!("block should be successfully processed");
             }
@@ -198,6 +213,7 @@ where
 
     fn add_attestations_to_op_pool(
         &self,
+        attestation_strategy: AttestationStrategy,
         state: &BeaconState<E>,
         head_block_root: Hash256,
         head_block_slot: Slot,
@@ -205,6 +221,10 @@ where
         let spec = &self.spec;
         let fork = &state.fork;
 
+        let attesting_validators: Vec<usize> = match attestation_strategy {
+            AttestationStrategy::AllValidators => (0..self.keypairs.len()).collect(),
+        };
+
         state
             .get_crosslink_committees_at_slot(state.slot)
             .expect("should get committees")
@@ -213,6 +233,9 @@ where
                 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 data = self
                             .chain
                             .produce_attestation_data_for_block(
@@ -237,7 +260,8 @@ where
                             }
                             .tree_hash_root();
 
-                        let domain = spec.get_domain(data.target_epoch, Domain::Attestation, fork);
+                            let domain =
+                                spec.get_domain(data.target_epoch, Domain::Attestation, fork);
 
                             let mut agg_sig = AggregateSignature::new();
                             agg_sig.add(&Signature::new(
@@ -260,6 +284,7 @@ where
                             .process_attestation(attestation)
                             .expect("should process attestation");
                     }
+                }
             });
     }
 
@@ -295,7 +320,11 @@ mod test {
 
         let harness = get_harness(VALIDATOR_COUNT);
 
-        harness.extend_chain(BuildStrategy::OnCanonicalHead, num_blocks_produced as usize);
+        harness.extend_chain(
+            BuildStrategy::OnCanonicalHead,
+            num_blocks_produced as usize,
+            AttestationStrategy::AllValidators,
+        );
 
         let state = &harness.chain.head().beacon_state;