foxhunt/services/api_gateway/load_tests/src/main.rs

#![allow(clippy::integer_division)]

use anyhow::Result;
use clap::{Parser, Subcommand};
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

mod clients;
mod config;
mod metrics;
mod orchestrator;
mod reporting;
mod scenarios;

#[derive(Parser)]
#[command(name = "load_test_runner")]
#[command(about = "API Gateway Load Testing Framework", long_about = None)]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Run normal load test (1K concurrent clients for 60s)
    Normal {
        #[arg(long, default_value = "http://localhost:50050")]
        gateway_url: String,
        #[arg(long, default_value = "1000")]
        num_clients: usize,
        #[arg(long, default_value = "60")]
        duration_secs: u64,
    },
    /// Run spike load test (0→10K clients in 10s, sustain 60s)
    Spike {
        #[arg(long, default_value = "http://localhost:50050")]
        gateway_url: String,
        #[arg(long, default_value = "10000")]
        target_clients: usize,
        #[arg(long, default_value = "10")]
        ramp_up_secs: u64,
        #[arg(long, default_value = "60")]
        sustain_secs: u64,
    },
    /// Run sustained load test (100 clients for 24h)
    Sustained {
        #[arg(long, default_value = "http://localhost:50050")]
        gateway_url: String,
        #[arg(long, default_value = "100")]
        num_clients: usize,
        #[arg(long, default_value = "86400")]
        duration_secs: u64,
    },
    /// Run stress test (incrementally increase until failure)
    Stress {
        #[arg(long, default_value = "http://localhost:50050")]
        gateway_url: String,
        #[arg(long, default_value = "100")]
        initial_clients: usize,
        #[arg(long, default_value = "100")]
        increment: usize,
        #[arg(long, default_value = "60")]
        increment_interval_secs: u64,
        #[arg(long, default_value = "50.0")]
        max_p99_latency_ms: f64,
        #[arg(long, default_value = "5.0")]
        max_error_rate_pct: f64,
    },
    /// Run all scenarios sequentially
    All {
        #[arg(long, default_value = "http://localhost:50050")]
        gateway_url: String,
    },
}

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize tracing
    tracing_subscriber::registry()
        .with(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "info,load_test_runner=debug".into()),
        )
        .with(tracing_subscriber::fmt::layer())
        .init();

    let cli = Cli::parse();

    match cli.command {
        Commands::Normal {
            gateway_url,
            num_clients,
            duration_secs,
        } => {
            tracing::info!(
                "Running NORMAL load test: {} clients for {}s",
                num_clients,
                duration_secs
            );
            let report =
                scenarios::normal_load::run(gateway_url, num_clients, duration_secs).await?;
            reporting::generate_html_report("normal_load_report.html", report)?;
        },
        Commands::Spike {
            gateway_url,
            target_clients,
            ramp_up_secs,
            sustain_secs,
        } => {
            tracing::info!(
                "Running SPIKE load test: 0→{} clients in {}s, sustain {}s",
                target_clients,
                ramp_up_secs,
                sustain_secs
            );
            let report =
                scenarios::spike_load::run(gateway_url, target_clients, ramp_up_secs, sustain_secs)
                    .await?;
            reporting::generate_html_report("spike_load_report.html", report)?;
        },
        Commands::Sustained {
            gateway_url,
            num_clients,
            duration_secs,
        } => {
            tracing::info!(
                "Running SUSTAINED load test: {} clients for {}s ({}h)",
                num_clients,
                duration_secs,
                duration_secs / 3600
            );
            let report =
                scenarios::sustained_load::run(gateway_url, num_clients, duration_secs).await?;
            reporting::generate_html_report("sustained_load_report.html", report)?;
        },
        Commands::Stress {
            gateway_url,
            initial_clients,
            increment,
            increment_interval_secs,
            max_p99_latency_ms,
            max_error_rate_pct,
        } => {
            tracing::info!(
                "Running STRESS test: start {} clients, increment by {} every {}s",
                initial_clients,
                increment,
                increment_interval_secs
            );
            let report = scenarios::stress_test::run(
                gateway_url,
                initial_clients,
                increment,
                increment_interval_secs,
                max_p99_latency_ms,
                max_error_rate_pct,
            )
            .await?;
            reporting::generate_html_report("stress_test_report.html", report)?;
        },
        Commands::All { gateway_url } => {
            tracing::info!("Running ALL load test scenarios sequentially");

            // Normal load
            let normal_report =
                scenarios::normal_load::run(gateway_url.clone(), 1000_usize, 60).await?;
            reporting::generate_html_report("normal_load_report.html", normal_report)?;

            // Wait between tests
            tokio::time::sleep(tokio::time::Duration::from_secs(30)).await;

            // Spike load
            let spike_report =
                scenarios::spike_load::run(gateway_url.clone(), 10000_usize, 10_u64, 60).await?;
            reporting::generate_html_report("spike_load_report.html", spike_report)?;

            // Wait between tests
            tokio::time::sleep(tokio::time::Duration::from_secs(30)).await;

            // Stress test (short version)
            let stress_report = scenarios::stress_test::run(
                gateway_url.clone(),
                100_usize,
                100_usize,
                60_u64,
                50.0,
                5.0,
            )
            .await?;
            reporting::generate_html_report("stress_test_report.html", stress_report)?;

            tracing::info!("All scenarios complete! Reports generated.");
        },
    }

    Ok(())
}