goldentests 1.4.4

use crate::{
    config::TestConfig,
    diff_printer::DiffPrinter,
    error::{InnerTestError, TestResult},
    glob::Globs,
};

use colored::Colorize;
use similar::TextDiff;

#[cfg(feature = "parallel")]
use rayon::iter::IntoParallelIterator;
#[cfg(feature = "parallel")]
use rayon::iter::ParallelIterator;

#[cfg(feature = "progress-bar")]
use indicatif::ProgressBar;

use std::{
    fs::File,
    io::{Read, Write},
    path::{Path, PathBuf},
    process::{Command, Output},
};

type InnerTestResult<T> = Result<T, InnerTestError>;

/// Data captured during a parallel test run that may be needed to overwrite
/// a failing test file later if `--interactive` was passed.
struct PendingOverwrite {
    file: PathBuf,
    output: Output,
    test: Test,
}

struct Test {
    path: PathBuf,
    command_line_args: String,
    command_line_args_after: String,
    expected_stdout: String,
    expected_stderr: String,
    expected_exit_status: Option<i32>,
    rest: String,
}

#[derive(PartialEq)]
enum TestParseState {
    Neutral,
    ReadingExpectedStdout,
    ReadingExpectedStderr,
}

fn find_tests(test_path: &Path, globs: &Globs) -> (Vec<PathBuf>, Vec<InnerTestError>) {
    let mut tests = vec![];
    let mut errors = vec![];

    if test_path.is_dir() {
        let read_dir = match std::fs::read_dir(test_path) {
            Ok(dir) => dir,
            Err(err) => return (tests, vec![InnerTestError::IoError(test_path.to_owned(), err)]),
        };

        for entry in read_dir {
            let path = match entry {
                Ok(entry) => entry.path(),
                Err(err) => {
                    errors.push(InnerTestError::IoError(test_path.to_owned(), err));
                    continue;
                }
            };

            if path.is_dir() {
                let (mut more_tests, mut more_errors) = find_tests(&path, globs);
                tests.append(&mut more_tests);
                errors.append(&mut more_errors);
            } else if globs.is_match(&path) {
                tests.push(path);
            }
        }
    } else if globs.is_match(&test_path) {
        tests.push(test_path.into());
    }

    (tests, errors)
}

fn strip_prefix<'a>(s: &'a str, prefix: &str) -> &'a str {
    s.strip_prefix(prefix).unwrap_or(s)
}

fn append_line(s: &mut String, line: &str) {
    *s += line;
    *s += "\n";
}

fn parse_test(test_path: &Path, config: &TestConfig) -> InnerTestResult<Test> {
    let mut command_line_args = String::new();
    let mut command_line_args_after = String::new();
    let mut expected_stdout = String::new();
    let mut expected_stderr = String::new();
    let mut expected_exit_status = None;
    let mut rest = String::new();

    let mut file = File::open(test_path).map_err(|err| InnerTestError::IoError(test_path.to_owned(), err))?;
    let mut contents = String::new();
    file.read_to_string(&mut contents)
        .map_err(|err| InnerTestError::IoError(test_path.to_owned(), err))?;

    let mut state = TestParseState::Neutral;
    for line in contents.lines() {
        if line.starts_with(&config.test_line_prefix) {
            // If we're currently reading stdout or stderr, append the line to the expected output
            if state == TestParseState::ReadingExpectedStdout {
                append_line(&mut expected_stdout, strip_prefix(line, &config.test_line_prefix))
            } else if state == TestParseState::ReadingExpectedStderr {
                append_line(&mut expected_stderr, strip_prefix(line, &config.test_line_prefix));

            // Otherwise, look to see if the line begins with a keyword and if so change state
            // (stdout/stderr) or parse an argument to the keyword (args/exit status).

            // args:
            } else if line.starts_with(&config.test_args_prefix) {
                command_line_args = strip_prefix(line, &config.test_args_prefix).to_string();

            // args after:
            } else if line.starts_with(&config.test_args_after_prefix) {
                command_line_args_after = strip_prefix(line, &config.test_args_after_prefix).to_string();

            // expected stdout:
            } else if line.starts_with(&config.test_stdout_prefix) {
                state = TestParseState::ReadingExpectedStdout;
                // Append the remainder of the line to the expected stdout.
                // Both expected_stdout and expected_stderr are trimmed so it
                // has no effect if the rest of this line is empty
                append_line(&mut expected_stdout, strip_prefix(line, &config.test_stdout_prefix));

            // expected stderr:
            } else if line.starts_with(&config.test_stderr_prefix) {
                state = TestParseState::ReadingExpectedStderr;
                append_line(&mut expected_stderr, strip_prefix(line, &config.test_stderr_prefix));

            // expected exit status:
            } else if line.starts_with(&config.test_exit_status_prefix) {
                let status = strip_prefix(line, &config.test_exit_status_prefix).trim();
                expected_exit_status = Some(status.parse().map_err(|err| {
                    InnerTestError::ErrorParsingExitStatus(test_path.to_owned(), status.to_owned(), err)
                })?);
            } else {
                append_line(&mut rest, line);
            }
        } else {
            // Both expected_stdout and expected_stderr need a blank line at the end,
            // the order here implicitly skips that newline.
            if state == TestParseState::Neutral {
                append_line(&mut rest, line);
            }
            state = TestParseState::Neutral;
        }
    }

    // Remove \r from strings for windows compatibility. This means we
    // also can't test for any string containing "\r" unless this check
    // is improved to be more clever (e.g. only removing at the end of a line).
    let expected_stdout = expected_stdout.replace("\r", "");
    let expected_stderr = expected_stderr.replace("\r", "");

    Ok(Test {
        path: test_path.to_owned(),
        command_line_args,
        command_line_args_after,
        expected_stdout,
        expected_stderr,
        expected_exit_status,
        rest,
    })
}

fn write_expected_output_for_stream(
    file: &mut File,
    prefix: &str,
    marker: &str,
    expected: &[u8],
) -> std::io::Result<()> {
    // Doesn't handle \r correctly!
    // Strip leading and trailing newlines from the output
    let expected_stdout = String::from_utf8_lossy(expected).replace("\r", "");
    let lines: Vec<&str> = expected_stdout.trim().split('\n').collect();
    match lines.len() {
        // Don't write if there's nothing to write
        0 => Ok(()),
        1 if lines[0].len() == 0 => Ok(()),
        // If the line is short and nice, write that line
        1 if lines[0].len() < 80 => {
            writeln!(file, "{} {}", marker, lines[0])
        }
        // Otherwise we write it more longform
        _ => {
            writeln!(file, "{}", marker)?;
            for line in lines {
                file.write_all(prefix.as_bytes())?;
                file.write_all(line.as_bytes())?;
                writeln!(file, "")?;
            }
            writeln!(file, "")
        }
    }
}

fn overwrite_test(test_path: &PathBuf, config: &TestConfig, output: &Output, test: &Test) -> std::io::Result<()> {
    // Maybe copy the file so we don't remove it if we fail here?
    let mut file = File::create(test_path)?;

    file.write_all(test.rest.trim_end().as_bytes())?;
    writeln!(file, "")?;
    writeln!(file, "")?;

    if !test.command_line_args.is_empty() {
        writeln!(file, "{} {}", config.test_args_prefix, test.command_line_args.trim())?;
    }

    if !test.command_line_args_after.is_empty() {
        writeln!(
            file,
            "{} {}",
            config.test_args_after_prefix,
            test.command_line_args_after.trim()
        )?;
    }

    if Some(0) != output.status.code() {
        writeln!(
            file,
            "{} {}",
            config.test_exit_status_prefix,
            output.status.code().unwrap_or(0)
        )?;
    }

    write_expected_output_for_stream(
        &mut file,
        &config.test_line_prefix,
        &config.test_stdout_prefix,
        &output.stdout,
    )?;
    // Write an empty line between expected stdout and stderr to separate
    // the two multi-line outputs if necessary
    if !output.stdout.is_empty() && !output.stderr.is_empty() {
        writeln!(file, "")?;
    }
    write_expected_output_for_stream(
        &mut file,
        &config.test_line_prefix,
        &config.test_stderr_prefix,
        &output.stderr,
    )
}

/// Diff the given "stream" and expected contents of the stream.
/// Returns non-zero on error.
fn check_for_differences_in_stream(name: &str, stream: &[u8], expected: &str, errors: &mut Vec<String>) {
    let output_string = String::from_utf8_lossy(stream).replace("\r", "");
    let output = output_string.trim();
    let expected = expected.trim();

    let differences = TextDiff::from_lines(expected, output);
    if differences.ratio() != 1.0 {
        errors.push(format!(
            "Actual {} differs from expected {}:\n{}",
            name,
            name,
            DiffPrinter(differences)
        ));
    }
}

fn check_exit_status(output: &Output, expected_status: Option<i32>, errors: &mut Vec<String>) {
    if let Some(expected_status) = expected_status {
        if let Some(actual_status) = output.status.code() {
            if expected_status != actual_status {
                errors.push(format!(
                    "Expected an exit status of {} but process returned {}\n",
                    expected_status, actual_status,
                ));
            }
        } else {
            errors.push(format!(
                "Expected an exit status of {} but process was terminated by signal instead\n",
                expected_status
            ));
        }
    }
}

fn check_for_differences(path: &Path, output: &Output, test: &Test) -> InnerTestResult<()> {
    let mut errors = vec![];
    check_exit_status(output, test.expected_exit_status, &mut errors);
    check_for_differences_in_stream("stdout", &output.stdout, &test.expected_stdout, &mut errors);
    check_for_differences_in_stream("stderr", &output.stderr, &test.expected_stderr, &mut errors);

    if errors.is_empty() {
        Ok(())
    } else {
        let path = path.to_owned();
        Err(InnerTestError::TestFailed { path, errors })
    }
}

#[cfg(feature = "parallel")]
fn into_iter<T: IntoParallelIterator>(value: T) -> T::Iter {
    value.into_par_iter()
}

#[cfg(not(feature = "parallel"))]
fn into_iter<T: IntoIterator>(value: T) -> T::IntoIter {
    value.into_iter()
}

impl TestConfig {
    fn test_all(&self, test_sources: Vec<PathBuf>) -> Vec<(InnerTestResult<()>, Option<PendingOverwrite>)> {
        #[cfg(feature = "progress-bar")]
        let progress = ProgressBar::new(test_sources.len() as u64);

        let results = into_iter(test_sources)
            .map(|file| {
                #[cfg(feature = "progress-bar")]
                progress.inc(1);
                let run = || -> InnerTestResult<(InnerTestResult<()>, Option<PendingOverwrite>)> {
                    let test = parse_test(&file, self)?;

                    let mut args = Self::split_args(&self.base_args, &file)?;
                    args.extend(Self::split_args(&test.command_line_args, &file)?);

                    args.push(test.path.to_string_lossy().to_string());

                    args.extend(Self::split_args(&self.base_args_after, &file)?);
                    args.extend(Self::split_args(&test.command_line_args_after, &file)?);

                    let mut command = Command::new(&self.binary_path);
                    command.args(args);

                    let output =
                        command.output().map_err(|err| InnerTestError::CommandError(file.clone(), command, err))?;

                    let differences = check_for_differences(&test.path, &output, &test);
                    if self.overwrite_tests {
                        if let Err(InnerTestError::TestFailed { path, errors }) = differences {
                            overwrite_test(&file, self, &output, &test)
                                .map_err(|err| InnerTestError::IoError(file.to_owned(), err))?;

                            return Ok((Err(InnerTestError::TestUpdated { path, errors }), None));
                        }
                    } else if self.interactive {
                        if let Err(InnerTestError::TestFailed { .. }) = &differences {
                            // Defer the overwrite decision to the serial review phase in `run_tests`.
                            return Ok((differences, Some(PendingOverwrite { file, output, test })));
                        }
                    }
                    Ok((differences, None))
                };

                // Lift any early-return error into the result slot so it is reported like before.
                run().unwrap_or_else(|err| (Err(err), None))
            })
            .collect();

        #[cfg(feature = "progress-bar")]
        progress.finish_and_clear();
        results
    }

    /// Splits a string into separate command-line args.
    /// Usually this means separating by spaces.
    fn split_args(s: &str, file: &Path) -> Result<Vec<String>, InnerTestError> {
        shlex::split(s).ok_or_else(|| InnerTestError::ErrorParsingArgs(file.to_path_buf(), s.to_owned()))
    }

    /// Recurse through all the files in self.path, parse them all,
    /// and run the target program with the arguments specified in the file.
    pub fn run_tests(&self) -> TestResult<()> {
        let globs = Globs::new(&self.glob)?;
        let (tests, path_errors) = find_tests(&self.test_path, &globs);
        let mut outputs = self.test_all(tests);

        for error in path_errors {
            eprintln!("{}", error);
        }

        if self.interactive {
            self.review_interactively(&mut outputs);
        }

        let total_tests = outputs.len();
        let mut failing_tests = 0;
        let mut can_be_fixed_with_overwrite_tests = 0;
        let mut updated_tests = 0;
        for (result, _) in &outputs {
            match result {
                Ok(_) => {}
                Err(InnerTestError::TestUpdated { .. }) => {
                    updated_tests += 1;
                }

                Err(InnerTestError::TestFailed { .. }) => {
                    can_be_fixed_with_overwrite_tests += 1;
                    failing_tests += 1;
                }

                Err(
                    InnerTestError::IoError(_, _)
                    | InnerTestError::CommandError(_, _, _)
                    | InnerTestError::ErrorParsingExitStatus(_, _, _)
                    | InnerTestError::ErrorParsingArgs(_, _),
                ) => {
                    failing_tests += 1;
                }
            }

            if let Err(err) = result {
                let already_shown = self.interactive
                    && matches!(
                        err,
                        InnerTestError::TestFailed { .. } | InnerTestError::TestUpdated { .. }
                    );
                if !already_shown {
                    eprintln!("{}", err)
                }
            }
        }

        if updated_tests == 0 && !self.overwrite_tests {
            println!(
                "ran {} {} tests with {} and {}\n",
                total_tests,
                "golden".bright_yellow(),
                format!("{} passing", total_tests - failing_tests).green(),
                format!("{} failing", failing_tests).red(),
            );
        } else {
            println!(
                "ran {} {} tests with {}, {} and {}\n",
                total_tests,
                "golden".bright_yellow(),
                format!("{} passing", total_tests - failing_tests).green(),
                format!("{} failing", failing_tests).red(),
                format!("{} updated", updated_tests).cyan(),
            );
        }

        if can_be_fixed_with_overwrite_tests > 0 && !self.interactive {
            println!("Looks like you have failing tests. Review the output of each and fix any unexpected differences. When finished, you can use the --overwrite flag to automatically write the new output to all {} failing test file(s) or the --interactive (-i) flag to review each file's diff one by one.", can_be_fixed_with_overwrite_tests);
        }

        if failing_tests != 0 {
            Err(())
        } else {
            Ok(())
        }
    }

    /// Review each failing test one by one, letting the user accept or reject the change, quit,
    /// or accept all changes. Accepted changes will overwrite the source file's expected output with the new output.
    fn review_interactively(&self, outputs: &mut [(InnerTestResult<()>, Option<PendingOverwrite>)]) {
        let mut accept_all = false;

        for (result, pending) in outputs.iter_mut() {
            let pending = match pending.take() {
                Some(pending) => pending,
                None => continue,
            };

            if let Err(file_diff) = result {
                println!("{file_diff}");
            }

            let accept = accept_all || {
                match prompt_review() {
                    ReviewChoice::Yes => true,
                    ReviewChoice::No => false,
                    ReviewChoice::AcceptAll => {
                        accept_all = true;
                        true
                    }
                    ReviewChoice::Quit => break,
                }
            };

            if accept {
                match overwrite_test(&pending.file, self, &pending.output, &pending.test) {
                    Ok(()) => {
                        if let Err(InnerTestError::TestFailed { path, errors }) = result {
                            *result = Err(InnerTestError::TestUpdated {
                                path: std::mem::take(path),
                                errors: std::mem::take(errors),
                            });
                        }
                    }
                    Err(err) => {
                        *result = Err(InnerTestError::IoError(pending.file.clone(), err));
                    }
                }
            }
        }
    }
}

enum ReviewChoice {
    Yes,
    No,
    Quit,
    AcceptAll,
}

/// Prompt the user on stdin for how to handle the current failing test file
fn prompt_review() -> ReviewChoice {
    use std::io::BufRead;

    let stdin = std::io::stdin();
    loop {
        print!("Accept these changes? [y]es / [n]o / [q]uit / [a]ccept all: ");
        // Ensure the above is printed before we ask for input below
        let _ = std::io::stdout().flush();

        let mut line = String::new();
        match stdin.lock().read_line(&mut line) {
            Ok(0) => return ReviewChoice::No, // EOF
            Ok(_) => match line.trim() {
                "y" | "Y" | "yes" => return ReviewChoice::Yes,
                "n" | "N" | "no" => return ReviewChoice::No,
                "q" | "Q" | "quit" => return ReviewChoice::Quit,
                "a" | "A" | "accept" | "accept all" => return ReviewChoice::AcceptAll,
                _ => continue,
            },
            Err(_) => return ReviewChoice::No,
        }
    }
}