autopilot/crates/mission_executor/tests/state_machine.rs

//! AZ-648 acceptance criteria.
//!
//! AC-1 / AC-2 — happy-path multirotor / fixed-wing flow with a fake
//! driver. The driver stands in for the SITL conformance target; the
//! state graph and event publication are what the AC asserts.
//!
//! AC-3 — bounded retry on mission-upload rejection: first attempt
//! rejected, second succeeds, FSM proceeds.
//!
//! AC-4 — cap exhaustion: all 3 default attempts rejected → FSM pauses,
//! health → red, transition event published, no transition past
//! `MissionUploaded`.

use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use std::time::Duration;

use async_trait::async_trait;
use mission_executor::{
    DriverError, MissionDriver, MissionExecutor, MissionExecutorConfig, MissionState, StepOutcome,
    Telemetry, TransitionKey, Variant, DEFAULT_RETRY_CAP,
};
use shared::health::HealthLevel;
use shared::models::mission::{MavCommand, MavFrame, MissionWaypoint};
use tokio::sync::watch;
use tokio::time::timeout;

/// Configurable in-memory driver. Counts every action and can be
/// scripted to reject the next N upload calls.
struct ScriptedDriver {
    arm_calls: AtomicU32,
    takeoff_calls: AtomicU32,
    upload_calls: AtomicU32,
    set_auto_calls: AtomicU32,
    post_flight_calls: AtomicU32,
    reject_first_n_uploads: AtomicU32,
}

impl ScriptedDriver {
    fn new() -> Arc<Self> {
        Arc::new(Self {
            arm_calls: AtomicU32::new(0),
            takeoff_calls: AtomicU32::new(0),
            upload_calls: AtomicU32::new(0),
            set_auto_calls: AtomicU32::new(0),
            post_flight_calls: AtomicU32::new(0),
            reject_first_n_uploads: AtomicU32::new(0),
        })
    }

    fn reject_next_uploads(&self, n: u32) {
        self.reject_first_n_uploads.store(n, Ordering::SeqCst);
    }
}

#[async_trait]
impl MissionDriver for ScriptedDriver {
    async fn arm(&self) -> Result<(), DriverError> {
        self.arm_calls.fetch_add(1, Ordering::SeqCst);
        Ok(())
    }

    async fn takeoff(&self, _altitude_m: f32) -> Result<(), DriverError> {
        self.takeoff_calls.fetch_add(1, Ordering::SeqCst);
        Ok(())
    }

    async fn upload_mission(&self, _items: &[MissionWaypoint]) -> Result<(), DriverError> {
        self.upload_calls.fetch_add(1, Ordering::SeqCst);
        loop {
            let remaining = self.reject_first_n_uploads.load(Ordering::SeqCst);
            if remaining == 0 {
                return Ok(());
            }
            if self
                .reject_first_n_uploads
                .compare_exchange(remaining, remaining - 1, Ordering::SeqCst, Ordering::SeqCst)
                .is_ok()
            {
                return Err(DriverError::Rejected("simulated".into()));
            }
        }
    }

    async fn set_auto_mode(&self) -> Result<(), DriverError> {
        self.set_auto_calls.fetch_add(1, Ordering::SeqCst);
        Ok(())
    }

    async fn post_flight_sync(&self) -> Result<(), DriverError> {
        self.post_flight_calls.fetch_add(1, Ordering::SeqCst);
        Ok(())
    }
}

/// Drive `telemetry_rx` forward through a script while polling the
/// executor until `target` is reached. The script entries are applied
/// in order — each one waits up to `step_timeout` for the FSM to
/// advance past `prev`, then publishes the next telemetry snapshot.
async fn await_state(
    handle: &mission_executor::MissionExecutorHandle,
    target: MissionState,
    overall: Duration,
) {
    let res = timeout(overall, async {
        loop {
            if handle.state().await == target {
                return;
            }
            tokio::time::sleep(Duration::from_millis(5)).await;
        }
    })
    .await;
    if res.is_err() {
        let actual = handle.state().await;
        panic!("FSM did not reach {target:?}; stuck at {actual:?}");
    }
}

fn one_waypoint() -> Vec<MissionWaypoint> {
    vec![MissionWaypoint {
        seq: 0,
        frame: MavFrame::MavFrameGlobalRelativeAlt,
        command: MavCommand::MavCmdNavWaypoint,
        current: true,
        auto_continue: true,
        param_1: 0.0,
        param_2: 0.0,
        param_3: 0.0,
        param_4: 0.0,
        lat_deg_e7: 0,
        lon_deg_e7: 0,
        alt_m: 50.0,
    }]
}

#[tokio::test]
async fn ac1_multirotor_happy_path_reaches_done() {
    // Arrange
    let driver = ScriptedDriver::new();
    let exec = MissionExecutor::new(MissionExecutorConfig {
        tick_interval: Duration::from_millis(5),
        ..MissionExecutorConfig::multirotor(10.0)
    });
    let (tx, rx) = watch::channel(Telemetry::default());
    let (handle, join) = exec.run(driver.clone(), one_waypoint(), rx);
    let mut events = handle.subscribe();

    // Act / Assert — step the telemetry script.
    tx.send(Telemetry {
        link_up: true,
        ..Telemetry::default()
    })
    .unwrap();
    await_state(&handle, MissionState::Connected, Duration::from_secs(1)).await;

    tx.send(Telemetry {
        link_up: true,
        health_ok: true,
        ..Telemetry::default()
    })
    .unwrap();
    await_state(&handle, MissionState::HealthOk, Duration::from_secs(1)).await;

    tx.send(Telemetry {
        link_up: true,
        health_ok: true,
        bit_ok: true,
        ..Telemetry::default()
    })
    .unwrap();
    await_state(&handle, MissionState::BitOk, Duration::from_secs(1)).await;
    await_state(&handle, MissionState::Armed, Duration::from_secs(1)).await;

    tx.send(Telemetry {
        link_up: true,
        health_ok: true,
        bit_ok: true,
        armed: true,
        ..Telemetry::default()
    })
    .unwrap();
    await_state(&handle, MissionState::TakeOff, Duration::from_secs(1)).await;

    tx.send(Telemetry {
        link_up: true,
        health_ok: true,
        bit_ok: true,
        armed: true,
        takeoff_complete: true,
        ..Telemetry::default()
    })
    .unwrap();
    await_state(
        &handle,
        MissionState::MissionUploaded,
        Duration::from_secs(1),
    )
    .await;

    tx.send(Telemetry {
        link_up: true,
        health_ok: true,
        bit_ok: true,
        armed: true,
        takeoff_complete: true,
        flight_mode_auto: true,
        ..Telemetry::default()
    })
    .unwrap();
    await_state(&handle, MissionState::FlyMission, Duration::from_secs(1)).await;

    tx.send(Telemetry {
        link_up: true,
        health_ok: true,
        bit_ok: true,
        armed: true,
        takeoff_complete: true,
        flight_mode_auto: true,
        mission_reached_final: true,
        ..Telemetry::default()
    })
    .unwrap();
    await_state(&handle, MissionState::Land, Duration::from_secs(1)).await;

    tx.send(Telemetry {
        link_up: true,
        health_ok: true,
        bit_ok: true,
        armed: false,
        takeoff_complete: true,
        flight_mode_auto: true,
        mission_reached_final: true,
        landed_disarmed: true,
    })
    .unwrap();
    await_state(
        &handle,
        MissionState::PostFlightSync,
        Duration::from_secs(1),
    )
    .await;
    await_state(&handle, MissionState::Done, Duration::from_secs(1)).await;

    // Assert — health is green at Done, driver saw exactly one of each action.
    let health = handle.health().await;
    assert_eq!(health.level, HealthLevel::Green);
    assert_eq!(driver.arm_calls.load(Ordering::SeqCst), 1);
    assert_eq!(driver.takeoff_calls.load(Ordering::SeqCst), 1);
    assert_eq!(driver.upload_calls.load(Ordering::SeqCst), 1);
    assert_eq!(driver.post_flight_calls.load(Ordering::SeqCst), 1);
    // No fixed-wing action on a multirotor flow.
    assert_eq!(driver.set_auto_calls.load(Ordering::SeqCst), 0);

    // Drain the event stream — count distinct transitions; we expect
    // every state above to appear in order.
    let mut observed = Vec::new();
    while let Ok(evt) = events.try_recv() {
        observed.push((evt.from, evt.to));
    }
    assert!(observed.contains(&(MissionState::Disconnected, MissionState::Connected)));
    assert!(observed.contains(&(MissionState::PostFlightSync, MissionState::Done)));

    let _ = join.await;
}

#[tokio::test]
async fn ac2_fixed_wing_happy_path_reaches_done() {
    // Arrange — fixed-wing skips Armed/TakeOff. The operator sets AUTO
    // externally; we model that by flipping `flight_mode_auto` while
    // the FSM is parked in WaitAuto.
    let driver = ScriptedDriver::new();
    let exec = MissionExecutor::new(MissionExecutorConfig {
        tick_interval: Duration::from_millis(5),
        ..MissionExecutorConfig::fixed_wing()
    });
    let (tx, rx) = watch::channel(Telemetry::default());
    let (handle, join) = exec.run(driver.clone(), one_waypoint(), rx);

    // Act / Assert
    tx.send(Telemetry {
        link_up: true,
        health_ok: true,
        bit_ok: true,
        ..Telemetry::default()
    })
    .unwrap();
    await_state(&handle, MissionState::WaitAuto, Duration::from_secs(2)).await;
    assert_eq!(driver.arm_calls.load(Ordering::SeqCst), 0);
    assert_eq!(driver.takeoff_calls.load(Ordering::SeqCst), 0);
    assert_eq!(driver.upload_calls.load(Ordering::SeqCst), 1);

    tx.send(Telemetry {
        link_up: true,
        health_ok: true,
        bit_ok: true,
        flight_mode_auto: true,
        ..Telemetry::default()
    })
    .unwrap();
    await_state(&handle, MissionState::FlyMission, Duration::from_secs(1)).await;

    tx.send(Telemetry {
        link_up: true,
        health_ok: true,
        bit_ok: true,
        flight_mode_auto: true,
        mission_reached_final: true,
        ..Telemetry::default()
    })
    .unwrap();
    await_state(&handle, MissionState::Land, Duration::from_secs(1)).await;

    tx.send(Telemetry {
        link_up: true,
        health_ok: true,
        bit_ok: true,
        flight_mode_auto: true,
        mission_reached_final: true,
        landed_disarmed: true,
        ..Telemetry::default()
    })
    .unwrap();
    await_state(&handle, MissionState::Done, Duration::from_secs(2)).await;

    assert_eq!(handle.health().await.level, HealthLevel::Green);
    let _ = join.await;
}

#[tokio::test]
async fn ac3_bounded_retry_then_success() {
    // Arrange — reject the first upload attempt, accept the second.
    let driver = ScriptedDriver::new();
    driver.reject_next_uploads(1);
    let exec = MissionExecutor::new(MissionExecutorConfig {
        tick_interval: Duration::from_millis(5),
        ..MissionExecutorConfig::fixed_wing()
    });
    let (tx, rx) = watch::channel(Telemetry {
        link_up: true,
        health_ok: true,
        bit_ok: true,
        ..Telemetry::default()
    });
    let (handle, join) = exec.run(driver.clone(), one_waypoint(), rx);

    // Act
    await_state(
        &handle,
        MissionState::MissionUploaded,
        Duration::from_secs(2),
    )
    .await;

    // Assert — driver was called twice (one rejected + one accepted),
    // retry counter for that transition is 1, FSM proceeded.
    assert_eq!(driver.upload_calls.load(Ordering::SeqCst), 2);
    let retry = handle
        .retry_count(TransitionKey::new(
            MissionState::BitOk,
            MissionState::MissionUploaded,
        ))
        .await;
    // Successful advance clears the retry counter (per FSM design —
    // a fresh transition starts with a clean budget). The proof that
    // a retry happened is the double upload_calls.
    assert_eq!(retry, 0);
    assert!(matches!(
        handle.state().await,
        MissionState::WaitAuto | MissionState::MissionUploaded
    ));

    // Cleanup — drive to Done so the task exits cleanly.
    tx.send(Telemetry {
        link_up: true,
        health_ok: true,
        bit_ok: true,
        flight_mode_auto: true,
        mission_reached_final: true,
        landed_disarmed: true,
        ..Telemetry::default()
    })
    .unwrap();
    await_state(&handle, MissionState::Done, Duration::from_secs(2)).await;
    let _ = join.await;
}

#[tokio::test]
async fn ac4_cap_exhaustion_pauses_and_flips_health_red() {
    // Arrange — reject every upload attempt. With the default cap of 3
    // the FSM should pause on the 3rd failure.
    let driver = ScriptedDriver::new();
    driver.reject_next_uploads(u32::MAX);
    let exec = MissionExecutor::new(MissionExecutorConfig {
        tick_interval: Duration::from_millis(5),
        ..MissionExecutorConfig::fixed_wing()
    });
    let (_tx, rx) = watch::channel(Telemetry {
        link_up: true,
        health_ok: true,
        bit_ok: true,
        ..Telemetry::default()
    });
    let (handle, join) = exec.run(driver.clone(), one_waypoint(), rx);
    let mut events = handle.subscribe();

    // Act
    await_state(&handle, MissionState::Paused, Duration::from_secs(2)).await;

    // Assert
    assert_eq!(
        driver.upload_calls.load(Ordering::SeqCst),
        DEFAULT_RETRY_CAP,
        "driver should have been called exactly cap times"
    );
    let health = handle.health().await;
    assert_eq!(health.level, HealthLevel::Red);
    let reason = handle.paused_reason().await.expect("paused reason");
    assert!(
        reason.contains("UploadMission") || reason.contains("cap-exhausted"),
        "reason should mention the failed action: got {reason}"
    );

    // A `→ Paused` event must have been published.
    let mut saw_pause_event = false;
    while let Ok(evt) = events.try_recv() {
        if evt.to == MissionState::Paused {
            saw_pause_event = true;
            assert_eq!(evt.variant, Variant::FixedWing);
            break;
        }
    }
    assert!(
        saw_pause_event,
        "expected a transition event with to=Paused"
    );

    // FSM does not advance past MissionUploaded — we never reached it.
    // Task exits because the state is terminal.
    let final_state = handle.state().await;
    assert_eq!(final_state, MissionState::Paused);
    let final_outcome = StepOutcome::Paused {
        reason: reason.clone(),
    };
    assert!(matches!(final_outcome, StepOutcome::Paused { .. }));

    let _ = join.await;
}