harmony/harmony_cli/src/lib.rs

use clap::Parser;
use clap::builder::ArgPredicate;
use harmony::instrumentation;
use harmony::inventory::Inventory;
use harmony::maestro::Maestro;
use harmony::{score::Score, topology::Topology};
use inquire::Confirm;
use log::debug;

pub mod cli_logger; // FIXME: Don't make me pub
pub mod progress;
pub mod theme;

#[cfg(feature = "tui")]
#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
pub struct Args {
    /// Run score(s) without prompt
    #[arg(short, long, default_value_t = false, conflicts_with = "interactive")]
    pub yes: bool,

    /// Filter query
    #[arg(short, long, conflicts_with = "interactive")]
    pub filter: Option<String>,

    /// Run interactive TUI or not
    #[arg(short, long, default_value_t = false)]
    pub interactive: bool,

    /// Run all or nth, defaults to all
    #[arg(
        short,
        long,
        default_value_t = true,
        default_value_if("number", ArgPredicate::IsPresent, "false"),
        conflicts_with = "number",
        conflicts_with = "interactive"
    )]
    pub all: bool,

    /// Run nth matching, zero indexed
    #[arg(short, long, default_value_t = 0, conflicts_with = "interactive")]
    pub number: usize,

    /// list scores, will also be affected by run filter
    #[arg(short, long, default_value_t = false, conflicts_with = "interactive")]
    pub list: bool,
}

fn maestro_scores_filter<T: Topology>(
    maestro: &harmony::maestro::Maestro<T>,
    all: bool,
    filter: Option<String>,
    number: usize,
) -> Vec<Box<dyn Score<T>>> {
    let scores = maestro.scores();
    let scores_read = scores.read().expect("Should be able to read scores");
    let mut scores_vec: Vec<Box<dyn Score<T>>> = match filter {
        Some(f) => scores_read
            .iter()
            .filter(|s| s.name().contains(&f))
            .map(|s| s.clone_box())
            .collect(),
        None => scores_read.iter().map(|s| s.clone_box()).collect(),
    };

    if !all {
        let score = scores_vec.get(number);
        match score {
            Some(s) => scores_vec = vec![s.clone_box()],
            None => return vec![],
        }
    };

    scores_vec
}

// TODO: consider adding doctest for this function
fn list_scores_with_index<T: Topology>(scores_vec: &Vec<Box<dyn Score<T>>>) -> String {
    let mut display_str = String::new();
    for (i, s) in scores_vec.iter().enumerate() {
        let name = s.name();
        display_str.push_str(&format!("\n{i}: {name}"));
    }
    display_str
}

pub async fn run<T: Topology + Send + Sync + 'static>(
    inventory: Inventory,
    topology: T,
    scores: Vec<Box<dyn Score<T>>>,
    args_struct: Option<Args>,
) -> Result<(), Box<dyn std::error::Error>> {
    let args = match args_struct {
        Some(args) => args,
        None => Args::parse(),
    };

    #[cfg(not(feature = "tui"))]
    if args.interactive {
        return Err("Not compiled with interactive support".into());
    }

    #[cfg(feature = "tui")]
    if args.interactive {
        return harmony_tui::run(inventory, topology, scores).await;
    }

    run_cli(inventory, topology, scores, args).await
}

pub async fn run_cli<T: Topology + Send + Sync + 'static>(
    inventory: Inventory,
    topology: T,
    scores: Vec<Box<dyn Score<T>>>,
    args: Args,
) -> Result<(), Box<dyn std::error::Error>> {
    cli_logger::init();

    let mut maestro = Maestro::initialize(inventory, topology).await.unwrap();
    maestro.register_all(scores);

    let result = init(maestro, args).await;

    instrumentation::instrument(instrumentation::HarmonyEvent::HarmonyFinished).unwrap();
    result
}

async fn init<T: Topology + Send + Sync + 'static>(
    maestro: harmony::maestro::Maestro<T>,
    args: Args,
) -> Result<(), Box<dyn std::error::Error>> {
    let _ = env_logger::builder().try_init();

    let scores_vec = maestro_scores_filter(&maestro, args.all, args.filter, args.number);

    if scores_vec.is_empty() {
        return Err("No score found".into());
    }

    // if list option is specified, print filtered list and exit
    if args.list {
        let num_scores = scores_vec.len();
        println!("Available scores {num_scores}:");
        println!("{}\n\n", list_scores_with_index(&scores_vec));
        return Ok(());
    }

    // prompt user if --yes is not specified
    if !args.yes {
        let confirmation = Confirm::new(
            format!(
                "This will run the following scores: {}\n",
                list_scores_with_index(&scores_vec)
            )
            .as_str(),
        )
        .with_default(true)
        .prompt()
        .expect("Unexpected prompt error");

        if !confirmation {
            return Ok(());
        }
    }

    // Run filtered scores
    for s in scores_vec {
        debug!("Running: {}", s.name());
        maestro.interpret(s).await?;
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use harmony::{
        inventory::Inventory,
        maestro::Maestro,
        modules::dummy::{ErrorScore, PanicScore, SuccessScore},
        topology::HAClusterTopology,
    };

    fn init_test_maestro() -> Maestro<HAClusterTopology> {
        let inventory = Inventory::autoload();
        let topology = HAClusterTopology::autoload();
        let mut maestro = Maestro::new_without_initialization(inventory, topology);

        maestro.register_all(vec![
            Box::new(SuccessScore {}),
            Box::new(ErrorScore {}),
            Box::new(PanicScore {}),
        ]);

        maestro
    }

    #[tokio::test]
    async fn test_init_success_score() {
        let maestro = init_test_maestro();
        let res = crate::init(
            maestro,
            crate::Args {
                yes: true,
                filter: Some("SuccessScore".to_owned()),
                interactive: false,
                all: true,
                number: 0,
                list: false,
            },
        )
        .await;

        assert!(res.is_ok());
    }

    #[tokio::test]
    async fn test_init_error_score() {
        let maestro = init_test_maestro();

        let res = crate::init(
            maestro,
            crate::Args {
                yes: true,
                filter: Some("ErrorScore".to_owned()),
                interactive: false,
                all: true,
                number: 0,
                list: false,
            },
        )
        .await;

        assert!(res.is_err());
    }

    #[tokio::test]
    async fn test_init_number_score() {
        let maestro = init_test_maestro();

        let res = crate::init(
            maestro,
            crate::Args {
                yes: true,
                filter: None,
                interactive: false,
                all: false,
                number: 0,
                list: false,
            },
        )
        .await;

        assert!(res.is_ok());
    }

    #[tokio::test]
    async fn test_filter_fn_all() {
        let maestro = init_test_maestro();

        let res = crate::maestro_scores_filter(&maestro, true, None, 0);

        assert!(res.len() == 3);
    }

    #[tokio::test]
    async fn test_filter_fn_all_success() {
        let maestro = init_test_maestro();

        let res = crate::maestro_scores_filter(&maestro, true, Some("Success".to_owned()), 0);

        assert!(res.len() == 1);

        assert!(
            maestro
                .interpret(res.first().unwrap().clone_box())
                .await
                .is_ok()
        );
    }

    #[tokio::test]
    async fn test_filter_fn_all_error() {
        let maestro = init_test_maestro();

        let res = crate::maestro_scores_filter(&maestro, true, Some("Error".to_owned()), 0);

        assert!(res.len() == 1);

        assert!(
            maestro
                .interpret(res.first().unwrap().clone_box())
                .await
                .is_err()
        );
    }

    #[tokio::test]
    async fn test_filter_fn_all_score() {
        let maestro = init_test_maestro();

        let res = crate::maestro_scores_filter(&maestro, true, Some("Score".to_owned()), 0);

        assert!(res.len() == 3);

        assert!(
            maestro
                .interpret(res.first().unwrap().clone_box())
                .await
                .is_ok()
        );
        assert!(
            maestro
                .interpret(res.get(1).unwrap().clone_box())
                .await
                .is_err()
        );
    }

    #[tokio::test]
    async fn test_filter_fn_number() {
        let maestro = init_test_maestro();

        let res = crate::maestro_scores_filter(&maestro, false, None, 0);

        assert!(res.len() == 1);

        assert!(
            maestro
                .interpret(res.first().unwrap().clone_box())
                .await
                .is_ok()
        );
    }

    #[tokio::test]
    async fn test_filter_fn_number_invalid() {
        let maestro = init_test_maestro();

        let res = crate::maestro_scores_filter(&maestro, false, None, 11);

        assert!(res.is_empty());
    }
}