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>>;
#[allow(dead_code)]
pub fn build_configuration(server: &MockServer, custom: Vec<&str>) -> GooseConfiguration {
let mut configuration: Vec<&str> = vec![];
let server_url = server.base_url();
configuration.extend_from_slice(&custom);
if !configuration.contains(&"--worker") {
if !configuration.contains(&"--host") {
configuration.extend_from_slice(&["--host", &server_url]);
}
if !configuration.contains(&"--users") {
configuration.extend_from_slice(&["--users", "1"]);
}
if !configuration.contains(&"--hatch-rate") {
configuration.extend_from_slice(&["--hatch-rate", "1"]);
}
if !configuration.contains(&"--run-time") {
configuration.extend_from_slice(&["--run-time", "1"]);
}
if !configuration.contains(&"--co-mitigation") {
configuration.extend_from_slice(&["--co-mitigation", "disabled"]);
}
}
configuration.extend_from_slice(&["--quiet"]);
GooseConfiguration::parse_args_default(&configuration)
.expect("failed to parse options and generate a configuration")
}
#[allow(dead_code)]
pub fn launch_gaggle_workers<F: Fn() -> GooseAttack>(
expect_workers: usize,
goose_attack_provider: F,
) -> WorkerHandles {
let mut worker_handles = Vec::new();
for _ in 0..expect_workers {
let worker_goose_attack = goose_attack_provider();
worker_handles.push(tokio::spawn(run_load_test(worker_goose_attack, None)));
}
worker_handles
}
#[allow(dead_code)]
pub fn build_load_test(
configuration: GooseConfiguration,
scenarios: Vec<Scenario>,
start_transaction: Option<&Transaction>,
stop_transaction: Option<&Transaction>,
) -> GooseAttack {
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
}
pub async fn run_load_test(
goose_attack: GooseAttack,
worker_handles: Option<WorkerHandles>,
) -> GooseMetrics {
let goose_metrics = goose_attack.execute().await.unwrap();
if let Some(handles) = worker_handles {
join_all(handles).await;
}
goose_metrics
}
#[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
}
}
#[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");
}
}
}
#[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)); })
}
#[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
);
}
#[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 {
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
}
}