From b21cc64949e5abeb6f9d9805daa4e5449477367e Mon Sep 17 00:00:00 2001
From: Michael Sproul <michael@sigmaprime.io>
Date: Tue, 2 Apr 2019 18:46:08 +1100
Subject: [PATCH] state transition tests: check expected state

---
 eth2/state_processing/tests/tests.rs | 52 ++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/eth2/state_processing/tests/tests.rs b/eth2/state_processing/tests/tests.rs
index 6ea8863b8..193511852 100644
--- a/eth2/state_processing/tests/tests.rs
+++ b/eth2/state_processing/tests/tests.rs
@@ -10,6 +10,45 @@ use types::*;
 #[allow(unused_imports)]
 use yaml_utils;
 
+#[derive(Debug, Deserialize)]
+pub struct ExpectedState {
+    pub slot: Option<Slot>,
+    pub genesis_time: Option<u64>,
+    pub fork: Option<Fork>,
+    pub validator_registry: Option<Vec<Validator>>,
+    pub validator_balances: Option<Vec<u64>>,
+}
+
+impl ExpectedState {
+    // Return a list of fields that differ, and a string representation of the beacon state's field.
+    fn check(&self, state: &BeaconState) -> Vec<(&str, String)> {
+        // Check field equality
+        macro_rules! cfe {
+            ($field_name:ident) => {
+                if self.$field_name.as_ref().map_or(true, |$field_name| {
+                    println!("  > Checking {}", stringify!($field_name));
+                    $field_name == &state.$field_name
+                }) {
+                    vec![]
+                } else {
+                    vec![(stringify!($field_name), format!("{:#?}", state.$field_name))]
+                }
+            };
+        }
+
+        vec![
+            cfe!(slot),
+            cfe!(genesis_time),
+            cfe!(fork),
+            cfe!(validator_registry),
+            cfe!(validator_balances),
+        ]
+        .into_iter()
+        .flat_map(|x| x)
+        .collect()
+    }
+}
+
 #[derive(Debug, Deserialize)]
 pub struct TestCase {
     pub name: String,
@@ -17,6 +56,7 @@ pub struct TestCase {
     pub verify_signatures: bool,
     pub initial_state: BeaconState,
     pub blocks: Vec<BeaconBlock>,
+    pub expected_state: ExpectedState,
 }
 
 #[derive(Debug, Deserialize)]
@@ -70,6 +110,18 @@ fn run_state_transition_test(test_name: &str) {
                 ok = false;
             };
         }
+
+        let mismatched_fields = test_case.expected_state.check(&state);
+        if !mismatched_fields.is_empty() {
+            println!(
+                "Error in expected state, these fields didn't match: {:?}",
+                mismatched_fields.iter().map(|(f, _)| f).collect::<Vec<_>>()
+            );
+            for (field_name, state_val) in mismatched_fields {
+                println!("state.{} was: {}", field_name, state_val);
+            }
+            ok = false;
+        }
     }
 
     assert!(ok, "one or more tests failed, see above");