use std::ffi;
use std::io;
use std::io::{Read, Write};
use std::ops::Deref;
use std::path;
use std::process;
use crate::assert::Assert;
use crate::assert::OutputAssertExt;
use crate::output::DebugBuffer;
use crate::output::DebugBytes;
use crate::output::OutputError;
use crate::output::OutputOkExt;
use crate::output::OutputResult;
#[derive(Debug)]
pub struct Command {
cmd: process::Command,
stdin: Option<bstr::BString>,
timeout: Option<std::time::Duration>,
}
impl Command {
pub fn from_std(cmd: process::Command) -> Self {
Self {
cmd,
stdin: None,
timeout: None,
}
}
pub fn cargo_bin<S: AsRef<str>>(name: S) -> Result<Self, crate::cargo::CargoError> {
let cmd = crate::cargo::cargo_bin_cmd(name)?;
Ok(Self::from_std(cmd))
}
pub fn write_stdin<S>(&mut self, buffer: S) -> &mut Self
where
S: Into<Vec<u8>>,
{
self.stdin = Some(bstr::BString::from(buffer.into()));
self
}
pub fn timeout(&mut self, timeout: std::time::Duration) -> &mut Self {
self.timeout = Some(timeout);
self
}
pub fn pipe_stdin<P>(&mut self, file: P) -> io::Result<&mut Self>
where
P: AsRef<path::Path>,
{
let buffer = std::fs::read(file)?;
Ok(self.write_stdin(buffer))
}
pub fn ok(&mut self) -> OutputResult {
OutputOkExt::ok(self)
}
pub fn unwrap(&mut self) -> process::Output {
OutputOkExt::unwrap(self)
}
pub fn unwrap_err(&mut self) -> OutputError {
OutputOkExt::unwrap_err(self)
}
pub fn assert(&mut self) -> Assert {
OutputAssertExt::assert(self)
}
}
impl Command {
pub fn new<S: AsRef<ffi::OsStr>>(program: S) -> Self {
let cmd = process::Command::new(program);
Self::from_std(cmd)
}
pub fn arg<S: AsRef<ffi::OsStr>>(&mut self, arg: S) -> &mut Self {
self.cmd.arg(arg);
self
}
pub fn args<I, S>(&mut self, args: I) -> &mut Self
where
I: IntoIterator<Item = S>,
S: AsRef<ffi::OsStr>,
{
self.cmd.args(args);
self
}
pub fn env<K, V>(&mut self, key: K, val: V) -> &mut Self
where
K: AsRef<ffi::OsStr>,
V: AsRef<ffi::OsStr>,
{
self.cmd.env(key, val);
self
}
pub fn envs<I, K, V>(&mut self, vars: I) -> &mut Self
where
I: IntoIterator<Item = (K, V)>,
K: AsRef<ffi::OsStr>,
V: AsRef<ffi::OsStr>,
{
self.cmd.envs(vars);
self
}
pub fn env_remove<K: AsRef<ffi::OsStr>>(&mut self, key: K) -> &mut Self {
self.cmd.env_remove(key);
self
}
pub fn env_clear(&mut self) -> &mut Self {
self.cmd.env_clear();
self
}
pub fn current_dir<P: AsRef<path::Path>>(&mut self, dir: P) -> &mut Self {
self.cmd.current_dir(dir);
self
}
pub fn output(&mut self) -> io::Result<process::Output> {
let spawn = self.spawn()?;
Self::wait_with_input_output(spawn, self.stdin.as_deref().cloned(), self.timeout)
}
fn wait_with_input_output(
mut child: process::Child,
input: Option<Vec<u8>>,
timeout: Option<std::time::Duration>,
) -> io::Result<process::Output> {
let stdin = input.and_then(|i| {
child
.stdin
.take()
.map(|mut stdin| std::thread::spawn(move || stdin.write_all(&i)))
});
fn read<R>(mut input: R) -> std::thread::JoinHandle<io::Result<Vec<u8>>>
where
R: Read + Send + 'static,
{
std::thread::spawn(move || {
let mut ret = Vec::new();
input.read_to_end(&mut ret).map(|_| ret)
})
}
let stdout = child.stdout.take().map(read);
let stderr = child.stderr.take().map(read);
stdin.and_then(|t| t.join().unwrap().ok());
let status = if let Some(timeout) = timeout {
wait_timeout::ChildExt::wait_timeout(&mut child, timeout)
.transpose()
.unwrap_or_else(|| {
let _ = child.kill();
child.wait()
})
} else {
child.wait()
}?;
let stdout = stdout
.and_then(|t| t.join().unwrap().ok())
.unwrap_or_default();
let stderr = stderr
.and_then(|t| t.join().unwrap().ok())
.unwrap_or_default();
Ok(process::Output {
status,
stdout,
stderr,
})
}
fn spawn(&mut self) -> io::Result<process::Child> {
self.cmd.stdin(process::Stdio::piped());
self.cmd.stdout(process::Stdio::piped());
self.cmd.stderr(process::Stdio::piped());
self.cmd.spawn()
}
}
impl From<process::Command> for Command {
fn from(cmd: process::Command) -> Self {
Command::from_std(cmd)
}
}
impl<'c> OutputOkExt for &'c mut Command {
fn ok(self) -> OutputResult {
let output = self.output().map_err(OutputError::with_cause)?;
if output.status.success() {
Ok(output)
} else {
let error = OutputError::new(output).set_cmd(format!("{:?}", self.cmd));
let error = if let Some(stdin) = self.stdin.as_ref() {
error.set_stdin(stdin.deref().clone())
} else {
error
};
Err(error)
}
}
fn unwrap_err(self) -> OutputError {
match self.ok() {
Ok(output) => {
if let Some(stdin) = self.stdin.as_ref() {
panic!(
"Completed successfully:\ncommand=`{:?}`\nstdin=```{}```\nstdout=```{}```",
self.cmd,
DebugBytes::new(stdin),
DebugBytes::new(&output.stdout)
)
} else {
panic!(
"Completed successfully:\ncommand=`{:?}`\nstdout=```{}```",
self.cmd,
DebugBytes::new(&output.stdout)
)
}
}
Err(err) => err,
}
}
}
impl<'c> OutputAssertExt for &'c mut Command {
fn assert(self) -> Assert {
let output = match self.output() {
Ok(output) => output,
Err(err) => {
panic!("Failed to spawn {:?}: {}", self, err);
}
};
let assert = Assert::new(output).append_context("command", format!("{:?}", self.cmd));
if let Some(stdin) = self.stdin.as_ref() {
assert.append_context("stdin", DebugBuffer::new(stdin.deref().clone()))
} else {
assert
}
}
}