mecha10-runtime 0.7.0

use mecha10_runtime::supervisor::*;

use anyhow::Result;
use async_trait::async_trait;
use mecha10_runtime::node::NodeRunner;
use mecha10_runtime::{HealthChecker, RestartPolicy, ShutdownHandle};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;

struct TestNode {
    name: String,
    run_count: Arc<AtomicUsize>,
    should_fail: Arc<AtomicBool>,
    fail_count: Arc<AtomicUsize>,
}

#[async_trait]
impl NodeRunner for TestNode {
    fn name(&self) -> &str {
        &self.name
    }

    async fn run(&mut self) -> Result<()> {
        self.run_count.fetch_add(1, Ordering::SeqCst);

        if self.should_fail.load(Ordering::SeqCst) {
            self.fail_count.fetch_add(1, Ordering::SeqCst);
            anyhow::bail!("Intentional test failure");
        }

        tokio::time::sleep(Duration::from_millis(100)).await;
        Ok(())
    }
}

#[tokio::test]
async fn test_successful_node() {
    let health_checker = Arc::new(HealthChecker::new(Duration::from_secs(30)));
    let shutdown = ShutdownHandle::new();
    let supervisor = Supervisor::new(health_checker.clone(), shutdown.clone(), RestartPolicy::Never);

    let run_count = Arc::new(AtomicUsize::new(0));
    let node = TestNode {
        name: "test_node".to_string(),
        run_count: run_count.clone(),
        should_fail: Arc::new(AtomicBool::new(false)),
        fail_count: Arc::new(AtomicUsize::new(0)),
    };

    let handle = supervisor.launch_node(Box::new(node));
    assert_eq!(handle.name(), "test_node");

    handle.wait().await.unwrap();

    assert_eq!(run_count.load(Ordering::SeqCst), 1);
}

#[tokio::test]
async fn test_node_restart_on_failure() {
    let health_checker = Arc::new(HealthChecker::new(Duration::from_secs(30)));
    let shutdown = ShutdownHandle::new();
    let supervisor = Supervisor::new(
        health_checker.clone(),
        shutdown.clone(),
        RestartPolicy::OnFailure {
            max_retries: 3,
            backoff: Duration::from_millis(10),
        },
    );

    let run_count = Arc::new(AtomicUsize::new(0));
    let fail_count = Arc::new(AtomicUsize::new(0));
    let should_fail = Arc::new(AtomicBool::new(true));

    let node = TestNode {
        name: "failing_node".to_string(),
        run_count: run_count.clone(),
        should_fail: should_fail.clone(),
        fail_count: fail_count.clone(),
    };

    let handle = supervisor.launch_node(Box::new(node));

    tokio::time::sleep(Duration::from_millis(300)).await;

    assert!(
        fail_count.load(Ordering::SeqCst) >= 3,
        "Should have at least 3 failures"
    );

    handle.abort();
}

#[tokio::test]
async fn test_graceful_shutdown() {
    let health_checker = Arc::new(HealthChecker::new(Duration::from_secs(30)));
    let shutdown = ShutdownHandle::new();
    let supervisor = Supervisor::new(health_checker.clone(), shutdown.clone(), RestartPolicy::Never);

    let run_count = Arc::new(AtomicUsize::new(0));
    let node = TestNode {
        name: "shutdown_node".to_string(),
        run_count: run_count.clone(),
        should_fail: Arc::new(AtomicBool::new(false)),
        fail_count: Arc::new(AtomicUsize::new(0)),
    };

    let handle = supervisor.launch_node(Box::new(node));

    tokio::time::sleep(Duration::from_millis(10)).await;

    shutdown.shutdown();

    handle.wait().await.unwrap();

    assert_eq!(run_count.load(Ordering::SeqCst), 1);
}