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
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151

use crate::{assert_not_contains, handle_failed_output};
use std::ffi::OsStr;
use std::io::Write;
use std::ops::{Deref, DerefMut};
use std::process::{Command as StdCommand, ExitStatus, Output, Stdio};

/// This is a custom command wrapper that simplifies working with commands
/// and makes it easier to ensure that we check the exit status of executed
/// processes.
#[derive(Debug)]
pub struct Command {
    cmd: StdCommand,
    stdin: Option<Box<[u8]>>,
}

impl Command {
    pub fn new<S: AsRef<OsStr>>(program: S) -> Self {
        Self { cmd: StdCommand::new(program), stdin: None }
    }

    pub fn set_stdin(&mut self, stdin: Box<[u8]>) {
        self.stdin = Some(stdin);
    }

    /// Run the constructed command and assert that it is successfully run.
    #[track_caller]
    pub fn run(&mut self) -> CompletedProcess {
        let caller_location = std::panic::Location::caller();
        let caller_line_number = caller_location.line();

        let output = self.command_output();
        if !output.status().success() {
            handle_failed_output(&self, output, caller_line_number);
        }
        output
    }

    /// Run the constructed command and assert that it does not successfully run.
    #[track_caller]
    pub fn run_fail(&mut self) -> CompletedProcess {
        let caller_location = std::panic::Location::caller();
        let caller_line_number = caller_location.line();

        let output = self.command_output();
        if output.status().success() {
            handle_failed_output(&self, output, caller_line_number);
        }
        output
    }

    #[track_caller]
    pub(crate) fn command_output(&mut self) -> CompletedProcess {
        // let's make sure we piped all the input and outputs
        self.cmd.stdin(Stdio::piped());
        self.cmd.stdout(Stdio::piped());
        self.cmd.stderr(Stdio::piped());

        let output = if let Some(input) = &self.stdin {
            let mut child = self.cmd.spawn().unwrap();

            {
                let mut stdin = child.stdin.take().unwrap();
                stdin.write_all(input.as_ref()).unwrap();
            }

            child.wait_with_output().expect("failed to get output of finished process")
        } else {
            self.cmd.output().expect("failed to get output of finished process")
        };
        output.into()
    }
}

impl Deref for Command {
    type Target = StdCommand;

    fn deref(&self) -> &Self::Target {
        &self.cmd
    }
}

impl DerefMut for Command {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.cmd
    }
}

/// Represents the result of an executed process.
/// The various `assert_` helper methods should preferably be used for
/// checking the contents of stdout/stderr.
pub struct CompletedProcess {
    output: Output,
}

impl CompletedProcess {
    pub fn stdout_utf8(&self) -> String {
        String::from_utf8(self.output.stdout.clone()).expect("stdout is not valid UTF-8")
    }

    pub fn stderr_utf8(&self) -> String {
        String::from_utf8(self.output.stderr.clone()).expect("stderr is not valid UTF-8")
    }

    pub fn status(&self) -> ExitStatus {
        self.output.status
    }

    /// Checks that trimmed `stdout` matches trimmed `content`.
    #[track_caller]
    pub fn assert_stdout_equals<S: AsRef<str>>(self, content: S) -> Self {
        assert_eq!(self.stdout_utf8().trim(), content.as_ref().trim());
        self
    }

    #[track_caller]
    pub fn assert_stdout_not_contains<S: AsRef<str>>(self, needle: S) -> Self {
        assert_not_contains(&self.stdout_utf8(), needle.as_ref());
        self
    }

    /// Checks that trimmed `stderr` matches trimmed `content`.
    #[track_caller]
    pub fn assert_stderr_equals<S: AsRef<str>>(self, content: S) -> Self {
        assert_eq!(self.stderr_utf8().trim(), content.as_ref().trim());
        self
    }

    #[track_caller]
    pub fn assert_stderr_contains<S: AsRef<str>>(self, needle: S) -> Self {
        assert!(self.stderr_utf8().contains(needle.as_ref()));
        self
    }

    #[track_caller]
    pub fn assert_stderr_not_contains<S: AsRef<str>>(self, needle: S) -> Self {
        assert_not_contains(&self.stdout_utf8(), needle.as_ref());
        self
    }

    #[track_caller]
    pub fn assert_exit_code(self, code: i32) -> Self {
        assert!(self.output.status.code() == Some(code));
        self
    }
}

impl From<Output> for CompletedProcess {
    fn from(output: Output) -> Self {
        Self { output }
    }
}