fix tests

EspressoSystems · Mar 27, 2024 · aff1443 · aff1443
1 parent 86e2d43
commit aff1443
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 7 deletions.
diff --git a/crates/macros/src/lib.rs b/crates/macros/src/lib.rs
@@ -68,7 +68,8 @@ pub fn test_scripts(input: proc_macro::TokenStream) -> TokenStream {
     use hotshot_testing::predicates::Predicate;
     use async_broadcast::broadcast;
     use hotshot_task_impls::events::HotShotEvent;
-
+    use std::time::Duration;
+    use async_compatibility_layer::art::async_timeout;
     use hotshot_task::task::{Task, TaskRegistry, TaskState};
     use hotshot_types::traits::node_implementation::NodeType;
     use std::sync::Arc;
@@ -130,7 +131,7 @@ pub fn test_scripts(input: proc_macro::TokenStream) -> TokenStream {
                         #task_names.state().handle_result(&res).await;
                     }
 
-                    while let Ok(received_output) = test_receiver.try_recv() {
+                    while let Ok(Ok(received_output)) = async_timeout(Duration::from_millis(250), test_receiver.recv_direct()).await {
                         tracing::debug!("Test received: {:?}", received_output);
 
                         let output_asserts = &#task_expectations[stage_number].output_asserts;

diff --git a/crates/testing/src/script.rs b/crates/testing/src/script.rs
@@ -2,9 +2,10 @@ use crate::predicates::Predicate;
 use async_broadcast::broadcast;
 use hotshot_task_impls::events::HotShotEvent;
 
+use async_compatibility_layer::art::async_timeout;
 use hotshot_task::task::{Task, TaskRegistry, TaskState};
 use hotshot_types::traits::node_implementation::NodeType;
-use std::sync::Arc;
+use std::{sync::Arc, time::Duration};
 
 pub struct TestScriptStage<TYPES: NodeType, S: TaskState<Event = Arc<HotShotEvent<TYPES>>>> {
     pub inputs: Vec<HotShotEvent<TYPES>>,
@@ -107,7 +108,9 @@ pub async fn run_test_script<TYPES, S: TaskState<Event = Arc<HotShotEvent<TYPES>
         }
 
         for assert in &stage.outputs {
-            if let Ok(received_output) = test_receiver.try_recv() {
+            if let Ok(Ok(received_output)) =
+                async_timeout(Duration::from_millis(250), test_receiver.recv_direct()).await
+            {
                 tracing::debug!("Test received: {:?}", received_output);
                 validate_output_or_panic(stage_number, &received_output, assert);
             } else {

diff --git a/crates/testing/tests/consensus_task.rs b/crates/testing/tests/consensus_task.rs
@@ -57,6 +57,7 @@ async fn test_consensus_task() {
             QuorumProposalRecv(proposals[0].clone(), leaders[0]),
             DACRecv(dacs[0].clone()),
             VidDisperseRecv(vids[0].0.clone()),
+            QuorumProposalValidated(proposals[0].data.clone()),
         ],
         outputs: vec![
             exact(ViewChange(ViewNumber::new(1))),
@@ -75,11 +76,12 @@ async fn test_consensus_task() {
             QCFormed(either::Left(cert)),
             // We must have a payload commitment and metadata to propose.
             SendPayloadCommitmentAndMetadata(payload_commitment, (), ViewNumber::new(2)),
+            QuorumProposalValidated(proposals[1].data.clone()),
         ],
         outputs: vec![
             exact(ViewChange(ViewNumber::new(2))),
-            exact(QuorumProposalValidated(proposals[1].data.clone())),
             quorum_proposal_send(),
+            exact(QuorumProposalValidated(proposals[1].data.clone())),
         ],
         asserts: vec![is_at_view_number(2)],
     };

diff --git a/crates/testing/tests/proposal_ordering.rs b/crates/testing/tests/proposal_ordering.rs
@@ -53,6 +53,7 @@ async fn test_ordering_with_specific_order(input_permutation: Vec<usize>) {
             QuorumProposalRecv(proposals[0].clone(), leaders[0]),
             DACRecv(dacs[0].clone()),
             VidDisperseRecv(vids[0].0.clone()),
+            QuorumProposalValidated(proposals[0].data.clone()),
         ],
         outputs: vec![
             exact(ViewChange(ViewNumber::new(1))),
@@ -83,8 +84,8 @@ async fn test_ordering_with_specific_order(input_permutation: Vec<usize>) {
     } else {
         vec![
             exact(ViewChange(ViewNumber::new(2))),
-            exact(QuorumProposalValidated(proposals[1].data.clone())),
             quorum_proposal_send(),
+            exact(QuorumProposalValidated(proposals[1].data.clone())),
         ]
     };
 

diff --git a/crates/testing/tests/upgrade_task.rs b/crates/testing/tests/upgrade_task.rs
@@ -69,11 +69,12 @@ async fn test_upgrade_task() {
             QuorumProposalRecv(proposals[0].clone(), leaders[0]),
             VidDisperseRecv(vids[0].0.clone()),
             DACRecv(dacs[0].clone()),
+            QuorumProposalValidated(proposals[0].data.clone()),
         ],
         outputs: vec![
             exact(ViewChange(ViewNumber::new(1))),
-            exact(QuorumProposalValidated(proposals[0].data.clone())),
             exact(QuorumVoteSend(votes[0].clone())),
+            exact(QuorumProposalValidated(proposals[0].data.clone())),
         ],
         asserts: vec![],
     };