use std::collections::BTreeMap;
use crate::activity::{ActivityRegistry, HandlerFuture};
use crate::error::WorkerError;
use super::action::{ShellAction, ShellOutcome};
pub type DeclaredArguments = BTreeMap<String, serde_json::Value>;
pub trait DeclaredCommands: Sized {
fn register_declared_command(
self,
activity_type: impl Into<String>,
action: ShellAction,
) -> Result<Self, WorkerError>;
}
impl DeclaredCommands for ActivityRegistry {
fn register_declared_command(
self,
activity_type: impl Into<String>,
action: ShellAction,
) -> Result<Self, WorkerError> {
self.register_activity(
activity_type,
move |arguments: DeclaredArguments, context| -> HandlerFuture<'_, ShellOutcome> {
let action = action.clone();
Box::pin(async move { action.run(&arguments, context).await })
},
)
}
}
#[cfg(test)]
mod tests {
use super::{DeclaredArguments, DeclaredCommands};
use crate::activity::ActivityRegistry;
use crate::runtime::{ActivityDispatcher, DispatchOutcome, decode_payload, encode_payload};
use crate::shell::action::{FailureMode, ShellAction, ShellOutcome};
use aion_core::{ActivityErrorKind, ActivityId, ContentType, WorkflowId};
type TestResult = Result<(), Box<dyn std::error::Error>>;
fn arguments(pairs: &[(&str, serde_json::Value)]) -> DeclaredArguments {
pairs
.iter()
.map(|(name, value)| ((*name).to_owned(), value.clone()))
.collect()
}
async fn dispatch(
registry: &ActivityRegistry,
activity_type: &str,
call: &DeclaredArguments,
) -> Result<DispatchOutcome, Box<dyn std::error::Error>> {
let workflow_id = WorkflowId::new_v4();
let run_id = aion_core::RunId::new_v4();
let activity_id = ActivityId::from_sequence_position(1);
let task = crate::protocol::ActivityTask {
workflow_id: workflow_id.clone(),
activity_id: activity_id.clone(),
run_id: run_id.clone(),
activity_type: activity_type.to_owned(),
attempt: 1,
completion_token: "declared-command-test".to_owned(),
idempotency_key: "declared-command-test".to_owned(),
input: encode_payload(call)?,
labels: std::collections::BTreeMap::new(),
};
let (context, cancellation) =
crate::ActivityContext::new(workflow_id, run_id, activity_id, 1);
drop(cancellation);
Ok(registry.dispatch(task, context).await?)
}
#[tokio::test]
async fn a_declared_command_is_served_through_the_normal_dispatch_path() -> TestResult {
let registry = ActivityRegistry::new()
.register_declared_command("greet", ShellAction::new("echo {{greeting}}")?)?;
let outcome = dispatch(
®istry,
"greet",
&arguments(&[("greeting", serde_json::json!("hello there"))]),
)
.await?;
let DispatchOutcome::Completed { output } = outcome else {
return Err("a succeeding declared command must complete".into());
};
assert_eq!(output.content_type(), &ContentType::Json);
let result: ShellOutcome = decode_payload(&output)?;
assert_eq!(result.stdout, "hello there");
assert_eq!(result.exit_code, 0);
Ok(())
}
#[tokio::test]
async fn a_failing_declared_command_reports_its_classification() -> TestResult {
let registry = ActivityRegistry::new().register_declared_command(
"fails",
ShellAction::new("sh -c 'exit 4'")?.with_failure_mode(FailureMode::Terminal),
)?;
let outcome = dispatch(®istry, "fails", &arguments(&[])).await?;
let DispatchOutcome::Failed { failure } = outcome else {
return Err("a non-zero exit must fail the activity".into());
};
assert_eq!(failure.kind, ActivityErrorKind::Terminal);
assert!(
!failure.is_retryable(),
"a terminal declared-command failure must not be retryable"
);
Ok(())
}
#[tokio::test]
async fn a_declared_command_cannot_displace_a_registered_activity() -> TestResult {
let registry = ActivityRegistry::new()
.register_declared_command("shared", ShellAction::new("echo first")?)?;
let Err(error) =
registry.register_declared_command("shared", ShellAction::new("echo second")?)
else {
return Err("a duplicate activity type must be refused, not silently replaced".into());
};
assert!(
error.to_string().contains("shared"),
"the refusal must name the colliding activity type: {error}"
);
Ok(())
}
#[tokio::test]
async fn an_unnamed_argument_is_ignored_rather_than_smuggled_into_the_command() -> TestResult {
let registry = ActivityRegistry::new()
.register_declared_command("echo", ShellAction::new("echo {{wanted}}")?)?;
let outcome = dispatch(
®istry,
"echo",
&arguments(&[
("wanted", serde_json::json!("kept")),
("ignored", serde_json::json!("; echo PWNED")),
]),
)
.await?;
let DispatchOutcome::Completed { output } = outcome else {
return Err("the command must complete".into());
};
let result: ShellOutcome = decode_payload(&output)?;
assert_eq!(result.stdout, "kept");
Ok(())
}
}