use std::sync::atomic::{AtomicUsize, Ordering};
use async_trait::async_trait;
use graph_flow::{Context, NextAction, Task, TaskResult, error::Result};
use guardflow::Guard;
use guardflow::graphflow::GuardedTask;
use guardflow::validators::MinLength;
struct FlakyTask {
attempts: AtomicUsize,
succeeds_on: usize,
}
#[async_trait]
impl Task for FlakyTask {
fn id(&self) -> &str {
"flaky"
}
async fn run(&self, _context: Context) -> Result<TaskResult> {
let attempt = self.attempts.fetch_add(1, Ordering::SeqCst) + 1;
let response = if attempt >= self.succeeds_on {
"this response is long enough to pass"
} else {
"short"
};
Ok(TaskResult::new(
Some(response.to_string()),
NextAction::Continue,
))
}
}
#[tokio::test]
async fn guarded_task_retries_until_validator_passes() {
let inner = FlakyTask {
attempts: AtomicUsize::new(0),
succeeds_on: 3,
};
let guard = Guard::new().with(MinLength(20));
let guarded = GuardedTask::new(inner, guard).with_max_attempts(5);
let result = guarded.run(Context::new()).await.unwrap();
assert_eq!(
result.response,
Some("this response is long enough to pass".to_string())
);
}
#[tokio::test]
async fn guarded_task_gives_up_after_max_attempts() {
let inner = FlakyTask {
attempts: AtomicUsize::new(0),
succeeds_on: 100,
};
let guard = Guard::new().with(MinLength(20));
let guarded = GuardedTask::new(inner, guard).with_max_attempts(2);
let result = guarded.run(Context::new()).await.unwrap();
assert_eq!(result.response, Some("short".to_string()));
}