1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

// Copyright (c) 2021-2024 Espresso Systems (espressosys.com)
// This file is part of the HotShot repository.

// You should have received a copy of the MIT License
// along with the HotShot repository. If not, see <https://mit-license.org/>.

use std::{sync::Arc, time::Duration};

use hotshot_task_impls::events::HotShotEvent;
use hotshot_types::traits::node_implementation::NodeType;

use crate::predicates::{Predicate, PredicateResult};

pub enum InputOrder<TYPES: NodeType> {
    Random(Vec<HotShotEvent<TYPES>>),
    Serial(Vec<HotShotEvent<TYPES>>),
}

#[macro_export]
macro_rules! random {
    ($($x:expr),* $(,)?) => {
        {
            let inputs = vec![$($x),*];
            InputOrder::Random(inputs)
        }
    };
}

#[macro_export]
macro_rules! serial {
    ($($x:expr),* $(,)?) => {
        {
            let inputs = vec![$($x),*];
            InputOrder::Serial(inputs)
        }
    };
}

pub struct TaskScript<TYPES: NodeType, S> {
    /// The time to wait on the receiver for this script.
    pub timeout: Duration,
    pub state: S,
    pub expectations: Vec<Expectations<TYPES, S>>,
}

pub struct Expectations<TYPES: NodeType, S> {
    pub output_asserts: Vec<Box<dyn Predicate<Arc<HotShotEvent<TYPES>>>>>,
    pub task_state_asserts: Vec<Box<dyn Predicate<S>>>,
}

impl<TYPES: NodeType, S> Expectations<TYPES, S> {
    pub fn from_outputs(output_asserts: Vec<Box<dyn Predicate<Arc<HotShotEvent<TYPES>>>>>) -> Self {
        Self {
            output_asserts,
            task_state_asserts: vec![],
        }
    }
    pub fn from_outputs_and_task_states(
        output_asserts: Vec<Box<dyn Predicate<Arc<HotShotEvent<TYPES>>>>>,
        task_state_asserts: Vec<Box<dyn Predicate<S>>>,
    ) -> Self {
        Self {
            output_asserts,
            task_state_asserts,
        }
    }
}

pub fn panic_extra_output_in_script<S>(stage_number: usize, script_name: String, output: &S)
where
    S: std::fmt::Debug,
{
    let extra_output_error = format!(
        "Stage {} | Received unexpected additional output in {}:\n\n{:?}",
        stage_number, script_name, output
    );

    panic!("{}", extra_output_error);
}

pub fn panic_missing_output_in_script<S>(stage_number: usize, script_name: String, predicate: &S)
where
    S: std::fmt::Debug,
{
    let output_missing_error = format!(
        "Stage {} | Failed to receive output for predicate in {}: {:?}",
        stage_number, script_name, predicate
    );

    panic!("{}", output_missing_error);
}

pub async fn validate_task_state_or_panic_in_script<S>(
    stage_number: usize,
    script_name: String,
    state: &S,
    assert: &dyn Predicate<S>,
) {
    assert!(
        assert.evaluate(state).await == PredicateResult::Pass,
        "Stage {} | Task state in {} failed to satisfy: {:?}",
        stage_number,
        script_name,
        assert
    );
}

pub async fn validate_output_or_panic_in_script<S: std::fmt::Debug>(
    stage_number: usize,
    script_name: String,
    output: &S,
    assert: &dyn Predicate<S>,
) -> PredicateResult {
    let result = assert.evaluate(output).await;

    match result {
        PredicateResult::Pass => result,
        PredicateResult::Incomplete => result,
        PredicateResult::Fail => {
            panic!(
                "Stage {} | Output in {} failed to satisfy: {:?}.\n\nReceived:\n\n{:?}",
                stage_number, script_name, assert, output
            )
        }
    }
}