reifydb_testing/testscript/

runner.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2026 ReifyDB

#[allow(clippy::disallowed_types)]
use std::sync::Mutex;
use std::{
	collections::HashMap,
	env::temp_dir,
	error::Error,
	fs, io,
	io::Write as _,
	panic, path, process,
	sync::{Arc, LazyLock, PoisonError},
	time,
};

use fs::read_to_string;
use io::ErrorKind;
use panic::AssertUnwindSafe;
use path::Path;
use time::SystemTime;

use crate::{
	goldenfile::Mint,
	testscript::{
		command::{Block, Command},
		parser::parse,
	},
};

#[allow(clippy::disallowed_types)]
static SCRIPT_LOCKS: LazyLock<Mutex<HashMap<path::PathBuf, Arc<Mutex<()>>>>> =
	LazyLock::new(|| Mutex::new(HashMap::new()));

#[allow(clippy::disallowed_types)]
fn lock_for(path: &Path) -> Arc<Mutex<()>> {
	let mut map = SCRIPT_LOCKS.lock().unwrap_or_else(PoisonError::into_inner);
	map.entry(path.to_path_buf()).or_insert_with(|| Arc::new(Mutex::new(()))).clone()
}

pub trait Runner {
	fn run(&mut self, command: &Command) -> Result<String, Box<dyn Error>>;

	fn start_script(&mut self) -> Result<(), Box<dyn Error>> {
		Ok(())
	}

	fn end_script(&mut self) -> Result<(), Box<dyn Error>> {
		Ok(())
	}

	fn start_block(&mut self) -> Result<String, Box<dyn Error>> {
		Ok(String::new())
	}

	fn end_block(&mut self) -> Result<String, Box<dyn Error>> {
		Ok(String::new())
	}

	#[allow(unused_variables)]
	fn start_command(&mut self, command: &Command) -> Result<String, Box<dyn Error>> {
		Ok(String::new())
	}

	#[allow(unused_variables)]
	fn end_command(&mut self, command: &Command) -> Result<String, Box<dyn Error>> {
		Ok(String::new())
	}
}

pub fn run_path<R: Runner, P: AsRef<Path>>(runner: &mut R, path: P) -> io::Result<()> {
	let path = path.as_ref();
	let Some(dir) = path.parent() else {
		return Err(io::Error::new(ErrorKind::InvalidInput, format!("invalid path '{path:?}'")));
	};
	let Some(filename) = path.file_name() else {
		return Err(io::Error::new(ErrorKind::InvalidInput, format!("invalid path '{path:?}'")));
	};

	if filename.to_str().unwrap().ends_with(".skip") {
		return Ok(());
	}

	let lock = lock_for(path);
	let _guard = lock.lock().unwrap_or_else(PoisonError::into_inner);

	let input = read_to_string(dir.join(filename))?;
	let output = generate(runner, &input)?;

	Mint::new(dir).new_goldenfile(filename)?.write_all(output.as_bytes())
}

pub fn run<R: Runner, S: Into<String>>(runner: R, test: S) {
	try_run(runner, test).unwrap();
}

pub fn try_run<R: Runner, S: Into<String>>(mut runner: R, test: S) -> io::Result<()> {
	let input = test.into();

	let dir = temp_dir();
	#[allow(clippy::disallowed_methods)]
	let file_name = format!(
		"test-{}-{}.txt",
		process::id(),
		SystemTime::now().duration_since(time::UNIX_EPOCH).unwrap().as_nanos()
	);
	let file_path = dir.join(&file_name);

	let mut file = fs::File::create(&file_path)?;
	file.write_all(input.as_bytes())?;

	let output = generate(&mut runner, &input)?;
	Mint::new(dir).new_goldenfile(&file_name)?.write_all(output.as_bytes())
}

pub fn generate<R: Runner>(runner: &mut R, input: &str) -> io::Result<String> {
	let mut output = String::with_capacity(input.len());
	let eol = detect_eol(input);
	let blocks = parse_blocks(input)?;

	runner.start_script().map_err(|e| io::Error::other(format!("start_script failed: {e}")))?;

	for (i, block) in blocks.iter().enumerate() {
		if block.commands.is_empty() {
			output.push_str(&block.literal);
			continue;
		}
		let block_output = process_block(runner, block, eol)?;
		output.push_str(&format!("{}---{eol}{}", block.literal, block_output));
		if i < blocks.len() - 1 {
			output.push_str(eol);
		}
	}

	runner.end_script().map_err(|e| io::Error::other(format!("end_script failed: {e}")))?;
	Ok(output)
}

#[inline]
fn detect_eol(input: &str) -> &'static str {
	if input.contains("\r\n") {
		"\r\n"
	} else {
		"\n"
	}
}

#[inline]
fn parse_blocks(input: &str) -> io::Result<Vec<Block>> {
	parse(input).map_err(|e| {
		io::Error::new(
			ErrorKind::InvalidInput,
			format!(
				"parse error at line {} column {} for {:?}:\n{}\n{}^",
				e.input.location_line(),
				e.input.get_column(),
				e.code,
				String::from_utf8_lossy(e.input.get_line_beginning()),
				' '.to_string().repeat(e.input.get_utf8_column() - 1)
			),
		)
	})
}

fn process_block<R: Runner>(runner: &mut R, block: &Block, eol: &str) -> io::Result<String> {
	let mut block_output = String::new();
	block_output.push_str(&ensure_eol(
		runner.start_block().map_err(|e| {
			io::Error::other(format!("start_block failed at line {}: {e}", block.line_number))
		})?,
		eol,
	));
	for command in &block.commands {
		let command_output = process_command(runner, command, eol)?;
		block_output.push_str(&command_output);
	}
	block_output.push_str(&ensure_eol(
		runner.end_block().map_err(|e| {
			io::Error::other(format!("end_block failed at line {}: {e}", block.line_number))
		})?,
		eol,
	));
	if block_output.is_empty() {
		block_output.push_str("ok\n");
	}
	Ok(apply_blank_line_prefix(block_output))
}

fn process_command<R: Runner>(runner: &mut R, command: &Command, eol: &str) -> io::Result<String> {
	let mut command_output = String::new();
	command_output.push_str(&ensure_eol(
		runner.start_command(command).map_err(|e| {
			io::Error::other(format!("start_command failed at line {}: {e}", command.line_number))
		})?,
		eol,
	));
	command_output.push_str(&run_command_with_panic_handling(runner, command)?);
	command_output = ensure_eol(command_output, eol);
	command_output.push_str(&ensure_eol(
		runner.end_command(command).map_err(|e| {
			io::Error::other(format!("end_command failed at line {}: {e}", command.line_number))
		})?,
		eol,
	));
	if command.silent {
		command_output.clear();
	}
	if let Some(prefix) = &command.prefix
		&& !command_output.is_empty()
	{
		command_output = format!(
			"{prefix}: {}{eol}",
			command_output
				.strip_suffix(eol)
				.unwrap_or(command_output.as_str())
				.replace('\n', &format!("\n{prefix}: "))
		);
	}
	Ok(command_output)
}

fn run_command_with_panic_handling<R: Runner>(runner: &mut R, command: &Command) -> io::Result<String> {
	let run = AssertUnwindSafe(|| runner.run(command));
	match panic::catch_unwind(run) {
		Ok(Ok(output)) if command.fail => Err(io::Error::other(format!(
			"expected command '{}' to fail at line {}, succeeded with: {output}",
			command.name, command.line_number
		))),
		Ok(Ok(output)) => Ok(output),
		Ok(Err(e)) if command.fail => Ok(format!("{e}")),
		Ok(Err(e)) => Err(io::Error::other(format!(
			"command '{}' failed at line {}: {e}",
			command.name, command.line_number
		))),
		Err(panic) if command.fail => {
			let message = panic
				.downcast_ref::<&str>()
				.map(|s| s.to_string())
				.or_else(|| panic.downcast_ref::<String>().cloned())
				.unwrap_or_else(|| panic::resume_unwind(panic));
			Ok(format!("Panic: {message}"))
		}
		Err(panic) => panic::resume_unwind(panic),
	}
}

#[inline]
fn apply_blank_line_prefix(mut block_output: String) -> String {
	if block_output.starts_with('\n')
		|| block_output.starts_with("\r\n")
		|| block_output.contains("\n\n")
		|| block_output.contains("\n\r\n")
	{
		block_output = format!("> {}", block_output.replace('\n', "\n> "));
		block_output = block_output.replace("> \n", ">\n");
		block_output.pop();
		block_output.pop();
	}
	block_output
}

fn ensure_eol(mut s: String, eol: &str) -> String {
	if let Some(c) = s.chars().next_back()
		&& c != '\n'
	{
		s.push_str(eol)
	}
	s
}

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

	/// A runner which simply counts the number of times its hooks are
	/// called.
	#[derive(Default)]
	struct HookRunner {
		start_script_count: usize,
		end_script_count: usize,
		start_block_count: usize,
		end_block_count: usize,
		start_command_count: usize,
		end_command_count: usize,
	}

	impl Runner for HookRunner {
		fn run(&mut self, _: &Command) -> Result<String, Box<dyn Error>> {
			Ok(String::new())
		}

		fn start_script(&mut self) -> Result<(), Box<dyn Error>> {
			self.start_script_count += 1;
			Ok(())
		}

		fn end_script(&mut self) -> Result<(), Box<dyn Error>> {
			self.end_script_count += 1;
			Ok(())
		}

		fn start_block(&mut self) -> Result<String, Box<dyn Error>> {
			self.start_block_count += 1;
			Ok(String::new())
		}

		fn end_block(&mut self) -> Result<String, Box<dyn Error>> {
			self.end_block_count += 1;
			Ok(String::new())
		}

		fn start_command(&mut self, _: &Command) -> Result<String, Box<dyn Error>> {
			self.start_command_count += 1;
			Ok(String::new())
		}

		fn end_command(&mut self, _: &Command) -> Result<String, Box<dyn Error>> {
			self.end_command_count += 1;
			Ok(String::new())
		}
	}

	/// Tests that runner hooks are called as expected.
	#[test]
	fn hooks() {
		let mut runner = HookRunner::default();
		generate(
			&mut runner,
			r#"
command
---

command
command
---
"#,
		)
		.unwrap();

		assert_eq!(runner.start_script_count, 1);
		assert_eq!(runner.end_script_count, 1);
		assert_eq!(runner.start_block_count, 2);
		assert_eq!(runner.end_block_count, 2);
		assert_eq!(runner.start_command_count, 3);
		assert_eq!(runner.end_command_count, 3);
	}
}