From 2df71372d236efad4268fa9cd78dc2f8d12bfe07 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Thu, 16 Jan 2020 17:02:19 +1100
Subject: [PATCH] Add more fork choice tests

---
 .../ffg_updates.rs                            | 105 +++++++++++++++++-
 1 file changed, 104 insertions(+), 1 deletion(-)

diff --git a/eth2/proto_array_fork_choice/src/fork_choice_test_definition/ffg_updates.rs b/eth2/proto_array_fork_choice/src/fork_choice_test_definition/ffg_updates.rs
index b4589c6362..5a3cd26357 100644
--- a/eth2/proto_array_fork_choice/src/fork_choice_test_definition/ffg_updates.rs
+++ b/eth2/proto_array_fork_choice/src/fork_choice_test_definition/ffg_updates.rs
@@ -206,7 +206,7 @@ pub fn get_ffg_case_02_test_definition() -> ForkChoiceTestDefinition {
         finalized_epoch: Epoch::new(0),
     });
 
-    // Ensure that if we start at 0 we find 10.
+    // Ensure that if we start at 0 we find 10 (just: 0, fin: 0).
     //
     //           0  <-- start
     //          / \
@@ -227,6 +227,109 @@ pub fn get_ffg_case_02_test_definition() -> ForkChoiceTestDefinition {
         expected_head: get_hash(10),
     });
 
+    // Add a vote to 1.
+    //
+    //                 0
+    //                / \
+    //    +1 vote -> 1   2
+    //               |   |
+    //               3   4
+    //               |   |
+    //               5   6
+    //               |   |
+    //               7   8
+    //               |   |
+    //               9  10
+    ops.push(Operation::ProcessAttestation {
+        validator_index: 0,
+        block_root: get_hash(1),
+        target_epoch: Epoch::new(0),
+    });
+
+    // Ensure that if we start at 0 we find 9 (just: 0, fin: 0).
+    //
+    //           0  <-- start
+    //          / \
+    //         1   2
+    //         |   |
+    //         3   4
+    //         |   |
+    //         5   6
+    //         |   |
+    //         7   8
+    //         |   |
+    // head -> 9  10
+    ops.push(Operation::FindHead {
+        justified_epoch: Epoch::new(0),
+        justified_root: get_hash(0),
+        finalized_epoch: Epoch::new(0),
+        justified_state_balances: balances.clone(),
+        expected_head: get_hash(9),
+    });
+
+    // Add a vote to 2.
+    //
+    //                 0
+    //                / \
+    //               1   2 <- +1 vote
+    //               |   |
+    //               3   4
+    //               |   |
+    //               5   6
+    //               |   |
+    //               7   8
+    //               |   |
+    //               9  10
+    ops.push(Operation::ProcessAttestation {
+        validator_index: 1,
+        block_root: get_hash(2),
+        target_epoch: Epoch::new(0),
+    });
+
+    // Ensure that if we start at 0 we find 10 (just: 0, fin: 0).
+    //
+    //           0  <-- start
+    //          / \
+    //         1   2
+    //         |   |
+    //         3   4
+    //         |   |
+    //         5   6
+    //         |   |
+    //         7   8
+    //         |   |
+    //         9  10 <-- head
+    ops.push(Operation::FindHead {
+        justified_epoch: Epoch::new(0),
+        justified_root: get_hash(0),
+        finalized_epoch: Epoch::new(0),
+        justified_state_balances: balances.clone(),
+        expected_head: get_hash(10),
+    });
+
+    // Ensure that if we start at 1 we find 9 (just: 0, fin: 0).
+    //
+    //            0
+    //           / \
+    //          1   2
+    //          |   |
+    // start -> 3   4
+    //          |   |
+    //          5   6
+    //          |   |
+    //          7   8
+    //          |   |
+    //  head -> 9  10
+    ops.push(Operation::FindHead {
+        justified_epoch: Epoch::new(0),
+        justified_root: get_hash(0),
+        finalized_epoch: Epoch::new(0),
+        justified_state_balances: balances.clone(),
+        expected_head: get_hash(10),
+    });
+
+    // TODO: add more tests here.
+
     // END OF TESTS
     ForkChoiceTestDefinition {
         finalized_block_slot: Slot::new(0),