From 89cc92572a1764be16c12dcdd8b205fea1770b23 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Sat, 30 Mar 2019 13:03:05 +1100
Subject: [PATCH] Add `test_harness` tests for attestation count

---
 .../specs/validator_registry.yaml             |  2 ++
 .../test_harness/src/test_case/state_check.rs | 27 +++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/beacon_node/beacon_chain/test_harness/specs/validator_registry.yaml b/beacon_node/beacon_chain/test_harness/specs/validator_registry.yaml
index 0c4f5004b..1674ecffc 100644
--- a/beacon_node/beacon_chain/test_harness/specs/validator_registry.yaml
+++ b/beacon_node/beacon_chain/test_harness/specs/validator_registry.yaml
@@ -47,6 +47,8 @@ test_cases:
       states:
         - slot: 63
           num_validators: 1003
+          num_previous_epoch_attestations: 0
+          num_current_epoch_attestations: 10
           slashed_validators: [11, 12, 13, 14, 42]
           exited_validators: []
           exit_initiated_validators: [50]
diff --git a/beacon_node/beacon_chain/test_harness/src/test_case/state_check.rs b/beacon_node/beacon_chain/test_harness/src/test_case/state_check.rs
index 4d2bfd07d..7ac33c86c 100644
--- a/beacon_node/beacon_chain/test_harness/src/test_case/state_check.rs
+++ b/beacon_node/beacon_chain/test_harness/src/test_case/state_check.rs
@@ -16,6 +16,10 @@ pub struct StateCheck {
     pub slot: Slot,
     /// Checked against `beacon_state.validator_registry.len()`.
     pub num_validators: Option<usize>,
+    /// The number of pending attestations from the previous epoch that should be in the state.
+    pub num_previous_epoch_attestations: Option<usize>,
+    /// The number of pending attestations from the current epoch that should be in the state.
+    pub num_current_epoch_attestations: Option<usize>,
     /// A list of validator indices which have been penalized. Must be in ascending order.
     pub slashed_validators: Option<Vec<u64>>,
     /// A list of validator indices which have been fully exited. Must be in ascending order.
@@ -34,6 +38,8 @@ impl StateCheck {
         Self {
             slot: Slot::from(as_u64(&yaml, "slot").expect("State must specify slot")),
             num_validators: as_usize(&yaml, "num_validators"),
+            num_previous_epoch_attestations: as_usize(&yaml, "num_previous_epoch_attestations"),
+            num_current_epoch_attestations: as_usize(&yaml, "num_current_epoch_attestations"),
             slashed_validators: as_vec_u64(&yaml, "slashed_validators"),
             exited_validators: as_vec_u64(&yaml, "exited_validators"),
             exit_initiated_validators: as_vec_u64(&yaml, "exit_initiated_validators"),
@@ -58,6 +64,7 @@ impl StateCheck {
             "State slot is invalid."
         );
 
+        // Check the validator count
         if let Some(num_validators) = self.num_validators {
             assert_eq!(
                 state.validator_registry.len(),
@@ -67,6 +74,26 @@ impl StateCheck {
             info!("OK: num_validators = {}.", num_validators);
         }
 
+        // Check the previous epoch attestations
+        if let Some(n) = self.num_previous_epoch_attestations {
+            assert_eq!(
+                state.previous_epoch_attestations.len(),
+                n,
+                "previous epoch attestations count != expected."
+            );
+            info!("OK: num_previous_epoch_attestations = {}.", n);
+        }
+
+        // Check the current epoch attestations
+        if let Some(n) = self.num_current_epoch_attestations {
+            assert_eq!(
+                state.current_epoch_attestations.len(),
+                n,
+                "current epoch attestations count != expected."
+            );
+            info!("OK: num_current_epoch_attestations = {}.", n);
+        }
+
         // Check for slashed validators.
         if let Some(ref slashed_validators) = self.slashed_validators {
             let actually_slashed_validators: Vec<u64> = state