iocaine 3.4.0

// SPDX-FileCopyrightText: Gergely Nagy
// SPDX-FileContributor: Gergely Nagy
//
// SPDX-License-Identifier: MIT

use anyhow::Result;
use axum::extract::Query;
use clap::{Args as ClapArgs, Subcommand};
use std::collections::HashMap;
use std::error::Error;
use std::io::{self, Write};
use std::path::PathBuf;

use iocaine_powder::{
    acab::State,
    http::{
        header::{self, HeaderMap, HeaderName},
        uri::Uri,
    },
    sex_dungeon::{DungeonMaster, Request, Response},
};

use crate::{
    dominatrix::{self, Config},
    tenx_programmer::TenXProgrammer,
};

pub fn run(config: &Config, command: TestCommands) -> Result<()> {
    match command {
        TestCommands::Suite { handler_names } => run_tests(config, &handler_names),
        TestCommands::Decision(args) => run_test_request(config, None, args, false),
        TestCommands::Output { decision, args } => {
            run_test_request(config, decision.as_ref(), args, true)
        }
    }
}

fn run_tests(config: &Config, handler_names: &[String]) -> Result<()> {
    let handlers = find_handlers(config, handler_names)?;

    for (name, decl) in handlers {
        tracing::info!({ name }, "Running handler tests");
        println!("==> {name} <==");
        let mut handler = DungeonMaster::new(&config.initial_seed)
            .language(decl.language)
            .compiler(decl.compiler)
            .path(decl.path)
            .config(decl.config)
            .build(&TenXProgrammer::default().metrics, &State::default())
            .map_err(|e| anyhow::anyhow!("{e:#?}"))?;
        handler.run_tests().map_err(|e| anyhow::anyhow!("{e:#?}"))?;
        println!();
    }

    Ok(())
}

fn run_test_request(
    config: &Config,
    decision: Option<&String>,
    request_args: RequestArgs,
    with_output: bool,
) -> Result<()> {
    let handlers = find_handlers(config, &request_args.handler_names)?;
    let uri: Uri = request_args.url.parse()?;

    let mut headers = HeaderMap::new();
    headers.insert(header::HOST, uri.host().unwrap_or("example.com").parse()?);
    headers.insert(header::USER_AGENT, request_args.user_agent.parse()?);

    for (key, value) in request_args.headers {
        headers.insert(HeaderName::from_bytes(key.as_bytes())?, value.parse()?);
    }

    let request = Request {
        method: request_args.method,
        path: uri.path().to_owned(),
        headers,
        params: Query::try_from_uri(&uri)?.0,
    };

    for (name, decl) in handlers {
        let handler = DungeonMaster::new(&config.initial_seed)
            .language(decl.language)
            .compiler(request_args.compiler.as_ref().or(decl.compiler.as_ref()))
            .path(decl.path)
            .config(decl.config)
            .build(&TenXProgrammer::default().metrics, &State::default())
            .map_err(|e| anyhow::anyhow!("{e:#?}"))?;

        let decision = if let Some(v) = decision {
            v
        } else {
            tracing::info!({
                    handler_name = name,
                    method = request.method,
                    path = request.path,
                    headers = format!("{:?}", request.headers),
                    params = format!("{:?}", request.params)
                }, "Executing `decide()`");

            &handler
                .decide(request.clone().into())
                .map_err(|e| anyhow::anyhow!("{e:#?}"))?
        };

        if with_output {
            tracing::info!({
                handler_name = name,
                method = request.method,
                path = request.path,
                headers = format!("{:?}", request.headers),
                params = format!("{:?}", request.params)
            }, "Executing `output()`");
            if request_args.handler_names.len() > 1 || request_args.handler_names.is_empty() {
                println!("==> {name} <==");
            }

            let mut response = handler
                .output(request.clone().into(), Some(decision.to_owned()))
                .map_err(|e| anyhow::anyhow!("{e:#?}"))?;
            io::stdout().write_all(&http_response(&mut response))?;
        } else if request_args.handler_names.len() > 1 || request_args.handler_names.is_empty() {
            println!("{name}: {decision}");
        } else {
            println!("{decision}");
        }
    }

    Ok(())
}

fn http_response(response: &mut Response) -> Vec<u8> {
    let mut output = Vec::new();
    output.append(&mut format!("HTTP/1.1 {}\n", response.status_code).into());

    for (name, value) in &response.headers {
        output.append(&mut name.as_str().into());
        output.push(b':');
        output.push(b' ');
        output.append(&mut value.as_bytes().into());
        output.push(b'\n');
    }
    if !response.body.is_empty() {
        output.push(b'\n');
    }
    output.append(&mut response.body);
    output.push(b'\n');

    output
}

fn find_handlers(
    config: &Config,
    handler_names: &[String],
) -> Result<HashMap<String, dominatrix::Handler>> {
    let mut handlers;

    if handler_names.is_empty() {
        handlers = config.handlers.clone();
    } else {
        handlers = HashMap::new();
        for name in handler_names {
            let decl = config.handlers.get(name).ok_or_else(|| {
                tracing::error!({ name }, "handler not found");
                anyhow::anyhow!("handler not found: {name}")
            })?;
            handlers.insert(name.clone(), decl.clone());
        }
    }
    Ok(handlers)
}

#[derive(Debug, Subcommand)]
pub enum TestCommands {
    /// Run the test suite embedded in the request handlers.
    #[command(bin_name = "iocaine test suite")]
    Suite {
        /// Handlers whose test suite should be run. If not specified, runs the test suite of all declared handlers
        handler_names: Vec<String>,
    },
    /// Run a single, synthetic request through a request handler's decide function
    #[command(bin_name = "iocaine test decision")]
    Decision(RequestArgs),
    /// Run a single, synthetic request through a request handler's output function
    #[command(bin_name = "iocaine test output")]
    Output {
        /// Define a decision. If not specified, will run the decide function aswell
        #[arg(short = 'D', long)]
        decision: Option<String>,

        #[command(flatten)]
        args: RequestArgs,
    },
}

#[derive(Clone, Debug, ClapArgs)]
pub struct RequestArgs {
    /// Query method to make the test request with
    #[arg(short = 'M', long, default_value = "GET")]
    method: String,
    /// URL to make the test request with
    #[arg(short = 'U', long, default_value = "http://example.com")]
    url: String,
    /// The user-agent to make the test request with
    #[arg(short = 'A', long, default_value = "iocaine test ad-hoc")]
    user_agent: String,
    /// Additional HTTP headers for the test request, in `name=value` format. Can be specified multiple times
    #[arg(short = 'H', long = "header", value_parser = parse_key_val::<String, String>)]
    headers: Vec<(String, String)>,

    /// Compiler to use for the request handler. Only relevant for Fennel.
    #[arg(short = 'C', long)]
    compiler: Option<PathBuf>,

    /// Handlers the request should be ran through. If not specified, runs the request through all declared handlers
    handler_names: Vec<String>,
}

impl Default for TestCommands {
    fn default() -> Self {
        Self::Suite {
            handler_names: Vec::new(),
        }
    }
}

fn parse_key_val<T, U>(
    s: &str,
) -> std::result::Result<(T, U), Box<dyn Error + Send + Sync + 'static>>
where
    T: std::str::FromStr,
    T::Err: Error + Send + Sync + 'static,
    U: std::str::FromStr,
    U::Err: Error + Send + Sync + 'static,
{
    let pos = s
        .find('=')
        .ok_or_else(|| format!("invalid KEY=value: no `=` found in `{s}`"))?;
    Ok((s[..pos].parse()?, s[pos + 1..].parse()?))
}