pe-graph 0.1.0

//! Integration tests for RetryPolicy wiring into the Pregel BSP engine.
//!
//! Verifies that when a RetryPolicy is set on GraphConfig, failed nodes
//! with retryable errors are retried before the engine gives up.

use pe_core::error::PeError;
use pe_core::node::{NodeContext, NodeFn, NodeFuture, NodeResult};
use pe_core::state::{State, StateUpdate};
use pe_core::types::{END, START};
use pe_graph::{ExecutionOutcome, GraphConfig, RetryPolicy, StateGraph};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};
use std::time::Duration;

// ── Test Types ────────────────────────────────────────────────────────

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
struct TestState {
    messages: Vec<String>,
    counter: u32,
    thread_id: String,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
struct TestUpdate {
    messages: Option<Vec<String>>,
    counter: Option<u32>,
}

impl StateUpdate for TestUpdate {}

impl State for TestState {
    type Update = TestUpdate;
    fn apply(&mut self, update: TestUpdate) {
        if let Some(msgs) = update.messages {
            self.messages.extend(msgs);
        }
        if let Some(c) = update.counter {
            self.counter = c;
        }
    }
}

impl TestState {
    fn new() -> Self {
        Self {
            messages: Vec::new(),
            counter: 0,
            thread_id: "test".into(),
        }
    }
}

// ── Test Nodes ────────────────────────────────────────────────────────

/// A node that fails with a retryable error for the first N calls,
/// then succeeds.
struct FailThenSucceedNode {
    name: &'static str,
    fail_count: u32,
    attempts: Arc<AtomicU32>,
}

impl FailThenSucceedNode {
    fn new(name: &'static str, fail_count: u32, attempts: Arc<AtomicU32>) -> Self {
        Self {
            name,
            fail_count,
            attempts,
        }
    }
}

impl NodeFn<TestState> for FailThenSucceedNode {
    fn call(&self, _state: &TestState, _ctx: &NodeContext) -> NodeFuture<TestUpdate> {
        let count = self.attempts.fetch_add(1, Ordering::SeqCst);
        let fail_count = self.fail_count;
        Box::pin(async move {
            if count < fail_count {
                NodeResult::Error(PeError::Timeout { seconds: 1.0 })
            } else {
                NodeResult::Update(TestUpdate {
                    messages: Some(vec!["recovered".into()]),
                    counter: None,
                })
            }
        })
    }

    fn name(&self) -> &str {
        self.name
    }
}

/// A node that always fails with a retryable error.
struct AlwaysFailRetryableNode {
    name: &'static str,
    attempts: Arc<AtomicU32>,
}

impl AlwaysFailRetryableNode {
    fn new(name: &'static str, attempts: Arc<AtomicU32>) -> Self {
        Self { name, attempts }
    }
}

impl NodeFn<TestState> for AlwaysFailRetryableNode {
    fn call(&self, _state: &TestState, _ctx: &NodeContext) -> NodeFuture<TestUpdate> {
        self.attempts.fetch_add(1, Ordering::SeqCst);
        Box::pin(async { NodeResult::Error(PeError::Timeout { seconds: 1.0 }) })
    }

    fn name(&self) -> &str {
        self.name
    }
}

/// A node that fails with a non-retryable error.
struct NonRetryableFailNode {
    name: &'static str,
    attempts: Arc<AtomicU32>,
}

impl NonRetryableFailNode {
    fn new(name: &'static str, attempts: Arc<AtomicU32>) -> Self {
        Self { name, attempts }
    }
}

impl NodeFn<TestState> for NonRetryableFailNode {
    fn call(&self, _state: &TestState, _ctx: &NodeContext) -> NodeFuture<TestUpdate> {
        self.attempts.fetch_add(1, Ordering::SeqCst);
        Box::pin(async {
            NodeResult::Error(PeError::PermissionDenied {
                action: "test".into(),
            })
        })
    }

    fn name(&self) -> &str {
        self.name
    }
}

/// Simple node that appends a message.
struct AppendNode {
    name: &'static str,
    msg: &'static str,
}

impl NodeFn<TestState> for AppendNode {
    fn call(&self, _state: &TestState, _ctx: &NodeContext) -> NodeFuture<TestUpdate> {
        let msg = self.msg.to_string();
        Box::pin(async move {
            NodeResult::Update(TestUpdate {
                messages: Some(vec![msg]),
                counter: None,
            })
        })
    }

    fn name(&self) -> &str {
        self.name
    }
}

// ── Tests ─────────────────────────────────────────────────────────────

#[tokio::test]
async fn test_retry_recovers_after_transient_failure() {
    // Node fails once (retryable), then succeeds on retry.
    let attempts = Arc::new(AtomicU32::new(0));
    let node = FailThenSucceedNode::new("flaky", 1, Arc::clone(&attempts));

    let graph = StateGraph::new()
        .add_node("flaky", node)
        .add_edge(START, "flaky")
        .add_edge("flaky", END)
        .compile()
        .unwrap();

    let config = GraphConfig::default().with_retry_policy(RetryPolicy {
        max_attempts: 3,
        initial_interval: Duration::from_millis(1),
        jitter: false,
        ..Default::default()
    });

    let outcome = graph.invoke(TestState::new(), config).await.unwrap();
    match outcome {
        ExecutionOutcome::Completed(state) => {
            assert_eq!(state.messages, vec!["recovered"]);
        }
        other => panic!("Expected Completed, got {:?}", other),
    }

    // 1 initial + 1 retry = 2 total calls
    assert_eq!(attempts.load(Ordering::SeqCst), 2);
}

#[tokio::test]
async fn test_retry_exhausts_all_attempts_then_errors() {
    // Node always fails. With max_attempts=2, total calls = 1 initial + 2 retries = 3.
    let attempts = Arc::new(AtomicU32::new(0));
    let node = AlwaysFailRetryableNode::new("always_fail", Arc::clone(&attempts));

    let graph = StateGraph::new()
        .add_node("always_fail", node)
        .add_edge(START, "always_fail")
        .add_edge("always_fail", END)
        .compile()
        .unwrap();

    let config = GraphConfig::default().with_retry_policy(RetryPolicy {
        max_attempts: 2,
        initial_interval: Duration::from_millis(1),
        jitter: false,
        ..Default::default()
    });

    let err = graph.invoke(TestState::new(), config).await.unwrap_err();
    assert!(matches!(err, PeError::Timeout { .. }));

    // 1 initial call (in execute_parallel) + 2 retries = 3 total
    assert_eq!(attempts.load(Ordering::SeqCst), 3);
}

#[tokio::test]
async fn test_non_retryable_error_not_retried() {
    // Non-retryable errors should fail immediately even with a retry policy.
    let attempts = Arc::new(AtomicU32::new(0));
    let node = NonRetryableFailNode::new("perm_fail", Arc::clone(&attempts));

    let graph = StateGraph::new()
        .add_node("perm_fail", node)
        .add_edge(START, "perm_fail")
        .add_edge("perm_fail", END)
        .compile()
        .unwrap();

    let config = GraphConfig::default().with_retry_policy(RetryPolicy {
        max_attempts: 5,
        initial_interval: Duration::from_millis(1),
        jitter: false,
        ..Default::default()
    });

    let err = graph.invoke(TestState::new(), config).await.unwrap_err();
    assert!(matches!(err, PeError::PermissionDenied { .. }));

    // Only 1 call — no retries for non-retryable errors
    assert_eq!(attempts.load(Ordering::SeqCst), 1);
}

#[tokio::test]
async fn test_no_retry_policy_errors_immediately() {
    // Without a retry policy, retryable errors still fail immediately.
    let attempts = Arc::new(AtomicU32::new(0));
    let node = AlwaysFailRetryableNode::new("fail_no_policy", Arc::clone(&attempts));

    let graph = StateGraph::new()
        .add_node("fail_no_policy", node)
        .add_edge(START, "fail_no_policy")
        .add_edge("fail_no_policy", END)
        .compile()
        .unwrap();

    let config = GraphConfig::default(); // No retry policy

    let err = graph.invoke(TestState::new(), config).await.unwrap_err();
    assert!(matches!(err, PeError::Timeout { .. }));
    assert_eq!(attempts.load(Ordering::SeqCst), 1);
}

#[tokio::test]
async fn test_retry_succeeds_on_last_attempt() {
    // Node fails exactly max_attempts times, then succeeds on the last retry.
    let attempts = Arc::new(AtomicU32::new(0));
    let node = FailThenSucceedNode::new("last_chance", 3, Arc::clone(&attempts));

    let graph = StateGraph::new()
        .add_node("last_chance", node)
        .add_edge(START, "last_chance")
        .add_edge("last_chance", END)
        .compile()
        .unwrap();

    let config = GraphConfig::default().with_retry_policy(RetryPolicy {
        max_attempts: 3,
        initial_interval: Duration::from_millis(1),
        jitter: false,
        ..Default::default()
    });

    let outcome = graph.invoke(TestState::new(), config).await.unwrap();
    match outcome {
        ExecutionOutcome::Completed(state) => {
            assert_eq!(state.messages, vec!["recovered"]);
        }
        other => panic!("Expected Completed, got {:?}", other),
    }

    // 1 initial + 3 retries = 4 total (succeeds on 4th call)
    assert_eq!(attempts.load(Ordering::SeqCst), 4);
}

#[tokio::test]
async fn test_retry_only_retries_failed_node_in_multi_node_graph() {
    // In a graph with multiple sequential nodes, only the failing node
    // should be retried. Other nodes should run normally.
    let flaky_attempts = Arc::new(AtomicU32::new(0));

    let graph = StateGraph::new()
        .add_node(
            "first",
            AppendNode {
                name: "first",
                msg: "step1",
            },
        )
        .add_node(
            "flaky",
            FailThenSucceedNode::new("flaky", 1, Arc::clone(&flaky_attempts)),
        )
        .add_edge(START, "first")
        .add_edge("first", "flaky")
        .add_edge("flaky", END)
        .compile()
        .unwrap();

    let config = GraphConfig::default().with_retry_policy(RetryPolicy {
        max_attempts: 3,
        initial_interval: Duration::from_millis(1),
        jitter: false,
        ..Default::default()
    });

    let outcome = graph.invoke(TestState::new(), config).await.unwrap();
    match outcome {
        ExecutionOutcome::Completed(state) => {
            // first node's message + flaky node's recovery message
            assert_eq!(state.messages, vec!["step1", "recovered"]);
        }
        other => panic!("Expected Completed, got {:?}", other),
    }

    // flaky node: 1 initial (fail) + 1 retry (success) = 2
    assert_eq!(flaky_attempts.load(Ordering::SeqCst), 2);
}