treadmill/

lib.rs

use std::{
    collections::{hash_map::Entry, HashMap},
    ffi::OsString,
    io::{BufReader, ErrorKind},
    path::{Path, PathBuf},
    process::Stdio,
};

use cargo_metadata::{
    diagnostic::{Diagnostic, DiagnosticLevel},
    BuildFinished, Edition, Message, MetadataCommand,
};

#[derive(thiserror::Error, Debug)]
pub enum Error {
    #[error("Missing environment variable: {0}. Expected to be run by cargo.")]
    MissingEnvVar(&'static str),

    #[error("Failed to get cargo metadata: {0}")]
    CargoMetadata(#[from] cargo_metadata::Error),

    #[error("Failed to find package: {0}")]
    NoPackage(std::io::Error),

    #[error("Failed to get path to test's executable: {0}")]
    NoCurrentExe(std::io::Error),

    #[error("Failed create cargo package for test binary: {0}")]
    CreatingCargoPackage(std::io::Error),

    #[error("Failed to run cargo build for test binary: {0}")]
    CargoBuildCommand(std::io::Error),

    #[error("Failed to build test binary.")]
    BuildError(HashMap<DiagnosticLevel, Vec<Diagnostic>>),
}

type Result<T, E = Error> = std::result::Result<T, E>;

#[derive(Debug)]
pub struct DetectedEnvironment {
    target: String,
    package_name: String,
    package_manifest: PathBuf,
    test_name: String,
    edition: Edition,
    target_directory: PathBuf,
}

#[derive(Debug)]
pub struct Program {
    /// Path to program's directory in the target folder.
    path: PathBuf,

    name: String,

    source: String,

    dependencies: Vec<String>,

    detected: DetectedEnvironment,
}

impl Program {
    pub fn new(name: String, source: String) -> Result<Self> {
        let package_name = package_name()?;

        let (base_dir, edition) = base_dir_and_edition(&package_name)?;

        let target = target();

        let package_manifest = package_manifest()?;

        let test_name = test_name()?;

        let path = base_dir.join(&package_name).join(&test_name).join(&name);

        let path_dependency = format!(
            "{} = {{ path = \"{}\"}}",
            package_name,
            package_manifest.to_string_lossy()
        );

        Ok(Program {
            path,
            name,
            source,
            dependencies: vec![path_dependency],
            detected: DetectedEnvironment {
                target: target.to_owned(),
                package_name,
                package_manifest,
                test_name,
                edition,
                target_directory: base_dir
                    .parent()
                    .expect("base dir should have parent")
                    .to_path_buf(),
            },
        })
    }

    pub fn add_dependency(&mut self, value: String) {
        self.dependencies.push(value);
    }

    pub fn compile(self) -> Result<Binary> {
        let cargo_path = cargo_path()?;

        // Crate directory for cargo package.
        std::fs::create_dir_all(&self.path)
            .map_err(Error::CreatingCargoPackage)?;

        // Create cargo manifest.
        // cargo won't rebuild if it didn't change.
        std::fs::write(self.path.join("Cargo.toml"), format!(r#"
[package]
name = "{name}"
version = "0.1.0"
edition = "{edition}"
publish = false

[[bin]]
name = "{name}"
path = "main.rs"

[dependencies]
{dependencies}

[package.metadata.treadmill]
notice = "This package was automatically generated by the `treadmill` crate for use with tests."
"#,
            name = &self.name,
            dependencies = self.dependencies.join("\n"),
            edition = self.detected.edition.as_str(),
        )).map_err(Error::CreatingCargoPackage)?;

        // cargo won't rebuild if it didn't change.
        std::fs::write(self.path.join("main.rs"), self.source)
            .map_err(Error::CreatingCargoPackage)?;

        // Set cargo command.
        let mut command = std::process::Command::new(cargo_path);
        command
            .args(["build", "--message-format=json", "-q", "--manifest-path"])
            .arg(self.path.join("Cargo.toml"))
            .arg("--bin")
            .arg(&self.name)
            .arg("--target")
            .arg(&self.detected.target)
            .arg("--target-dir")
            .arg(&self.detected.target_directory)
            .stdout(Stdio::piped());

        let mut spawned = command.spawn().map_err(Error::CargoBuildCommand)?;

        let reader = BufReader::new(
            spawned
                .stdout
                .take()
                .expect("spawned should have stdout buffer"),
        );

        let messages = Message::parse_stream(reader);

        let mut executable_path = None;
        let mut diagnostics: HashMap<DiagnosticLevel, Vec<Diagnostic>> =
            HashMap::new();
        let mut build_finished = false;

        // Record the path to the artifact and any diagnostics.
        for message in messages {
            match message {
                Ok(Message::CompilerArtifact(artifact))
                    if artifact.target.name == self.name =>
                {
                    if let Some(path) = artifact.executable {
                        executable_path = Some(path);
                    }
                }
                Ok(Message::CompilerMessage(msg))
                    if msg.target.name == self.name =>
                {
                    match diagnostics.entry(msg.message.level) {
                        Entry::Occupied(occupied) => {
                            occupied.into_mut().push(msg.message)
                        }
                        Entry::Vacant(vacant) => {
                            vacant.insert(vec![msg.message]);
                        }
                    }
                }
                Ok(Message::BuildFinished(BuildFinished {
                    success: true,
                    ..
                })) => build_finished = true,
                _ => {}
            }
        }

        if build_finished {
            if let Some(path) = executable_path {
                Ok(Binary {
                    path: path.into_std_path_buf(),
                })
            } else {
                Err(Error::CargoBuildCommand(std::io::Error::new(
                    ErrorKind::Other,
                    "No executable was generated.",
                )))
            }
        } else {
            Err(Error::BuildError(diagnostics))
        }
    }
}

fn environment_var(name: &'static str) -> Result<OsString> {
    std::env::var_os(name).ok_or(Error::MissingEnvVar(name))
}

fn base_dir_and_edition(package_name: &str) -> Result<(PathBuf, Edition)> {
    let manifest_path = package_manifest()?.join("Cargo.toml");

    // cargo_metadata will find the target folder for us.
    let metadata =
        MetadataCommand::new().manifest_path(manifest_path).exec()?;

    let package = metadata
        .packages
        .iter()
        .find(|package| package.name == package_name)
        .ok_or_else(|| {
            Error::NoPackage(std::io::Error::new(ErrorKind::Other, ""))
        })?;

    let edition = package.edition.clone();

    // ttp stands for "treadmill test programs" but is short to minimize the length
    // of the paths for systems like Windows.
    Ok((
        metadata.target_directory.join("ttp").into_std_path_buf(),
        edition,
    ))
}

fn cargo_path() -> Result<PathBuf> {
    // cargo will set this environment variable when using `cargo test`.
    // https://doc.rust-lang.org/cargo/reference/environment-variables.html#environment-variables-cargo-sets-for-crates
    environment_var("CARGO").map(PathBuf::from)
}

fn target() -> &'static str {
    // This is set by the build script.
    env!("TREADMILL_TARGET")
}

fn package_manifest() -> Result<PathBuf> {
    // cargo will set this environment variable when using `cargo test`.
    // https://doc.rust-lang.org/cargo/reference/environment-variables.html#environment-variables-cargo-sets-for-crates
    environment_var("CARGO_MANIFEST_DIR").map(PathBuf::from)
}

fn package_name() -> Result<String> {
    // cargo will set this environment variable when using `cargo test`.
    // https://doc.rust-lang.org/cargo/reference/environment-variables.html#environment-variables-cargo-sets-for-crates
    environment_var("CARGO_PKG_NAME")
        .map(|name| name.to_string_lossy().into_owned())
}

fn test_name() -> Result<String> {
    // We will get the name of the test from the executable.
    // We can do this because cargo names test executables like `test_name-hash`.
    let current_exe = std::env::current_exe().map_err(Error::NoCurrentExe)?;

    let name = current_exe
        .file_name()
        .ok_or_else(|| {
            Error::NoCurrentExe(std::io::Error::new(
                ErrorKind::Other,
                "File doesn't have a file name.",
            ))
        })?
        .to_string_lossy()
        .into_owned();

    // Remove the hash. We just want the name of the test.
    if let Some((test_name, _)) = name.rsplit_once('-') {
        Ok(test_name.to_owned())
    } else {
        Ok(name)
    }
}

pub struct Binary {
    path: PathBuf,
}

impl Binary {
    pub fn command(&self) -> std::process::Command {
        std::process::Command::new(&self.path)
    }
}