1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

use anyhow::anyhow;
use std::ffi::OsStr;
use std::fmt;
use std::process::Command;
use std::str::from_utf8;

/// An execution environment, consisting of environment variables
/// which are provided on the launch of each new process.
pub struct Executor<K, V>
where
    K: AsRef<OsStr> + Clone,
    V: AsRef<OsStr> + Clone,
{
    env: Vec<(K, V)>,
}

impl<K, V> Executor<K, V>
where
    K: AsRef<OsStr> + Clone,
    V: AsRef<OsStr> + Clone,
{
    /// Initializes a new Executor.
    ///
    /// All environment variables are provided to processes launched
    /// with the `run` method.
    pub fn new(env: Vec<(K, V)>) -> Self {
        Executor { env }
    }

    /// Launches a new subprocess and awaits its completion.
    ///
    /// Pretty-prints stdout/stderr on failure.
    ///
    /// # Panics
    ///
    /// This method is a little aggressive about panicking; it
    /// can totally evolve structured errors if that would be useful.
    /// However, given that the primary purpose is testing, this
    /// behavior is *currently* acceptable.
    ///
    /// Panics if...
    /// - `args` is empty.
    /// - The sub-process fails to execute.
    /// - The execution of the sub-process returns a non-zero exit code.
    /// - The sub-process writes invalid UTF-8 stdout/stderr.
    pub fn run<I, S>(&self, args: I)
    where
        I: IntoIterator<Item = S>,
        S: AsRef<OsStr>,
    {
        Execution::run(args, self.env.clone())
    }
}

struct Execution<S: AsRef<OsStr>> {
    cmd: S,
    args: Vec<S>,
    result: Option<std::process::Output>,
}

impl<S: AsRef<OsStr>> Execution<S> {
    fn run<I, K, V, E>(args: I, envs: E)
    where
        I: IntoIterator<Item = S>,
        K: AsRef<OsStr>,
        V: AsRef<OsStr>,
        E: IntoIterator<Item = (K, V)>,
    {
        let mut iter = args.into_iter();
        let mut exec = Execution {
            cmd: iter
                .next()
                .ok_or_else(|| anyhow!("Missing command"))
                .unwrap(),
            args: iter.collect::<Vec<S>>(),
            result: None,
        };

        exec.result = Some(
            Command::new(&exec.cmd)
                .args(&exec.args)
                .envs(envs)
                .output()
                .expect("Failed to execute command"),
        );
        assert!(
            exec.result.as_ref().unwrap().status.success(),
            format!("{}", exec)
        );
    }
}

impl<S: AsRef<OsStr>> fmt::Display for Execution<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut cmd = Vec::new();
        cmd.push(self.cmd.as_ref());
        for arg in &self.args {
            cmd.push(arg.as_ref());
        }
        let cmd: Vec<String> = cmd
            .into_iter()
            .map(|osstr| osstr.to_string_lossy().to_string())
            .collect();
        write!(f, "\x1b[95m{}\x1b[0m", cmd.join(" "))?;
        if let Some(out) = self.result.as_ref() {
            if !out.status.success() {
                write!(f, "\n{}", out.status)?;
            }
            if !out.stdout.is_empty() {
                write!(f, "\n\x1b[92m{}\x1b[0m", from_utf8(&out.stdout).unwrap())?;
            }
            if !out.stderr.is_empty() {
                write!(f, "\n\x1b[91m{}\x1b[0m", from_utf8(&out.stderr).unwrap())?;
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn export_right_environment_variable() {
        let executor = Executor::new(vec![("FOO", "BAR")]);
        executor.run(vec!["/bin/bash", "-c", "[ \"$FOO\" == \"BAR\" ]"]);
    }

    #[test]
    #[should_panic]
    fn export_wrong_environment_variable() {
        let executor = Executor::new(vec![("FOO", "BAZINGA")]);
        executor.run(vec!["/bin/bash", "-c", "[ \"$FOO\" == \"BAR\" ]"]);
    }
}