Add tests for the upgrade task (#2642)

* Adds stronger testing harness

* Adds an iterator to generate test data for views

* Adds a test to validate that consensus correctly updates its internal state when receiving an UpgradeCertificate

* Correctly a mistake in consensus logic that took the UpgradeCertificate from the received QuourmProposal rather than the decided leaf.

Author: ss-es

hotshot/src/lib.rs

@@ -234,6 +234,7 @@ impl<TYPES: NodeType, I: NodeImplementation<TYPES>> SystemContext<TYPES, I> {
             saved_leaves,
             saved_payloads,
             saved_da_certs: HashMap::new(),
+            saved_upgrade_certs: HashMap::new(),
             // TODO this is incorrect
             // https://github.com/EspressoSystems/HotShot/issues/560
             locked_view: anchored_leaf.get_view_number(),

task-impls/src/consensus.rs

@@ -498,7 +498,13 @@ impl<TYPES: NodeType, I: NodeImplementation<TYPES>, A: ConsensusApi<TYPES, I> +
                 // We should just make sure we don't *sign* an UpgradeCertificate for an upgrade
                 // that we do not support.
                 if let Some(ref upgrade_cert) = proposal.data.upgrade_certificate {
-                    if !upgrade_cert.is_valid_cert(self.quorum_membership.as_ref()) {
+                    if upgrade_cert.is_valid_cert(self.quorum_membership.as_ref()) {
+                        self.consensus
+                            .write()
+                            .await
+                            .saved_upgrade_certs
+                            .insert(view, upgrade_cert.clone());
+                    } else {
                         error!("Invalid upgrade_cert in proposal for view {}", *view);
                         return;
                     }
@@ -740,7 +746,7 @@ impl<TYPES: NodeType, I: NodeImplementation<TYPES>, A: ConsensusApi<TYPES, I> +
                                         .last_synced_block_height
                                         .set(usize::try_from(leaf.get_height()).unwrap_or(0));
                                 }
-                                if let Some(ref upgrade_cert) = proposal.data.upgrade_certificate {
+                                if let Some(upgrade_cert) = consensus.saved_upgrade_certs.get(&leaf.get_view_number()) {
                                     info!("Updating consensus state with decided upgrade certificate: {:?}", upgrade_cert);
                                     self.decided_upgrade_cert = Some(upgrade_cert.clone());
                                 }

task-impls/src/upgrade.rs

@@ -205,7 +205,7 @@ impl<TYPES: NodeType, I: NodeImplementation<TYPES>, A: ConsensusApi<TYPES, I> +
     }
 }
 
-/// task state implementation for DA Task
+/// task state implementation for the upgrade task
 impl<TYPES: NodeType, I: NodeImplementation<TYPES>, A: ConsensusApi<TYPES, I> + 'static> TaskState
     for UpgradeTaskState<TYPES, I, A>
 {

task/src/task.rs

@@ -166,6 +166,11 @@ impl<S: TaskState + Send + 'static> Task<S> {
     pub fn state_mut(&mut self) -> &mut S {
         &mut self.state
     }
+    /// Get an immutable reference to this tasks state
+    pub fn state(&self) -> &S {
+        &self.state
+    }
+
     /// Spawn a new task adn register it.  It will get all events not seend
     /// by the task creating it.
     pub async fn run_sub_task(&self, state: S) {

testing/src/lib.rs

@@ -33,6 +33,15 @@ pub mod spinning_task;
 /// task for checking if view sync got activated
 pub mod view_sync_task;
 
+/// predicates to use in tests
+pub mod predicates;
+
+/// scripting harness for tests
+pub mod script;
+
+/// view generator for tests
+pub mod view_generator;
+
 /// global event at the test level
 #[derive(Clone, Debug)]
 pub enum GlobalTestEvent {

testing/src/predicates.rs

@@ -0,0 +1,78 @@
+use hotshot_task_impls::{
+    consensus::ConsensusTaskState, events::HotShotEvent, events::HotShotEvent::*,
+};
+use hotshot_types::traits::node_implementation::NodeType;
+
+use hotshot::types::SystemContextHandle;
+
+use hotshot_example_types::node_types::{MemoryImpl, TestTypes};
+
+pub struct Predicate<INPUT> {
+    pub function: Box<dyn Fn(&INPUT) -> bool>,
+    pub info: String,
+}
+
+pub fn exact<TYPES>(event: HotShotEvent<TYPES>) -> Predicate<HotShotEvent<TYPES>>
+where
+    TYPES: NodeType,
+{
+    let info = format!("{:?}", event);
+
+    Predicate {
+        function: Box::new(move |e| e == &event),
+        info,
+    }
+}
+
+pub fn leaf_decided<TYPES>() -> Predicate<HotShotEvent<TYPES>>
+where
+    TYPES: NodeType,
+{
+    let info = "LeafDecided".to_string();
+    let function = |e: &_| matches!(e, LeafDecided(_));
+
+    Predicate {
+        function: Box::new(function),
+        info,
+    }
+}
+
+pub fn quorum_vote_send<TYPES>() -> Predicate<HotShotEvent<TYPES>>
+where
+    TYPES: NodeType,
+{
+    let info = "QuorumVoteSend".to_string();
+    let function = |e: &_| matches!(e, QuorumVoteSend(_));
+
+    Predicate {
+        function: Box::new(function),
+        info,
+    }
+}
+
+type ConsensusTaskTestState =
+    ConsensusTaskState<TestTypes, MemoryImpl, SystemContextHandle<TestTypes, MemoryImpl>>;
+
+pub fn consensus_predicate(
+    function: Box<dyn for<'a> Fn(&'a ConsensusTaskTestState) -> bool>,
+    info: &str,
+) -> Predicate<ConsensusTaskTestState> {
+    Predicate {
+        function,
+        info: info.to_string(),
+    }
+}
+
+pub fn no_decided_upgrade_cert() -> Predicate<ConsensusTaskTestState> {
+    consensus_predicate(
+        Box::new(|state| state.decided_upgrade_cert.is_none()),
+        "expected decided_upgrade_cert to be None",
+    )
+}
+
+pub fn decided_upgrade_cert() -> Predicate<ConsensusTaskTestState> {
+    consensus_predicate(
+        Box::new(|state| state.decided_upgrade_cert.is_some()),
+        "expected decided_upgrade_cert to be Some(_)",
+    )
+}

testing/src/script.rs

@@ -0,0 +1,106 @@
+use crate::predicates::Predicate;
+use async_broadcast::broadcast;
+use hotshot_task_impls::events::HotShotEvent;
+
+use hotshot_task::task::{Task, TaskRegistry, TaskState};
+use hotshot_types::traits::node_implementation::NodeType;
+use std::sync::Arc;
+
+/// A `TestScript` is a sequence of triples (input sequence, output sequence, assertions).
+type TestScript<TYPES, S> = Vec<TestScriptStage<TYPES, S>>;
+
+pub struct TestScriptStage<TYPES: NodeType, S: TaskState<Event = HotShotEvent<TYPES>>> {
+    pub inputs: Vec<HotShotEvent<TYPES>>,
+    pub outputs: Vec<Predicate<HotShotEvent<TYPES>>>,
+    pub asserts: Vec<Predicate<S>>,
+}
+
+impl<INPUT> std::fmt::Debug for Predicate<INPUT> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
+        write!(f, "{}", self.info)
+    }
+}
+
+/// `run_test_script` reads a triple (inputs, outputs, asserts) in a `TestScript`,
+/// It broadcasts all given inputs (in order) and waits to receive all outputs (in order).
+/// Once the expected outputs have been received, it validates the task state at that stage
+/// against the given assertions.
+///
+/// If all assertions pass, it moves onto the next stage. If it receives an unexpected output
+/// or fails to receive an output, the test fails immediately with a panic.
+///
+/// Note: the task is not spawned with an async thread; instead, the harness just calls `handle_event`.
+/// This has a few implications, e.g. shutting down tasks doesn't really make sense,
+/// and event ordering is deterministic.
+pub async fn run_test_script<TYPES, S: TaskState<Event = HotShotEvent<TYPES>>>(
+    mut script: TestScript<TYPES, S>,
+    state: S,
+) where
+    TYPES: NodeType,
+    S: Send + 'static,
+{
+    let registry = Arc::new(TaskRegistry::default());
+
+    let (test_input, task_receiver) = broadcast(1024);
+    // let (task_input, mut test_receiver) = broadcast(1024);
+
+    let task_input = test_input.clone();
+    let mut test_receiver = task_receiver.clone();
+
+    let mut task = Task::new(
+        task_input.clone(),
+        task_receiver.clone(),
+        registry.clone(),
+        state,
+    );
+
+    for (stage_number, stage) in script.iter_mut().enumerate() {
+        tracing::debug!("Beginning test stage {}", stage_number);
+        for input in &mut *stage.inputs {
+            if !task.state_mut().filter(input) {
+                tracing::debug!("Test sent: {:?}", input);
+
+                if let Some(res) = S::handle_event(input.clone(), &mut task).await {
+                    task.state_mut().handle_result(&res).await;
+                }
+            }
+        }
+
+        for expected in &stage.outputs {
+            let output_missing_error = format!(
+                "Stage {} | Failed to receive output for predicate: {:?}",
+                stage_number, expected
+            );
+
+            if let Ok(received_output) = test_receiver.try_recv() {
+                tracing::debug!("Test received: {:?}", received_output);
+                assert!(
+                    (expected.function)(&received_output),
+                    "Stage {} | Output failed to satisfy {:?}",
+                    stage_number,
+                    expected
+                );
+            } else {
+                panic!("{}", output_missing_error);
+            }
+        }
+
+        for assert in &stage.asserts {
+            assert!(
+                (assert.function)(task.state()),
+                "Stage {} | Assertion on task state failed: {:?}",
+                stage_number,
+                assert
+            );
+        }
+
+        if let Ok(received_output) = test_receiver.try_recv() {
+            let extra_output_error = format!(
+                "Stage {} | Received unexpected additional output: {:?}",
+                stage_number, received_output
+            );
+
+            panic!("{}", extra_output_error);
+        }
+    }
+}