goose 0.18.1

A load testing framework inspired by Locust.
Documentation
use futures::future::join_all;
use gumdrop::Options;
use httpmock::MockServer;
use std::io::{self, BufRead};

use goose::config::GooseConfiguration;
use goose::goose::{Scenario, Transaction};
use goose::metrics::GooseMetrics;
use goose::GooseAttack;

type WorkerHandles = Vec<tokio::task::JoinHandle<GooseMetrics>>;

/// Not all functions are used by all tests, so we enable allow(dead_code) to avoid
/// compiler warnings during testing.
///
/// The following options are configured by default, if not set to a custom value
/// and if not building a Worker configuration:
///  --host <mock-server>
///  --users 1
///  --hatch-rate 1
///  --run-time 1
///  --co-mitigation disabled
#[allow(dead_code)]
pub fn build_configuration(server: &MockServer, custom: Vec<&str>) -> GooseConfiguration {
    // Start with an empty configuration.
    let mut configuration: Vec<&str> = vec![];
    // Declare server_url here no matter what, so its lifetime is sufficient when needed.
    let server_url = server.base_url();

    // Merge in all custom options first.
    configuration.extend_from_slice(&custom);

    // If not building a Worker configuration, set some defaults.
    if !configuration.contains(&"--worker") {
        // Default to using mock server if not otherwise configured.
        if !configuration.contains(&"--host") {
            configuration.extend_from_slice(&["--host", &server_url]);
        }

        // Default to testing with 1 user if not otherwise configured.
        if !configuration.contains(&"--users") {
            configuration.extend_from_slice(&["--users", "1"]);
        }

        // Default to hatch 1 user per second if not otherwise configured.
        if !configuration.contains(&"--hatch-rate") {
            configuration.extend_from_slice(&["--hatch-rate", "1"]);
        }

        // Default to running for 1 second if not otherwise configured.
        if !configuration.contains(&"--run-time") {
            configuration.extend_from_slice(&["--run-time", "1"]);
        }

        // Default to disabling coordinated omission mitigation if not otherwise configured
        if !configuration.contains(&"--co-mitigation") {
            configuration.extend_from_slice(&["--co-mitigation", "disabled"]);
        }
    }

    // Disable verbose output when running tests.
    configuration.extend_from_slice(&["--quiet"]);

    // Parse these options to generate a GooseConfiguration.
    GooseConfiguration::parse_args_default(&configuration)
        .expect("failed to parse options and generate a configuration")
}

/// Launch each Worker in its own thread, and return a vector of Worker handles.
#[allow(dead_code)]
pub fn launch_gaggle_workers<F: Fn() -> GooseAttack>(
    // The number of Workers to launch.
    expect_workers: usize,
    // A goose attack object which is cloned for each Worker.
    goose_attack_provider: F,
) -> WorkerHandles {
    // Launch each worker in its own thread, storing the join handles.
    let mut worker_handles = Vec::new();
    for _ in 0..expect_workers {
        let worker_goose_attack = goose_attack_provider();
        // Start worker instance of the load test.
        worker_handles.push(tokio::spawn(run_load_test(worker_goose_attack, None)));
    }

    worker_handles
}

// Create a GooseAttack object from the configuration, Scenarios, and optional start and
// stop Transactions.
#[allow(dead_code)]
pub fn build_load_test(
    configuration: GooseConfiguration,
    scenarios: Vec<Scenario>,
    start_transaction: Option<&Transaction>,
    stop_transaction: Option<&Transaction>,
) -> GooseAttack {
    // First set up the common base configuration.
    let mut goose = GooseAttack::initialize_with_config(configuration).unwrap();

    for scenario in scenarios {
        goose = goose.register_scenario(scenario.clone());
    }

    if let Some(transaction) = start_transaction {
        goose = goose.test_start(transaction.clone());
    }

    if let Some(transaction) = stop_transaction {
        goose = goose.test_stop(transaction.clone());
    }

    goose
}

/// Run the actual load test, returning the GooseMetrics.
pub async fn run_load_test(
    goose_attack: GooseAttack,
    worker_handles: Option<WorkerHandles>,
) -> GooseMetrics {
    // Execute the load test.
    let goose_metrics = goose_attack.execute().await.unwrap();

    // If this is a Manager test, first wait for the Workers to exit to return.
    if let Some(handles) = worker_handles {
        // Wait for both worker threads to finish and exit.
        join_all(handles).await;
    }

    goose_metrics
}

/// Helper to count the number of lines in a test artifact.
#[allow(dead_code)]
pub fn file_length(file_name: &str) -> usize {
    if let Ok(file) = std::fs::File::open(std::path::Path::new(file_name)) {
        io::BufReader::new(file).lines().count()
    } else {
        0
    }
}

/// Helper to delete test artifacts, if existing.
#[allow(dead_code)]
pub fn cleanup_files(files: Vec<&str>) {
    for file in files {
        if std::path::Path::new(file).exists() {
            std::fs::remove_file(file).expect("failed to remove file");
        }
    }
}

/// Configure mock server to trigger CO events predictably.
#[allow(dead_code)]
pub fn setup_co_triggering_server(server: &MockServer) -> httpmock::Mock<'_> {
    use httpmock::Method::GET;
    use std::time::Duration;

    server.mock(|when, then| {
        when.method(GET).path("/");
        then.status(200).delay(Duration::from_millis(50)); // Base delay that should trigger CO events
    })
}

/// Validate CO metrics meet expected criteria.
#[allow(dead_code)]
pub fn validate_co_metrics(
    metrics: &goose::metrics::GooseMetrics,
    expected_events: usize,
    min_synthetic_pct: f32,
) {
    assert!(
        metrics.coordinated_omission_metrics.is_some(),
        "CO metrics should be present"
    );
    let co_metrics = metrics.coordinated_omission_metrics.as_ref().unwrap();
    assert!(
        co_metrics.co_events.len() >= expected_events,
        "Expected at least {} CO events, found {}",
        expected_events,
        co_metrics.co_events.len()
    );
    assert!(
        co_metrics.synthetic_percentage >= min_synthetic_pct,
        "Expected at least {:.1}% synthetic requests, found {:.1}%",
        min_synthetic_pct,
        co_metrics.synthetic_percentage
    );
}

/// Helper to run load test with CO-specific setup.
#[allow(dead_code)]
pub fn run_load_test_with_co(
    config: goose::config::GooseConfiguration,
) -> impl std::future::Future<Output = goose::metrics::GooseMetrics> {
    use goose::prelude::*;

    async move {
        // Define a simple test transaction
        async fn get_index(user: &mut GooseUser) -> TransactionResult {
            let _goose = user.get("/").await?;
            Ok(())
        }

        let goose_attack = build_load_test(
            config,
            vec![scenario!("Test").register_transaction(transaction!(get_index))],
            None,
            None,
        );

        run_load_test(goose_attack, None).await
    }
}