Create test def concept

eth2/proto_array_fork_choice/tests/test.rs

@@ -3,14 +3,14 @@ use types::{Epoch, Hash256, Slot};
 
 #[test]
 fn no_votes() {
-    let tester = ForkChoiceTester::new(Slot::new(0), Epoch::new(1), Epoch::new(1), get_hash(0));
+    let test = get_no_votes_test_definition();
-    tester.process_operations(get_no_votes_ops())
+    test.run();
 }
 
 #[test]
 fn votes() {
-    let tester = ForkChoiceTester::new(Slot::new(0), Epoch::new(1), Epoch::new(1), get_hash(0));
+    let test = get_votes_test_definition();
-    tester.process_operations(get_votes_ops())
+    test.run();
 }
 
 #[derive(Debug, Clone)]
@@ -48,35 +48,25 @@ pub enum Operation {
     },
 }
 
-pub struct ForkChoiceTester {
+pub struct ForkChoiceTestDefinition {
-    fork_choice: ProtoArrayForkChoice,
+    pub finalized_block_slot: Slot,
+    pub justified_epoch: Epoch,
+    pub finalized_epoch: Epoch,
+    pub finalized_root: Hash256,
+    pub operations: Vec<Operation>,
 }
 
-impl ForkChoiceTester {
+impl ForkChoiceTestDefinition {
-    pub fn new(
+    fn run(self) {
-        finalized_block_slot: Slot,
+        let fork_choice = ProtoArrayForkChoice::new(
-        justified_epoch: Epoch,
+            self.finalized_block_slot,
-        finalized_epoch: Epoch,
+            self.justified_epoch,
-        finalized_root: Hash256,
+            self.finalized_epoch,
-    ) -> Self {
+            self.finalized_root,
-        Self {
-            fork_choice: ProtoArrayForkChoice::new(
-                finalized_block_slot,
-                justified_epoch,
-                finalized_epoch,
-                finalized_root,
         )
-            .expect("should create fork choice"),
+        .expect("should create fork choice struct");
-        }
-    }
 
-    pub fn process_operations(&self, ops: Vec<Operation>) {
+        for (op_index, op) in self.operations.into_iter().enumerate() {
-        for (i, op) in ops.into_iter().enumerate() {
-            self.process_operation(i, op)
-        }
-    }
-
-    fn process_operation(&self, op_index: usize, op: Operation) {
             match op.clone() {
                 Operation::FindHead {
                     justified_epoch,
@@ -85,8 +75,7 @@ impl ForkChoiceTester {
                     justified_state_balances,
                     expected_head,
                 } => {
-                let head = self
+                    let head = fork_choice
-                    .fork_choice
                         .find_head(
                             justified_epoch,
                             justified_root,
@@ -103,7 +92,7 @@ impl ForkChoiceTester {
                         "Operation at index {} failed checks. Operation: {:?}",
                         op_index, op
                     );
-                check_bytes_round_trip(&self.fork_choice);
+                    check_bytes_round_trip(&fork_choice);
                 }
                 Operation::InvalidFindHead {
                     justified_epoch,
@@ -111,7 +100,7 @@ impl ForkChoiceTester {
                     finalized_epoch,
                     justified_state_balances,
                 } => {
-                let result = self.fork_choice.find_head(
+                    let result = fork_choice.find_head(
                         justified_epoch,
                         justified_root,
                         finalized_epoch,
@@ -124,7 +113,7 @@ impl ForkChoiceTester {
                         op_index,
                         op
                     );
-                check_bytes_round_trip(&self.fork_choice);
+                    check_bytes_round_trip(&fork_choice);
                 }
                 Operation::ProcessBlock {
                     slot,
@@ -133,26 +122,26 @@ impl ForkChoiceTester {
                     justified_epoch,
                     finalized_epoch,
                 } => {
-                self.fork_choice
+                    fork_choice
                         .process_block(slot, root, parent_root, justified_epoch, finalized_epoch)
                         .expect(&format!(
                             "process_block op at index {} returned error",
                             op_index
                         ));
-                check_bytes_round_trip(&self.fork_choice);
+                    check_bytes_round_trip(&fork_choice);
                 }
                 Operation::ProcessAttestation {
                     validator_index,
                     block_root,
                     target_epoch,
                 } => {
-                self.fork_choice
+                    fork_choice
                         .process_attestation(validator_index, block_root, target_epoch)
                         .expect(&format!(
                             "process_attestation op at index {} returned error",
                             op_index
                         ));
-                check_bytes_round_trip(&self.fork_choice);
+                    check_bytes_round_trip(&fork_choice);
                 }
                 Operation::Prune {
                     finalized_epoch,
@@ -160,24 +149,25 @@ impl ForkChoiceTester {
                     prune_threshold,
                     expected_len,
                 } => {
-                self.fork_choice.set_prune_threshold(prune_threshold);
+                    fork_choice.set_prune_threshold(prune_threshold);
-                self.fork_choice
+                    fork_choice
                         .update_finalized_root(finalized_epoch, finalized_root)
                         .expect("update_finalized_root op at index {} returned error");
 
                     // Ensure that no pruning happened.
                     assert_eq!(
-                    self.fork_choice.len(),
+                        fork_choice.len(),
                         expected_len,
                         "Prune op at index {} failed with {} instead of {}",
                         op_index,
-                    self.fork_choice.len(),
+                        fork_choice.len(),
                         expected_len
                     );
                 }
             }
         }
     }
+}
 
 /// Gives a hash that is not the zero hash (unless i is `usize::max_value)`.
 fn get_hash(i: u64) -> Hash256 {
@@ -194,10 +184,10 @@ fn check_bytes_round_trip(original: &ProtoArrayForkChoice) {
     );
 }
 
-fn get_no_votes_ops() -> Vec<Operation> {
+fn get_no_votes_test_definition() -> ForkChoiceTestDefinition {
     let balances = vec![0; 16];
 
-    vec![
+    let operations = vec![
         // Check that the head is the finalized block.
         Operation::FindHead {
             justified_epoch: Epoch::new(1),
@@ -409,10 +399,18 @@ fn get_no_votes_ops() -> Vec<Operation> {
             justified_state_balances: balances.clone(),
             expected_head: get_hash(6),
         },
-    ]
+    ];
+
+    ForkChoiceTestDefinition {
+        finalized_block_slot: Slot::new(0),
+        justified_epoch: Epoch::new(1),
+        finalized_epoch: Epoch::new(1),
+        finalized_root: get_hash(0),
+        operations,
+    }
 }
 
-fn get_votes_ops() -> Vec<Operation> {
+fn get_votes_test_definition() -> ForkChoiceTestDefinition {
     let mut balances = vec![1; 2];
     let mut ops = vec![];
 
@@ -1091,5 +1089,11 @@ fn get_votes_ops() -> Vec<Operation> {
         expected_head: get_hash(11),
     });
 
-    ops
+    ForkChoiceTestDefinition {
+        finalized_block_slot: Slot::new(0),
+        justified_epoch: Epoch::new(1),
+        finalized_epoch: Epoch::new(1),
+        finalized_root: get_hash(0),
+        operations: ops,
+    }
 }