goose 0.15.1

use gumdrop::Options;
use httpmock::{Method::GET, Mock, MockServer};
use std::io::{Read, Write};
use std::net::TcpStream;
use std::{str, time};
use tokio_tungstenite::tungstenite::Message;

use goose::config::GooseConfiguration;
use goose::controller::{
    GooseControllerCommand, GooseControllerWebSocketRequest, GooseControllerWebSocketResponse,
};
use goose::prelude::*;

mod common;

// Paths used in load tests performed during these tests.
const INDEX_PATH: &str = "/";
const ABOUT_PATH: &str = "/about.html";

// Indexes to the above paths.
const INDEX_KEY: usize = 0;
const ABOUT_KEY: usize = 1;

const USERS: usize = 5;
const HATCH_RATE: usize = 10;
const RUN_TIME: usize = 10;

// There are multiple test variations in this file.
#[derive(Clone)]
enum TestType {
    // Enable --no-telnet.
    WebSocket,
    // Enable --no-websocket.
    Telnet,
}

// State machine for tracking Controller state during tests.
struct TestState {
    // A buffer for the telnet Controller.
    buf: [u8; 2048],
    // Track iterations through GooseControllerCommands.
    position: usize,
    // Track the steps within a given iteration.
    step: usize,
    // The Controller command currently being tested.
    command: GooseControllerCommand,
    // A TCP socket if testing the telnet Controller.
    telnet_stream: Option<TcpStream>,
    // A TCP socket if testing the WebSocket Controller.
    #[cfg(not(feature = "rustls-tls"))]
    websocket_stream: Option<tokio_tungstenite::tungstenite::WebSocket<std::net::TcpStream>>,
    #[cfg(feature = "rustls-tls")]
    websocket_stream: Option<
        tokio_tungstenite::tungstenite::WebSocket<
            tokio_tungstenite::tungstenite::stream::Stream<
                std::net::TcpStream,
                rustls::StreamOwned<rustls::ClientSession, TcpStream>,
            >,
        >,
    >,
    // A flag indicating whether or not to wait for a reply.
    websocket_expect_reply: bool,
    // A flag indicating whether or not the WebSocket controller is being tested.
    websocket_controller: bool,
}

// Test task.
pub async fn get_index(user: &mut GooseUser) -> GooseTaskResult {
    let _goose = user.get(INDEX_PATH).await?;
    Ok(())
}

// Test task.
pub async fn get_about(user: &mut GooseUser) -> GooseTaskResult {
    let _goose = user.get(ABOUT_PATH).await?;
    Ok(())
}

// All tests in this file run against the following common endpoints.
fn setup_mock_server_endpoints(server: &MockServer) -> Vec<Mock> {
    vec![
        // First set up INDEX_PATH, store in vector at INDEX_KEY.
        server.mock(|when, then| {
            when.method(GET).path(INDEX_PATH);
            then.status(200);
        }),
        // Next set up ABOUT_PATH, store in vector at ABOUT_KEY.
        server.mock(|when, then| {
            when.method(GET).path(ABOUT_PATH);
            then.status(200);
        }),
    ]
}

// Build appropriate configuration for these tests. Normally this also calls
// common::build_configuration() to get defaults most all tests needs, but
// for these tests we don't want a default configuration. We keep the signature
// the same to simplify reuse, accepting the MockServer but not using it.
fn common_build_configuration(_server: &MockServer, custom: &mut Vec<&str>) -> GooseConfiguration {
    // Common elements in all our tests.
    let mut configuration: Vec<&str> = vec!["--no-autostart", "--co-mitigation", "disabled"];

    // Custom elements in some tests.
    configuration.append(custom);

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

// Helper to confirm all variations generate appropriate results.
fn validate_one_taskset(
    goose_metrics: &GooseMetrics,
    mock_endpoints: &[Mock],
    configuration: &GooseConfiguration,
    _test_type: TestType,
) {
    //println!("goose_metrics: {:#?}", goose_metrics);
    //println!("configuration: {:#?}", configuration);

    // Confirm that we loaded the mock endpoints.
    assert!(mock_endpoints[INDEX_KEY].hits() > 0);
    assert!(mock_endpoints[ABOUT_KEY].hits() > 0);

    // Get index and about out of goose metrics.
    let index_metrics = goose_metrics
        .requests
        .get(&format!("GET {}", INDEX_PATH))
        .unwrap();
    let about_metrics = goose_metrics
        .requests
        .get(&format!("GET {}", ABOUT_PATH))
        .unwrap();

    // There should not have been any failures during this test.
    assert!(index_metrics.fail_count == 0);
    assert!(about_metrics.fail_count == 0);

    // Users were correctly configured through the controller.
    assert!(goose_metrics.users == USERS);

    // Host was not configured at start time.
    assert!(configuration.host.is_empty());

    // The load test was manually shut down instead of running to completion.
    assert!(goose_metrics.duration < RUN_TIME);
}

// Returns the appropriate taskset needed to build these tests.
fn get_tasks() -> GooseTaskSet {
    taskset!("LoadTest")
        .register_task(task!(get_index).set_weight(2).unwrap())
        .register_task(task!(get_about).set_weight(1).unwrap())
}

// Helper to run all standalone tests.
async fn run_standalone_test(test_type: TestType) {
    // Start the mock server.
    let server = MockServer::start();
    let server_url = server.base_url();

    // Setup the endpoints needed for this test on the mock server.
    let mock_endpoints = setup_mock_server_endpoints(&server);

    let mut configuration_flags = match &test_type {
        TestType::WebSocket => vec!["--no-telnet"],
        TestType::Telnet => vec!["--no-websocket"],
    };

    // Keep a copy for validation.
    let validate_test_type = test_type.clone();

    // Build common configuration elements.
    let configuration = common_build_configuration(&server, &mut configuration_flags);

    // Create a new thread from which to test the Controller.
    let _controller_handle = tokio::spawn(async move {
        // Sleep a half a second allowing the GooseAttack to start.
        tokio::time::sleep(time::Duration::from_millis(500)).await;

        // Initiailize the state engine.
        let mut test_state = update_state(None, &test_type);
        loop {
            // Process data received from the client in a loop.
            let response;
            let websocket_response: GooseControllerWebSocketResponse;
            if let Some(stream) = test_state.telnet_stream.as_mut() {
                let _ = match stream.read(&mut test_state.buf) {
                    Ok(data) => data,
                    Err(_) => {
                        panic!("ERROR: server disconnected!");
                    }
                };
                response = str::from_utf8(&test_state.buf).unwrap();
            // Process data received from the client in a loop.
            } else if let Some(stream) = test_state.websocket_stream.as_mut() {
                if !test_state.websocket_expect_reply {
                    response = "";
                    test_state.websocket_expect_reply = true;
                } else {
                    match stream.read_message() {
                        Ok(message) => {
                            if let Ok(r) = message.into_text() {
                                // Keep response around for the entire loop.
                                websocket_response = match serde_json::from_str(&r) {
                                    Ok(c) => c,
                                    Err(e) => panic!("invalid response from server: {}", e),
                                };
                                response = &websocket_response.response;
                            } else {
                                // @TODO: support non-text too
                                panic!("ERROR: invalid message type!");
                            }
                        }
                        Err(e) => {
                            panic!("error reading from server: {}", e);
                        }
                    }
                }
            } else {
                unreachable!();
            }

            //println!("{:?}: {}", test_state.command, response);
            match test_state.command {
                GooseControllerCommand::Exit => {
                    match test_state.step {
                        // Exit the Controller.
                        0 => {
                            make_request(&mut test_state, "exit\r\n");
                        }
                        // Confirm that the Controller exited.
                        _ => {
                            assert!(response.starts_with("goodbye!"));

                            // Re-connect to the Controller.
                            test_state = update_state(None, &test_type);

                            // Move onto the next command.
                            test_state = update_state(Some(test_state), &test_type);
                        }
                    }
                }
                GooseControllerCommand::Help => {
                    match test_state.step {
                        0 => {
                            // Request the help text.
                            make_request(&mut test_state, "help\r\n");
                        }
                        1 => {
                            // Be sure we actually received the help text.
                            assert!(response.contains("controller commands:"));

                            // Request the help text, using the short form.
                            make_request(&mut test_state, "?\r\n");
                        }
                        _ => {
                            // Be sure we actually received the help text.
                            assert!(response.contains("controller commands:"));

                            // Move onto the next command.
                            test_state = update_state(Some(test_state), &test_type);
                        }
                    }
                }
                GooseControllerCommand::Host => {
                    match test_state.step {
                        // Set the host to be load tested.
                        0 => {
                            make_request(&mut test_state, &["host ", &server_url, "\r\n"].concat());
                        }
                        // Confirm the host was configured.
                        1 => {
                            assert!(response.starts_with("host configured"));

                            // Then try and set an invalid host.
                            make_request(&mut test_state, "host foobar\r\n");
                        }
                        // Confirm that we can't configure an invalid host that doesn't
                        // match the regex.
                        2 => {
                            assert!(response.starts_with("unrecognized command"));

                            // Try again to set an invalid host.
                            make_request(&mut test_state, "host http://$[foo\r\n");
                        }
                        // Confirm that we can't configure an invalid host that does
                        // match the regex.
                        _ => {
                            assert!(response.starts_with("unrecognized command"));

                            // Move onto the next command.
                            test_state = update_state(Some(test_state), &test_type);
                        }
                    }
                }
                GooseControllerCommand::Users => {
                    match test_state.step {
                        // Reconfigure the number of users simulated by the load test.
                        0 => {
                            make_request(
                                &mut test_state,
                                &["users ", &USERS.to_string(), "\r\n"].concat(),
                            );
                        }
                        // Confirm that the users are reconfigured.
                        1 => {
                            assert!(response.starts_with("users configured"));

                            // Attempt to reconfigure users with bad data.
                            make_request(&mut test_state, "users 1.1\r\n");
                        }
                        // Confirm we can't configure users with a float.
                        _ => {
                            // The number of users started is verified when the load test finishes,
                            // so no further validation required here.
                            assert!(response.starts_with("unrecognized command"));

                            // Move onto the next command.
                            test_state = update_state(Some(test_state), &test_type);
                        }
                    }
                }
                GooseControllerCommand::HatchRate => {
                    match test_state.step {
                        // Configure a decimal hatch_rate.
                        0 => {
                            make_request(&mut test_state, "hatchrate .1\r\n");
                        }
                        // Confirm hatch_rate is configured.
                        1 => {
                            assert!(response.starts_with("hatch_rate configured"));

                            // Configure with leading and trailing zeros.
                            make_request(&mut test_state, "hatchrate 0.90\r\n");
                        }
                        // Confirm hatch_rate is configured.
                        2 => {
                            assert!(response.starts_with("hatch_rate configured"));

                            // Try to configure with an invalid decimal.
                            make_request(&mut test_state, "hatchrate 1.2.3\r\n");
                        }
                        // Confirm hatch_rate is not configured.
                        3 => {
                            assert!(response.starts_with("unrecognized command"));

                            // Configure hatch_rate with a single integer.
                            make_request(
                                &mut test_state,
                                &["hatchrate ", &HATCH_RATE.to_string(), "\r\n"].concat(),
                            );
                        }
                        // Confirm the final hatch_rate is configured.
                        _ => {
                            assert!(response.starts_with("hatch_rate configured"));

                            // The hatch_rate is verified when the load test finishes, so no
                            // further validation required here.

                            // Move onto the next command.
                            test_state = update_state(Some(test_state), &test_type);
                        }
                    }
                }
                GooseControllerCommand::RunTime => {
                    match test_state.step {
                        // Configure run_time using h:m:s format.
                        0 => {
                            // Set run_time with hours and minutes and seconds.
                            make_request(&mut test_state, "runtime 1h2m3s\r\n");
                        }
                        // Confirm the run_time is configured.
                        1 => {
                            assert!(response.starts_with("run_time configured"));

                            // Set run_time with hours and seconds.
                            make_request(&mut test_state, "run_time 1h2s\r\n");
                        }
                        // Confirm the run_time is configured.
                        2 => {
                            assert!(response.starts_with("run_time configured"));

                            // Set run_time with hours alone.
                            make_request(&mut test_state, "run-time 1h\r\n");
                        }
                        // Confirm the run_time is configured.
                        3 => {
                            assert!(response.starts_with("run_time configured"));

                            // Set run_time with seconds alone.
                            make_request(&mut test_state, "runtime 10s\r\n");
                        }
                        // Confirm the run_time is configured.
                        4 => {
                            assert!(response.starts_with("run_time configured"));

                            // Try to set run_time with unsupported value.
                            make_request(&mut test_state, "runtime 10d\r\n");
                        }
                        // Confirm the run_time is not configured.
                        5 => {
                            assert!(response.starts_with("unrecognized command"));

                            // Set run_time with seconds alone, and no "s".
                            make_request(
                                &mut test_state,
                                &["runtime ", &RUN_TIME.to_string(), "\r\n"].concat(),
                            );
                        }
                        // Confirm the run_time is configured.
                        _ => {
                            assert!(response.starts_with("run_time configured"));

                            // The run_time is verified when the load test finishes, so no
                            // further validation required here. Unfortunately, if this fails
                            // the load test could run forever.

                            // Move onto the next command.
                            test_state = update_state(Some(test_state), &test_type);
                        }
                    }
                }
                GooseControllerCommand::Config => {
                    match test_state.step {
                        // Request the configuration.
                        0 => {
                            make_request(&mut test_state, "config\r\n");
                        }
                        _ => {
                            // Confirm the configuration is returned in jsonformat.
                            if test_state.websocket_controller {
                                assert!(response
                                    .starts_with(r#"{"help":false,"version":false,"list":false,"#));
                            // Confirm the configuration object is returned.
                            } else {
                                assert!(response.starts_with(r"GooseConfiguration "));
                            }

                            // Move onto the next command.
                            test_state = update_state(Some(test_state), &test_type);
                        }
                    }
                }
                GooseControllerCommand::ConfigJson => {
                    match test_state.step {
                        // Request the configuration in json format.
                        0 => {
                            make_request(&mut test_state, "config-json\r\n");
                        }
                        // Confirm the configuration is returned in jsonformat.
                        _ => {
                            assert!(response
                                .starts_with(r#"{"help":false,"version":false,"list":false,"#));

                            // Move onto the next command.
                            test_state = update_state(Some(test_state), &test_type);
                        }
                    }
                }
                GooseControllerCommand::Metrics => {
                    match test_state.step {
                        // Request the running metrics.
                        0 => {
                            make_request(&mut test_state, "metrics\r\n");
                        }
                        _ => {
                            // Confirm the metrics are returned in json format.
                            if test_state.websocket_controller {
                                assert!(response.starts_with(r#"{"hash":0,"#));
                            }
                            // Confirm the metrics are returned and pretty-printed.
                            else {
                                assert!(response.contains("=== PER TASK METRICS ==="));
                            }

                            // Move onto the next command.
                            test_state = update_state(Some(test_state), &test_type);
                        }
                    }
                }
                GooseControllerCommand::MetricsJson => {
                    match test_state.step {
                        // Request the running metrics in json format.
                        0 => {
                            make_request(&mut test_state, "metrics-json\r\n");
                        }
                        // Confirm the metrics are returned in json format.
                        _ => {
                            assert!(response.starts_with(r#"{"hash":0,"#));

                            // Move onto the next command.
                            test_state = update_state(Some(test_state), &test_type);
                        }
                    }
                }
                GooseControllerCommand::Start => {
                    match test_state.step {
                        // Try to stop an idle load test.
                        0 => {
                            make_request(&mut test_state, "stop\r\n");
                        }
                        // Confirm an idle load test can not be stopped.
                        1 => {
                            assert!(response.starts_with("load test not running"));

                            // Send the start request.
                            make_request(&mut test_state, "start\r\n");
                        }
                        // Confirm an idle load test can be started.
                        2 => {
                            assert!(response.starts_with("load test started"));

                            // Send the start request again.
                            make_request(&mut test_state, "start\r\n");
                        }
                        // Confirm a running load test can not be started.
                        _ => {
                            assert!(response.starts_with("unable to start load test"));

                            // Move onto the next command.
                            test_state = update_state(Some(test_state), &test_type);
                        }
                    }
                }
                GooseControllerCommand::Stop => {
                    match test_state.step {
                        // Try to configure users on a running load test.
                        0 => {
                            make_request(&mut test_state, "users 1\r\n");
                        }
                        // Confirm users can not be configured on a running load test.
                        1 => {
                            assert!(response.starts_with("load test not idle"));

                            // Try to configure host on a running load test.
                            make_request(&mut test_state, "host http://localhost/\r\n");
                        }
                        // Confirm host can not be configured on a running load test.
                        2 => {
                            assert!(response.starts_with("failed to reconfigure host"));

                            // Try to stop a running load test.
                            make_request(&mut test_state, "stop\r\n");
                        }
                        // Confirm a running load test can be stopped.
                        _ => {
                            assert!(response.starts_with("load test stopped"));

                            // Give Goose a half second to stop before moving on.
                            tokio::time::sleep(time::Duration::from_millis(500)).await;

                            // Move onto the next command.
                            test_state = update_state(Some(test_state), &test_type);
                        }
                    }
                }
                GooseControllerCommand::Shutdown => {
                    match test_state.step {
                        // Shut down the load test.
                        0 => {
                            make_request(&mut test_state, "shutdown\r\n");
                        }
                        // Confirm that the load test shut down.
                        _ => {
                            assert!(response.starts_with("load test shut down"));

                            // Move onto the next command.
                            test_state = update_state(Some(test_state), &test_type);
                        }
                    }
                }
            }
            // Flush the buffer.
            test_state.buf = [0; 2048];

            // Give the parent process time to catch up.
            tokio::time::sleep(time::Duration::from_millis(100)).await;
        }
    });

    // Run the Goose Attack.
    let goose_metrics = common::run_load_test(
        common::build_load_test(configuration.clone(), &get_tasks(), None, None),
        None,
    )
    .await;

    // Confirm that the load test ran correctly.
    validate_one_taskset(
        &goose_metrics,
        &mock_endpoints,
        &configuration,
        validate_test_type,
    );
}

// Update (or create) the current testing state. A simple state maching for
// navigating through all supported Controller commands and test states.
fn update_state(test_state: Option<TestState>, test_type: &TestType) -> TestState {
    // The commands being tested, and the order they are tested.
    let commands_to_test = [
        GooseControllerCommand::Exit,
        GooseControllerCommand::Help,
        GooseControllerCommand::Host,
        GooseControllerCommand::Users,
        GooseControllerCommand::HatchRate,
        GooseControllerCommand::RunTime,
        GooseControllerCommand::Start,
        GooseControllerCommand::Config,
        GooseControllerCommand::ConfigJson,
        GooseControllerCommand::Metrics,
        GooseControllerCommand::MetricsJson,
        GooseControllerCommand::Stop,
        GooseControllerCommand::Shutdown,
    ];

    if let Some(mut state) = test_state {
        state.position += 1;
        state.step = 0;
        if let Some(command) = commands_to_test.get(state.position) {
            state.command = command.clone();
        }
        // Generate a new prompt.
        if let Some(stream) = state.telnet_stream.as_mut() {
            stream.write_all("\r\n".as_bytes()).unwrap();
        } else {
            state.websocket_expect_reply = false;
        }
        state
    } else {
        // Connect to telnet controller.
        let telnet_stream = match test_type {
            TestType::Telnet => Some(TcpStream::connect("127.0.0.1:5116").unwrap()),
            _ => None,
        };

        // Connect to WebSocket controller.
        let websocket_controller: bool;
        let websocket_stream = match test_type {
            TestType::WebSocket => {
                let (mut stream, _) =
                    tokio_tungstenite::tungstenite::client::connect("ws://127.0.0.1:5117").unwrap();
                // Send an empty message so the client performs a handshake.
                stream.write_message(Message::Text("".into())).unwrap();
                // Ignore the error that comes back.
                let _ = stream.read_message().unwrap();
                websocket_controller = true;
                Some(stream)
            }
            TestType::Telnet => {
                websocket_controller = false;
                None
            }
        };

        TestState {
            buf: [0; 2048],
            position: 0,
            step: 0,
            command: commands_to_test.first().unwrap().clone(),
            telnet_stream,
            websocket_stream,
            websocket_expect_reply: false,
            websocket_controller,
        }
    }
}

fn make_request(test_state: &mut TestState, command: &str) {
    //println!("making request: {}", command);
    if let Some(stream) = test_state.telnet_stream.as_mut() {
        stream.write_all(command.as_bytes()).unwrap()
    } else if let Some(stream) = test_state.websocket_stream.as_mut() {
        stream
            .write_message(Message::Text(
                serde_json::to_string(&GooseControllerWebSocketRequest {
                    request: command.to_string(),
                })
                .unwrap(),
            ))
            .unwrap()
    }
    test_state.step += 1;
}

// Test controlling a load test with Telnet.
#[tokio::test(flavor = "multi_thread", worker_threads = 8)]
async fn test_telnet_controller() {
    run_standalone_test(TestType::Telnet).await;
}

// Test controlling a load test with WebSocket controller.
#[tokio::test(flavor = "multi_thread", worker_threads = 8)]
async fn test_websocket_controller() {
    run_standalone_test(TestType::WebSocket).await;
}