#![cfg(feature = "boot")]
use std::sync::Arc;
use std::time::Duration;
use a3s_boot::{ModuleRef, Queue, QueueJobOptions, QueueOptions, QueueRetryPolicy};
#[cfg(feature = "sqlite")]
use a3s_flow::SqliteEventStore;
use a3s_flow::{
BootFlowTaskDeduplication, BootFlowTaskManager, BootFlowTaskPolicy, FlowEngine, FlowError,
FlowRuntime, FlowScheduler, FlowTask, RuntimeCommand, StepInvocation, WorkflowInvocation,
WorkflowRunStatus, WorkflowSpec,
};
use async_trait::async_trait;
use chrono::{DateTime, Duration as ChronoDuration, Utc};
use serde_json::json;
fn spec() -> WorkflowSpec {
WorkflowSpec::rust_embedded("boot.workflow", "0.1.0", "tests::boot", "main")
}
struct SleepRuntime;
#[async_trait]
impl FlowRuntime for SleepRuntime {
async fn run_workflow(
&self,
invocation: WorkflowInvocation,
) -> a3s_flow::Result<RuntimeCommand> {
let ctx = invocation.context();
if ctx.wait_completed("sleep") {
return Ok(ctx.complete(json!({ "slept": true })));
}
let resume_at = ctx.input()["resume_at"]
.as_str()
.ok_or_else(|| FlowError::Runtime("missing resume_at".to_string()))?
.parse::<DateTime<Utc>>()
.map_err(|error| FlowError::Runtime(format!("invalid resume_at: {error}")))?;
Ok(ctx.wait_until("sleep", resume_at))
}
async fn run_step(&self, _invocation: StepInvocation) -> a3s_flow::Result<serde_json::Value> {
unreachable!("sleep runtime does not schedule steps")
}
}
#[tokio::test]
async fn boot_task_manager_processes_scheduler_work_through_boot_lifecycle() {
let now = Utc::now();
let engine = FlowEngine::in_memory(Arc::new(SleepRuntime));
let run_id = engine
.start(
spec(),
json!({ "resume_at": (now - ChronoDuration::seconds(1)).to_rfc3339() }),
)
.await
.unwrap();
let queue = Arc::new(Queue::in_process_with_options(
"flow-tests",
QueueOptions::new()
.with_poll_interval(Duration::from_millis(5))
.with_lease_duration(Duration::from_secs(1)),
));
let manager = Arc::new(BootFlowTaskManager::new(engine.clone(), queue.clone()));
manager.register().unwrap();
queue.start(ModuleRef::new()).await.unwrap();
let scheduler = FlowScheduler::new(engine.clone(), manager.clone());
let tick = scheduler.enqueue_due_work(now).await.unwrap();
assert_eq!(tick.enqueued_tasks, 1);
assert_eq!(tick.due_waits, vec![(run_id.clone(), "sleep".to_string())]);
tokio::time::timeout(Duration::from_secs(1), async {
loop {
if engine.snapshot(&run_id).await.unwrap().status == WorkflowRunStatus::Completed {
break;
}
tokio::time::sleep(Duration::from_millis(5)).await;
}
})
.await
.expect("Boot queue should process the Flow task");
queue.shutdown().await.unwrap();
let stats = queue.stats().unwrap();
assert_eq!(stats.completed, 1);
assert_eq!(stats.pending, 0);
assert_eq!(stats.active, 0);
}
#[tokio::test]
async fn boot_task_manager_records_invalid_flow_payload_as_failed_job() {
let engine = FlowEngine::in_memory(Arc::new(SleepRuntime));
let queue = Arc::new(Queue::in_process_with_options(
"flow-invalid-task-tests",
QueueOptions::new().with_poll_interval(Duration::from_millis(5)),
));
let manager = BootFlowTaskManager::new(engine, queue.clone());
manager.register().unwrap();
queue.start(ModuleRef::new()).await.unwrap();
queue
.enqueue_value(manager.job_name(), json!({ "not": "a flow task" }))
.await
.unwrap();
tokio::time::timeout(Duration::from_secs(1), async {
loop {
if queue.stats().unwrap().failed == 1 {
break;
}
tokio::time::sleep(Duration::from_millis(5)).await;
}
})
.await
.expect("Boot queue should retain the failed task");
queue.shutdown().await.unwrap();
let failures = queue.failures().unwrap();
assert_eq!(failures.len(), 1);
assert!(failures[0].message.contains("invalid queued job data"));
}
#[test]
fn boot_task_manager_rejects_an_empty_job_name() {
let engine = FlowEngine::in_memory(Arc::new(SleepRuntime));
let queue = Arc::new(Queue::in_process("flow-tests"));
let result = BootFlowTaskManager::new(engine, queue).with_job_name(" ");
let error = match result {
Ok(_) => panic!("empty Boot job name should fail"),
Err(error) => error,
};
assert!(matches!(
error,
FlowError::InvalidWorkerConfiguration(message) if message.contains("job name")
));
}
#[test]
fn boot_task_policy_maps_typed_options_and_scan_targets() {
let engine = FlowEngine::in_memory(Arc::new(SleepRuntime));
let queue = Arc::new(Queue::in_process("flow-policy-tests"));
let policy = BootFlowTaskPolicy::new()
.with_retry_policy(QueueRetryPolicy::fixed(2, Duration::from_millis(25)))
.with_timeout(Duration::from_secs(30))
.with_max_stalled_count(4)
.remove_on_complete(true)
.remove_on_fail(true)
.with_deduplication(BootFlowTaskDeduplication::UntilTerminalOrTtl(
Duration::from_secs(60),
));
let manager = BootFlowTaskManager::new(engine, queue)
.with_task_policy(policy.clone())
.unwrap();
assert_eq!(manager.task_policy(), &policy);
let first_scan = FlowTask::ResumeDueWaits { now: Utc::now() };
let later_scan = FlowTask::ResumeDueWaits {
now: Utc::now() + ChronoDuration::seconds(1),
};
let retry_scan = FlowTask::ResumeDueRetries { now: Utc::now() };
let first_options = manager.job_options_for(&first_scan);
let later_options = manager.job_options_for(&later_scan);
let retry_options = manager.job_options_for(&retry_scan);
assert_eq!(first_options.retry_policy, policy.retry_policy().clone());
assert_eq!(first_options.timeout, Some(Duration::from_secs(30)));
assert_eq!(first_options.max_stalled_count, 4);
assert!(first_options.remove_on_complete);
assert!(first_options.remove_on_fail);
let first_deduplication = first_options.deduplication.unwrap();
let later_deduplication = later_options.deduplication.unwrap();
let retry_deduplication = retry_options.deduplication.unwrap();
assert_eq!(first_deduplication.id, later_deduplication.id);
assert_ne!(first_deduplication.id, retry_deduplication.id);
assert_eq!(first_deduplication.ttl, Some(Duration::from_secs(60)));
assert!(first_deduplication.keep_last_if_active);
}
#[test]
fn boot_task_deduplication_redacts_tokens_and_preserves_target_semantics() {
let engine = FlowEngine::in_memory(Arc::new(SleepRuntime));
let queue = Arc::new(Queue::in_process("flow-deduplication-tests"));
let manager = BootFlowTaskManager::new(engine, queue)
.with_task_policy(
BootFlowTaskPolicy::new().with_deduplication(BootFlowTaskDeduplication::UntilTerminal),
)
.unwrap();
let token = "public-callback-token-that-must-stay-secret";
let first_resume = manager.job_options_for(&FlowTask::ResumeHookByToken {
token: token.to_string(),
payload: json!({ "decision": "approve" }),
});
let duplicate_resume = manager.job_options_for(&FlowTask::ResumeHookByToken {
token: token.to_string(),
payload: json!({ "decision": "reject" }),
});
let dispose = manager.job_options_for(&FlowTask::DisposeHookByToken {
token: token.to_string(),
});
let drive = manager.job_options_for(&FlowTask::DriveRun {
run_id: "run-1".to_string(),
});
let wait = manager.job_options_for(&FlowTask::ResumeWait {
run_id: "run-1".to_string(),
wait_id: "wait-1".to_string(),
});
let first_resume = first_resume.deduplication.unwrap();
let duplicate_resume = duplicate_resume.deduplication.unwrap();
let dispose = dispose.deduplication.unwrap();
let drive = drive.deduplication.unwrap();
let wait = wait.deduplication.unwrap();
assert_eq!(first_resume.id, duplicate_resume.id);
assert_ne!(first_resume.id, dispose.id);
assert!(!first_resume.id.contains(token));
assert!(first_resume
.id
.starts_with("a3s-flow:resume_hook_by_token:"));
assert!(drive.keep_last_if_active);
assert!(!wait.keep_last_if_active);
}
#[test]
fn boot_task_manager_rejects_a_zero_deduplication_ttl() {
let engine = FlowEngine::in_memory(Arc::new(SleepRuntime));
let queue = Arc::new(Queue::in_process("flow-policy-tests"));
let result = BootFlowTaskManager::new(engine, queue).with_task_policy(
BootFlowTaskPolicy::new().with_deduplication(
BootFlowTaskDeduplication::UntilTerminalOrTtl(Duration::ZERO),
),
);
let error = match result {
Ok(_) => panic!("zero deduplication TTL should fail"),
Err(error) => error,
};
assert!(matches!(
error,
FlowError::InvalidWorkerConfiguration(message)
if message.contains("deduplication TTL")
));
}
#[tokio::test]
async fn boot_task_manager_deduplicates_scans_and_accepts_explicit_job_options() {
let engine = FlowEngine::in_memory(Arc::new(SleepRuntime));
let queue = Arc::new(Queue::in_process("flow-options-tests"));
let manager = BootFlowTaskManager::new(engine, queue.clone())
.with_task_policy(
BootFlowTaskPolicy::new().with_deduplication(BootFlowTaskDeduplication::UntilTerminal),
)
.unwrap();
let first = manager
.enqueue_with_receipt(FlowTask::ResumeDueWaits { now: Utc::now() })
.await
.unwrap();
let duplicate = manager
.enqueue_with_receipt(FlowTask::ResumeDueWaits {
now: Utc::now() + ChronoDuration::seconds(1),
})
.await
.unwrap();
assert_eq!(first.id, duplicate.id);
assert_eq!(queue.stats().unwrap().pending, 1);
let explicit = manager
.enqueue_with_options(
FlowTask::DriveRun {
run_id: "explicit-run".to_string(),
},
QueueJobOptions::new().with_job_id("flow-explicit-job"),
)
.await
.unwrap();
assert_eq!(explicit.id, "flow-explicit-job");
assert_eq!(queue.stats().unwrap().pending, 2);
}
#[cfg(feature = "sqlite")]
#[tokio::test]
async fn boot_task_manager_drives_an_orm_backed_engine_across_restart() {
let directory = tempfile::tempdir().unwrap();
let database_url = format!("sqlite://{}", directory.path().join("flow.db").display());
let now = Utc::now();
let run_id = {
let store = Arc::new(SqliteEventStore::connect(&database_url).await.unwrap());
let engine = FlowEngine::new(store, Arc::new(SleepRuntime));
let run_id = engine
.start(
spec(),
json!({ "resume_at": (now - ChronoDuration::seconds(1)).to_rfc3339() }),
)
.await
.unwrap();
assert_eq!(
engine.snapshot(&run_id).await.unwrap().status,
WorkflowRunStatus::Suspended
);
run_id
};
let store = Arc::new(SqliteEventStore::connect(&database_url).await.unwrap());
let restarted = FlowEngine::new(store, Arc::new(SleepRuntime));
assert_eq!(
restarted.snapshot(&run_id).await.unwrap().status,
WorkflowRunStatus::Suspended
);
let queue = Arc::new(Queue::in_process_with_options(
"flow-orm-tests",
QueueOptions::new().with_poll_interval(Duration::from_millis(5)),
));
let manager = Arc::new(BootFlowTaskManager::new(restarted.clone(), queue.clone()));
manager.register().unwrap();
queue.start(ModuleRef::new()).await.unwrap();
let tick = FlowScheduler::new(restarted.clone(), manager)
.enqueue_due_work(now)
.await
.unwrap();
assert_eq!(tick.enqueued_tasks, 1);
assert_eq!(tick.due_waits, vec![(run_id.clone(), "sleep".to_string())]);
tokio::time::timeout(Duration::from_secs(1), async {
loop {
if restarted.snapshot(&run_id).await.unwrap().status == WorkflowRunStatus::Completed {
break;
}
tokio::time::sleep(Duration::from_millis(5)).await;
}
})
.await
.expect("Boot should finish the ORM-backed Flow task after restart");
queue.shutdown().await.unwrap();
drop(restarted);
let store = Arc::new(SqliteEventStore::connect(&database_url).await.unwrap());
let verified = FlowEngine::new(store, Arc::new(SleepRuntime));
assert_eq!(
verified.snapshot(&run_id).await.unwrap().status,
WorkflowRunStatus::Completed
);
}