From aff144377c1d8739c77d940d3761001adb786d7e Mon Sep 17 00:00:00 2001
From: Brendon Fish <bfish713@gmail.com>
Date: Wed, 27 Mar 2024 13:23:10 -0400
Subject: [PATCH] fix tests

---
 crates/macros/src/lib.rs                  | 5 +++--
 crates/testing/src/script.rs              | 7 +++++--
 crates/testing/tests/consensus_task.rs    | 4 +++-
 crates/testing/tests/proposal_ordering.rs | 3 ++-
 crates/testing/tests/upgrade_task.rs      | 3 ++-
 5 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/crates/macros/src/lib.rs b/crates/macros/src/lib.rs
index 45f6f4805f..91e2dee0eb 100644
--- 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
index 3c98f5480b..5ed60d5b9c 100644
--- 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
index 9d83cc9247..0c8d4e24f0 100644
--- 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
index a96b16a9d2..1b38fc120c 100644
--- 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
index 4f9f10358e..fc9986c1a1 100644
--- 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![],
     };