assert_cmd.rs

// UPSTREAMING: Though it has seen some refactoring, this code has grown organically and doesn't use
// the latest and greatest assert_cmd APIs (actually it only uses assrt_fs). This file has been
// carelessly stripped from an internal project, don't expect it to compile as-is.

use assert_fs::{prelude::*, TempDir};
use libc::{c_int, kill, SIGCONT, SIGINT, SIGKILL, SIGTERM};
use log::*;
use regex::Regex;
use std::{collections::{HashMap, HashSet},
          env,
          ffi::OsStr,
          fmt::{Debug, Display, Formatter, Result as FmtResult},
          fs::File,
          io::{Read, Write},
          iter::FromIterator,
          net::TcpListener,
          ops::Range,
          os::unix::process::ExitStatusExt,
          path::Path,
          process::{Child, Command, Stdio},
          sync::atomic::{AtomicU16, Ordering},
          thread::sleep,
          time::{Duration, Instant}};


/// Test file content.
// UPSTREAMING: probably reinventing the wheel here, but it's a small one that fits well.
pub struct FilePredicate {
    /// Desciption for assert-logging purpose.
    desc: String,
    /// Which file to operate on.
    file: String,
    /// Closure that tests the content of the file.
    pred: Box<dyn Fn(&str) -> bool>,
}
impl FilePredicate {
    pub fn new(desc: &str, file: &str, p: impl Fn(&str) -> bool + 'static) -> Self {
        FilePredicate { desc: desc.to_owned(), file: file.to_owned(), pred: Box::new(p) }
    }
}
impl Debug for FilePredicate {
    fn fmt(&self, f: &mut Formatter) -> FmtResult {
        write!(f, "FilePredicate({}, {})", self.desc, self.file)
    }
}

/// Represents a file's name and its data (either from an existing file or directly as Vec<u8>).
#[derive(Debug)]
pub enum FileIn {
    From(&'static str, &'static str),
    Bin(&'static str, Vec<u8>),
}

/// Create a temp (delete-on-drop) working dir, seeded with some data.
///
/// All files will be deleted when the returned TempDir goes out of scope (unless env contains
/// ASSERTFS_KEEP_TEMP, but that's a user-level feature), and the caller is responsible for making
/// sure that TempDir lives long enough.
// UPSTREAMING: Needs a better persist_if() API. The 'testfailed' and 'testok' files are workarounds
// for the lack of advanced persist strategies.
pub fn workdir_init(files: &[FileIn]) -> TempDir {
    // Create temp working dir and decide if it'll delete on Drop.
    let wd = TempDir::new().unwrap().persist_if(env::var("ASSERTFS_KEEP_TEMP").is_ok());
    if env::var("ASSERTFS_KEEP_TEMP").is_ok() {
        warn!("Temp files can be viewed in {}/", wd.path().display());
    } else {
        warn!("Temp files will be deleted, set ASSERTFS_KEEP_TEMP=1 in your shell to keep them.");
    }
    wd.child("testfailed").touch().expect("couldn't touch testfailed");

    // Setup the input files, if any.
    for file in files {
        match file {
            FileIn::From(n, s) => wd.child(n)
                                    .write_file(&Path::new("test").join(s))
                                    .expect(&format!("Couldn't copy test/{:?}", s)),
            FileIn::Bin(n, s) => {
                wd.child(n).write_binary(s).expect(&format!("Couldn't write_binary {:?}", s))
            },
        }
    }

    // Done.
    wd
}

/// Check file contents inside workdir.
// TODO: Provide a nice way to check that no other file was created.
pub fn workdir_check(wd: TempDir, checks: &[FilePredicate]) {
    for c in checks {
        let mut sb = String::new();
        let f = wd.child(&c.file);
        let path = f.path();
        File::open(path).expect(&format!("Couldn't open {:?}", path))
                        .read_to_string(&mut sb)
                        .expect(&format!("Couldn't read {:?} as utf8", path));
        assert!((c.pred)(&sb), "Check {:?} failed: file={} content=>>>>{}<<<<", c.desc, c.file, sb);
    }
    wd.child("testok").touch().expect("couldn't touch testok");
    std::fs::remove_file(wd.child("testfailed").path()).expect("couldn't remove testfailed");
}

/// Convenience type to give program arguments as either a `&str` (will get split by whitespace) or
/// a `&[&str]` (Note that this is a slice, not an array. The easyest way to get create that is
/// `["a","b"].as_ref()`.).
pub struct ExecArgs(Vec<String>);
impl From<&str> for ExecArgs {
    fn from(s: &str) -> Self {
        ExecArgs(s.split_whitespace().map(str::to_string).collect())
    }
}
impl From<String> for ExecArgs {
    fn from(s: String) -> Self {
        ExecArgs(s.split_whitespace().map(str::to_string).collect())
    }
}
impl From<&[&str]> for ExecArgs {
    fn from(s: &[&str]) -> Self {
        ExecArgs(s.iter().map(|s| s.to_string()).collect())
    }
}

#[derive(Debug)]
enum CmdState {
    Wait(CmdCond),
    Started(Child, Instant),
    Done,
}

/// This is used by the run loop to decide when to start or signal a `Cmd`.
#[derive(Debug)]
pub enum CmdCond {
    /// Immediately true
    None,
    /// Duration elapsed
    Delay(Duration),
    /// Other Cmd exited
    Cmd(String),
    /// File-based predicate
    Predicate(FilePredicate),
}
impl CmdCond {
    fn check(&self, wd: &TempDir, elapsed: Duration, done: &[String]) -> bool {
        match self {
            CmdCond::None => true,
            CmdCond::Delay(d) => elapsed > *d,
            CmdCond::Cmd(c) => done.contains(c),
            CmdCond::Predicate(p) => {
                let mut sb = String::new();
                let f = wd.child(&p.file);
                let path = f.path();
                File::open(path).map(|mut f| f.read_to_string(&mut sb))
                                .map(|_| (p.pred)(&sb))
                                .unwrap_or(false)
            },
        }
    }
}
// UPSTREAMING: These are handy shortcuts (for example `cmd.after(("filename","regex"))`) but hard
// to discover, may need more verbose wrappers just for readability's sake.
impl From<u64> for CmdCond {
    fn from(d: u64) -> Self {
        CmdCond::Delay(Duration::from_millis(d))
    }
}
impl From<&str> for CmdCond {
    fn from(c: &str) -> Self {
        CmdCond::Cmd(String::from(c))
    }
}
impl From<(&str, &str)> for CmdCond {
    fn from(file_regex: (&str, &str)) -> Self {
        let re = Regex::new(file_regex.1).unwrap();
        CmdCond::Predicate(FilePredicate::new("match", file_regex.0, move |s| re.is_match(s)))
    }
}

/// This "unittest" actually behaves like the main binary (except some extra output on stdout, and
/// assuming test code doesn't seep into normal code). Use "--" to separate the test binary's
/// arguments from the fakemain binary's arguments:
///
/// `cargo test -- fakemain --ignored --nocapture -- the real args`
// UPSTREAMING: This is a kludge. It's a very useful one because it doesn't compile the normal
// binary at test runtime, saving time but more importantly being able to run tests on a host
// without a compiler. The proper (slower, less crosscompile-friendly) way to to use `lazy_static`
// and `escargot`:
// lazy_static! {
//     static ref CARGO_RUN: CargoRun = CargoBuild::new().bin(env!("CARGO_PKG_NAME"))
//                                                       .current_release()
//                                                       .current_target()
//                                                       .run()
//                                                       .unwrap();
// }
// But maybe everybody's need is different, and `asset_cmd` should just have doc examples of
// `Cmd.main()` instead of an actual implmentation.
#[test]
#[ignore]
fn fakemain() {
    use structopt::StructOpt;
    main_(Opt::from_iter(std::env::args_os().skip_while(|a| a != "--")))
}

#[derive(Debug)]
pub struct Cmd {
    /// Program to run.
    bin: Command,
    /// Command name for `after()`, std/err filname prefix, and logs.
    name: String,
    /// Expected exit status. Some(0)=success, Some(n)=failure, None=timeout.
    exit: Option<i32>,
    /// Fail if the cmd exits too early.
    mintime: Duration,
    /// Fail if the cmd run for too long.
    maxtime: Duration,
    /// Current state.
    state: CmdState,
    /// List of signals to send to the process after startup.
    signals: Vec<(CmdCond, c_int)>,
}
impl Cmd {
    /// Constructor using the crate's main binary.
    ///
    /// This actually runs the test binary's "fakemain" test, to save on compiling and simplify
    /// cross-compilation.
    pub fn main(name: &str, args: impl Into<ExecArgs>) -> Self {
        let mut bin = Command::new(std::env::args_os().nth(0).unwrap());
        bin.args(&["fakemain", "--ignored", "--nocapture", "--"]);
        bin.args(args.into().0);
        Cmd { bin,
              name: String::from(name),
              exit: Some(0),
              mintime: Duration::from_millis(0),
              maxtime: Duration::from_millis(2000),
              state: CmdState::Wait(CmdCond::None),
              signals: Vec::new() }
    }
    /// Constructor using any binary.
    pub fn any(name: &str, bin: &str, args: impl Into<ExecArgs>) -> Self {
        let mut bin = Command::new(bin);
        bin.args(args.into().0);
        Cmd { bin,
              name: String::from(name),
              exit: Some(0),
              mintime: Duration::from_millis(0),
              maxtime: Duration::from_millis(2000),
              state: CmdState::Wait(CmdCond::None),
              signals: Vec::new() }
    }
    // UPSTREAMING: I also have some specialized contructors out of scope for assert_cmd, but we
    // should make sure that the the API makes extensions like this easy.
    pub fn myconstructor(name: impl Into<String>, foo: &str, after: impl Into<String>) -> Self {
        unimplemented!()
    }
    /// Constructor using a shell (bash) command.
    // UPSTREAMING: No attempt was made to make this work on Windows or other non-*nix OSes.
    pub fn bash(name: &str, script: &str) -> Self {
        let mut bin = Command::new("bash");
        bin.args(&["-c", script]);
        Cmd { bin,
              name: String::from(name),
              exit: Some(0),
              mintime: Duration::from_millis(0),
              maxtime: Duration::from_millis(2000),
              state: CmdState::Wait(CmdCond::None),
              signals: Vec::new() }
    }
    /// Append more args to be passed to the binary.
    pub fn args(mut self, args: impl Into<ExecArgs>) -> Self {
        self.bin.args(args.into().0);
        self
    }
    /// Set env variable.
    pub fn env(mut self, key: impl AsRef<OsStr>, val: impl AsRef<OsStr>) -> Self {
        self.bin.env(key, val);
        self
    }
    /// Set env variable if not already set.
    pub fn env_default(mut self, key: impl AsRef<OsStr>, val: impl AsRef<OsStr>) -> Self {
        if env::var(&key).is_err() {
            self.bin.env(key, val);
        }
        self
    }
    /// Unset env variable.
    pub fn env_remove(mut self, key: impl AsRef<OsStr>) -> Self {
        self.bin.env_remove(key);
        self
    }
    /// Expected ExitStatus (default `Some(0)`). This will be the program's return value if it
    /// finishes voluntarily, the signal number if killed by a signal, and None if it timed out.
    // UPSTREAMING: The dual-meaning return_value/signal_number is handy but not very clean.
    pub fn exit(mut self, e: impl Into<Option<i32>>) -> Self {
        self.exit = e.into();
        self
    }
    /// Fail if programs exits before `t` ms (defaults to 0).
    pub fn mintime(mut self, t: u64) -> Self {
        self.mintime = Duration::from_millis(t);
        self
    }
    /// Fail if programs runs for longer than `t` ms (defaults to 2000).
    pub fn maxtime(mut self, t: u64) -> Self {
        self.maxtime = Duration::from_millis(t);
        self
    }
    /// Send signal to the program after CmdCond. Call this multiple times to send multiple signals.
    pub fn signal(mut self, sig: c_int, cond: impl Into<CmdCond>) -> Self {
        self.signals.push((cond.into(), sig));
        self
    }
    /// Wait for `CmdCond` before starting this one.
    pub fn after(mut self, cond: impl Into<CmdCond>) -> Self {
        self.state = CmdState::Wait(cond.into());
        self
    }
    // UPSTREAMING: Like myconstructor(), some setters are always the same and the API should
    // accomodate that. For example, delaying start of a command until a log file matches a regex.
    pub fn mysetter(mut self) -> Self {
        unimplemented!()
    }
}
impl Drop for Cmd {
    /// Make sure the process dies when the struct is droped. Needed when the Cmd timed out, when
    /// another Cmd triggered an assert, etc.
    fn drop(&mut self) {
        if let CmdState::Started(ref mut child, _) = self.state {
            if let None = child.try_wait().unwrap() {
                warn!("Cmd {} process alive when struct droped, killing {}", self.name, child.id());
                unsafe { kill(child.id() as i32, SIGKILL) };
            }
        }
    }
}

/// Run a list of external commands in a controlled environement, and check result.
/// See `Cmd`, `FileIn`, `FilePredicate`, and unittests for docs and examples.
// TODO: Support pipes (unless emulating by calling `sh -c 'cmd1|cmd2'` is enough).
#[derive(Debug)]
pub struct Exec {
    cmds: Vec<Cmd>,
    timeout: Duration,
    files: Vec<FileIn>,
    checks: Vec<FilePredicate>,
}
/// Create a new empty `Exec` (defined as a standalone fn to save on typing).
pub fn exec() -> Exec {
    log_init();
    Exec { cmds: vec![], timeout: Duration::from_millis(10000), files: vec![], checks: vec![] }
}
impl Exec {
    /// Add a new `Cmd` to the run list.
    pub fn cmd(mut self, c: Cmd) -> Self {
        self.cmds.push(c);
        self
    }
    /// Seed the run env with an existing file from the `test/` folder.
    pub fn infile(mut self, name: &'static str, src: &'static str) -> Self {
        self.files.push(FileIn::From(name, src));
        self
    }
    /// Seed the run env with a file with the given content.
    pub fn inbytes(mut self, name: &'static str, src: impl Into<Vec<u8>>) -> Self {
        self.files.push(FileIn::Bin(name, src.into()));
        self
    }
    /// Check file contents after programs have run. The files `{N}.out` and `{N}.err` correspond to
    /// the stdout and stderr of the command N (see `Cmd::set_name()`).
    pub fn check(mut self,
                 desc: &'static str, // Desciption for assert-logging purpose.
                 file: &'static str, // Which file to operate on.
                 pred: impl Fn(&str) -> bool + 'static // Closure that tests the content of the file.
    ) -> Self {
        self.checks.push(FilePredicate::new(desc, file, pred));
        self
    }
    /// Like `check()` but for the common usecase of checking against a regex.
    pub fn check_re(mut self,
                    desc: &'static str, // Desciption for assert-logging purpose.
                    file: &'static str, // Which file to operate on.
                    regex: &str         // Regex that the file content has to match.
    ) -> Self {
        let re = Regex::new(regex).unwrap();
        self.checks.push(FilePredicate::new(desc, file, move |s| re.is_match(s)));
        self
    }
    /// Like `check_re()` but checks that the number of matches is within a specified range.
    // UPSTREAMING: Should probably merge with `check_re()`.
    pub fn check_re_n(mut self,
                      desc: &'static str, // Description for assert-logging purpose.
                      file: &'static str, // Which file to operate on.
                      regex: &str,        // Regex that the file content has to match.
                      n: Range<usize>)
                      -> Self {
        let re = Regex::new(regex).unwrap();
        self.checks
            .push(FilePredicate::new(desc, file, move |s| n.contains(&re.find_iter(s).count())));
        self
    }
    /// Set overall run timeout (default 10s).
    pub fn timeout(mut self, t: u64) -> Self {
        self.timeout = Duration::from_millis(t);
        self
    }

    /// Actually run the test (setup files, launch programs, check outputs...).
    pub fn run(mut self) {
        // Init
        let wd = workdir_init(&self.files);
        let startall = Instant::now();
        // TODO: how can we get the test name ?
        let mut testdesc = File::create(wd.child("testdesc").path()).unwrap();
        writeln!(testdesc, "{:?}", self).unwrap();

        // Check that we're not waiting for an unknown cmd or self
        let names: Vec<String> = self.cmds.iter().map(|c| c.name.clone()).collect();
        for c in self.cmds.iter() {
            if let CmdState::Wait(CmdCond::Cmd(ref w)) = c.state {
                assert!(names.contains(w), "cmd {} start after {}, {:?}", c.name, w, names);
                assert!(*w != c.name, "cmd {} start after itself", c.name);
            }
            for s in c.signals.iter() {
                if let (CmdCond::Cmd(ref w), _) = s {
                    assert!(names.contains(w), "cmd {} signal after {}, {:?}", c.name, w, names);
                    assert!(*w != c.name, "cmd {} signal after itself", c.name);
                }
            }
        }

        // Until all Cmds are `Done`, handle each one in turn depending on their state.
        let mut done: Vec<String> = vec![];
        while self.cmds.len() > done.len() {
            let now = Instant::now();
            assert!(now < startall + self.timeout, "overall run timeout {:?}", now - startall);
            for c in self.cmds.iter_mut() {
                trace!("{:?} {:?}", now, c);
                match c.state {
                    // Cmd is ready to execute, spawn it as a child process.
                    CmdState::Wait(ref cond) if cond.check(&wd, now - startall, &done) => {
                        let so = File::create(wd.child(format!("{}.out", c.name)).path()).unwrap();
                        let se = File::create(wd.child(format!("{}.err", c.name)).path()).unwrap();
                        let cmd = c.bin
                                   .current_dir(wd.path())
                                   .stdin(Stdio::null()) // UPSTREAMING: should have more options for stdin
                                   .stdout(Stdio::from(so))
                                   .stderr(Stdio::from(se));
                        info!("Cmd {} start: {:?}", c.name, cmd);
                        c.state = CmdState::Started(cmd.spawn().expect("Failed to start"), now);
                    },
                    // Check if Cmd has exited with expected code, or timeouted
                    CmdState::Started(ref mut child, started) => {
                        let elapsed = now - started;
                        if let Some(exit) = child.try_wait().unwrap() {
                            info!("Cmd {} exit: {:?} after {:?}", c.name, exit, elapsed);
                            assert!(elapsed > c.mintime, "{} exit before {:?}", c.name, c.mintime);
                            assert!(c.exit == exit.signal().or(exit.code()),
                                    "{} expected exit {:?} got {:?}",
                                    c.name,
                                    c.exit,
                                    exit);
                            c.state = CmdState::Done;
                            done.push(c.name.clone());
                        } else if elapsed > c.maxtime {
                            info!("Cmd {} timeout after {:?}", c.name, elapsed);
                            assert!(None == c.exit,
                                    "{} expected exit {:?} got timeout",
                                    c.name,
                                    c.exit);
                            c.state = CmdState::Done;
                            done.push(c.name.clone());
                        } else {
                            // TODO: refactor with https://github.com/rust-lang/rust/issues/43244
                            let mut i = 0;
                            while i < c.signals.len() {
                                if c.signals[i].0.check(&wd, elapsed, &done) {
                                    let s = c.signals[i].1;
                                    info!("Cmd {} sending signal {} to {}", c.name, s, child.id());
                                    c.signals.remove(i);
                                    // std::process::Child::kill() exists, but doesn't allow specifying which signal.
                                    unsafe { assert_eq!(0, kill(child.id() as i32, s), "kill()") };
                                } else {
                                    i += 1;
                                }
                            }
                        }
                    },
                    // Nothing to do: Cmd either Done or pending
                    _ => (),
                }
            }
            // UPSTREAMING: Yeah, that's lazy coding. Should use futures and async instead.
            sleep(Duration::from_millis(10));
        }
        // Check results
        workdir_check(wd, &self.checks);
    }
}

// UPSTREAMING: Another pattern that I can't share the code for is elaborate builders for the
// `FilePredicate.pred` closure.