use std::fmt::Debug;
use trellis_core::{
GraphResult, InputNode, OutputFrameTrace, ResourceCommandTrace, Transaction, TransactionResult,
};
use crate::{ScenarioError, ScenarioTarget, StageOperation, TrellisHarness};
pub(crate) type InvariantCheck<G, C> = dyn Fn(&G, &TransactionResult<C>) -> bool + 'static;
pub(crate) struct NamedInvariantCheck<G, C> {
pub(crate) name: String,
pub(crate) check: Box<InvariantCheck<G, C>>,
}
pub struct HarnessStep<'harness, G, C> {
harness: &'harness mut TrellisHarness<G, C>,
name: String,
operations: Vec<Box<StageOperation<C>>>,
expected_resource_commands: Option<Vec<ResourceCommandTrace>>,
expected_output_frames: Option<Vec<OutputFrameTrace>>,
invariant_checks: Vec<NamedInvariantCheck<G, C>>,
}
impl<'harness, G, C> HarnessStep<'harness, G, C> {
pub(crate) fn new(harness: &'harness mut TrellisHarness<G, C>, name: String) -> Self {
Self {
harness,
name,
operations: Vec::new(),
expected_resource_commands: None,
expected_output_frames: None,
invariant_checks: Vec::new(),
}
}
}
impl<'harness, G, C> HarnessStep<'harness, G, C>
where
G: ScenarioTarget<C>,
C: Clone + Debug + PartialEq,
{
pub fn input<T>(mut self, input: InputNode<T>, value: T) -> Self
where
T: Clone + PartialEq + Send + Sync + 'static,
{
self.operations
.push(Box::new(move |tx| tx.set_input(input, value.clone())));
self
}
pub fn operation(
mut self,
operation: impl for<'tx> Fn(&mut Transaction<'tx, C>) -> GraphResult<()> + 'static,
) -> Self {
self.operations.push(Box::new(operation));
self
}
pub fn expect_plan(mut self, command: ResourceCommandTrace) -> Self {
self.expected_resource_commands
.get_or_insert_with(Vec::new)
.push(command);
self
}
pub fn expect_plans(
mut self,
commands: impl IntoIterator<Item = ResourceCommandTrace>,
) -> Self {
self.expected_resource_commands = Some(commands.into_iter().collect());
self
}
pub fn expect_output(mut self, frame: OutputFrameTrace) -> Self {
self.expected_output_frames
.get_or_insert_with(Vec::new)
.push(frame);
self
}
pub fn expect_outputs(mut self, frames: impl IntoIterator<Item = OutputFrameTrace>) -> Self {
self.expected_output_frames = Some(frames.into_iter().collect());
self
}
pub fn check(
mut self,
name: impl Into<String>,
check: impl Fn(&G, &TransactionResult<C>) -> bool + 'static,
) -> Self {
self.invariant_checks.push(NamedInvariantCheck {
name: name.into(),
check: Box::new(check),
});
self
}
pub fn commit(self) -> Result<&'harness mut TrellisHarness<G, C>, ScenarioError> {
self.harness.commit_operations(
&self.name,
&self.operations,
&self.invariant_checks,
self.expected_resource_commands.as_deref(),
self.expected_output_frames.as_deref(),
)?;
Ok(self.harness)
}
}