btdt-server 0.3.2

use reqwest::blocking::{Client, RequestBuilder};
use reqwest::{Certificate, Url};
use std::collections::BTreeMap;
use std::fmt::{Display, Formatter};
use std::path::PathBuf;
use std::process::{Child, Command, ExitStatus};
use std::time::{Duration, Instant};
use std::{env, fs, io};
use tempfile::{NamedTempFile, TempDir, tempdir};

#[allow(unused)]
pub static CERTIFICATE_PKCS12: &[u8] = include_bytes!("../../tls/leaf.p12");
pub static CERTIFICATE_PEM: &[u8] = include_bytes!("../../tls/ca.pem");

pub struct BtdtTestServer {
    _config_file: NamedTempFile,
    _private_key_dir: Option<TempDir>,
    process: Child,
    client: Client,
    base_url: Url,
}

impl Default for BtdtTestServer {
    fn default() -> Self {
        Self::new(&BTreeMap::default())
    }
}

impl BtdtTestServer {
    pub fn new(env: &BTreeMap<String, String>) -> Self {
        let config_file = tempfile::Builder::new().suffix(".toml").tempfile().unwrap();
        fs::write(
            config_file.path(),
            "\
                [caches]\n\
                test-cache = { type = 'InMemory' }\
        ",
        )
        .unwrap();

        let target_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("target-test");
        let mut build_command = Command::new("cargo");
        build_command.args(&[
            "build",
            "--profile",
            "test",
            "--package",
            "btdt-server",
            "--bin",
            "btdt-server",
            "--target-dir",
            target_dir.to_str().unwrap(),
        ]);
        let mut process = build_command.spawn().expect("failed to build btdt-server");
        if !process.wait().unwrap().success() {
            panic!("failed to build btdt-server");
        }

        static BIND_ADDR: &str = "127.0.0.1:8707";
        let mut command = Command::new("cargo");
        command.args(&[
            "run",
            "--profile",
            "test",
            "--package",
            "btdt-server",
            "--bin",
            "btdt-server",
            "--target-dir",
            target_dir.to_str().unwrap(),
        ]);
        command.env("CARGO_TARGET_DIR", "foo");
        command.env("BTDT_BIND_ADDRS", BIND_ADDR);
        command.env("BTDT_SERVER_CONFIG_FILE", config_file.path());
        for (key, value) in env {
            command.env(key, value);
        }

        let private_key_dir = if !env.contains_key("BTDT_AUTH_PRIVATE_KEY") {
            let private_key_dir = tempdir().unwrap();
            command.env(
                "BTDT_AUTH_PRIVATE_KEY",
                private_key_dir.path().join("auth-private-key"),
            );
            Some(private_key_dir)
        } else {
            None
        };

        let process = command.spawn().expect("failed to start btdt-server");
        let tls_enabled = env.contains_key("BTDT_TLS_KEYSTORE");
        Self {
            _config_file: config_file,
            _private_key_dir: private_key_dir,
            process,
            client: Client::builder()
                .add_root_certificate(Certificate::from_pem(CERTIFICATE_PEM).unwrap())
                .use_rustls_tls()
                .build()
                .unwrap(),
            base_url: Url::parse(&format!(
                "http{}://{BIND_ADDR}",
                if tls_enabled { "s" } else { "" }
            ))
            .expect("bind address did not form a valid URL"),
        }
    }
}

impl Drop for BtdtTestServer {
    fn drop(&mut self) {
        if let Err(e) = self.process.kill() {
            eprintln!("Failed to kill btdt-server: {}", e);
        }
    }
}

impl BtdtTestServer {
    pub fn base_url(&self) -> &Url {
        &self.base_url
    }

    pub fn get(&self, path: &str) -> RequestBuilder {
        let url = self.base_url.join(path).expect("Invalid path");
        self.client.get(url)
    }

    pub fn is_ready(&self) -> bool {
        self.get("/api/health")
            .send()
            .map_or(false, |r| r.error_for_status().is_ok())
    }

    pub fn wait_until_ready(self) -> Result<Self, WaitTimeout> {
        let start = Instant::now();
        while start.elapsed() < Duration::from_secs(5) {
            if self.is_ready() {
                return Ok(self);
            }
            std::thread::sleep(Duration::from_millis(100));
        }
        Err(WaitTimeout)
    }

    pub fn wait_for_shutdown(mut self) -> Result<io::Result<ExitStatus>, WaitTimeout> {
        let start = Instant::now();
        while start.elapsed() < Duration::from_secs(60) {
            if let Some(status) = self.process.try_wait().transpose() {
                return Ok(status);
            }
        }
        Err(WaitTimeout)
    }
}

#[derive(Debug)]
pub struct WaitTimeout;

impl Display for WaitTimeout {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.write_str("Wait timeout exceeded")
    }
}

impl std::error::Error for WaitTimeout {}