use std::collections::VecDeque;
use std::io::{self, Read, Write};
use std::sync::{Arc, Condvar, Mutex};
use std::time::Duration;
use super::handle::{
ExitStatus, ProcessError, ProcessHandle, ProcessKiller, ProcessSpawner, SpawnSpec,
};
#[derive(Clone, Debug)]
pub struct MockProcessConfig {
pub pid: u32,
pub pgid: Option<u32>,
pub stdout: Vec<u8>,
pub stderr: Vec<u8>,
pub exit_status: Option<ExitStatus>,
pub force_timeout: bool,
pub spawn_error: Option<ProcessError>,
}
impl Default for MockProcessConfig {
fn default() -> Self {
Self {
pid: 99_999,
pgid: Some(99_999),
stdout: Vec::new(),
stderr: Vec::new(),
exit_status: Some(ExitStatus::from_code(0)),
force_timeout: false,
spawn_error: None,
}
}
}
impl MockProcessConfig {
pub fn completed(exit_code: i32) -> Self {
Self {
exit_status: Some(ExitStatus::from_code(exit_code)),
..Self::default()
}
}
pub fn with_stdout(exit_code: i32, stdout: impl Into<Vec<u8>>) -> Self {
Self {
stdout: stdout.into(),
exit_status: Some(ExitStatus::from_code(exit_code)),
..Self::default()
}
}
pub fn running() -> Self {
Self {
exit_status: None,
..Self::default()
}
}
}
#[derive(Default)]
struct MockSpawnerInner {
queue: VecDeque<(MockProcessConfig, Arc<MockState>)>,
captured: Vec<SpawnSpec>,
last_controller: Option<MockHandleController>,
}
pub struct MockSpawner {
inner: Mutex<MockSpawnerInner>,
}
impl Default for MockSpawner {
fn default() -> Self {
Self::new()
}
}
impl MockSpawner {
pub fn new() -> Self {
Self {
inner: Mutex::new(MockSpawnerInner::default()),
}
}
pub fn enqueue(&self, config: MockProcessConfig) -> MockHandleController {
let state = Arc::new(MockState::new(&config));
let controller = MockHandleController {
state: Arc::clone(&state),
};
let mut inner = self.inner.lock().expect("MockSpawner mutex poisoned");
inner.queue.push_back((config, state));
inner.last_controller = Some(controller.clone());
controller
}
pub fn captured(&self) -> Vec<SpawnSpec> {
self.inner
.lock()
.expect("MockSpawner mutex poisoned")
.captured
.clone()
}
pub fn last_controller(&self) -> Option<MockHandleController> {
self.inner
.lock()
.expect("MockSpawner mutex poisoned")
.last_controller
.clone()
}
}
impl ProcessSpawner for MockSpawner {
fn spawn(&self, spec: SpawnSpec) -> Result<Box<dyn ProcessHandle>, ProcessError> {
let (config, state) = {
let mut inner = self.inner.lock().expect("MockSpawner mutex poisoned");
inner.captured.push(spec);
inner.queue.pop_front().expect(
"MockSpawner: spawn() called with no enqueued configuration. Call \
MockSpawner::enqueue(...) before each expected spawn.",
)
};
if let Some(err) = config.spawn_error {
return Err(err);
}
let killer: Arc<dyn ProcessKiller> = Arc::new(MockKiller {
state: Arc::clone(&state),
});
Ok(Box::new(MockProcess {
pid: config.pid,
pgid: config.pgid,
killer,
state,
stdin_taken: false,
stdout_taken: false,
stderr_taken: false,
}))
}
}
#[derive(Clone)]
pub struct MockHandleController {
state: Arc<MockState>,
}
impl MockHandleController {
pub fn append_stdout(&self, bytes: &[u8]) {
let mut data = self.state.stdout.lock().unwrap();
data.extend_from_slice(bytes);
self.state.stdout_cv.notify_all();
}
pub fn append_stderr(&self, bytes: &[u8]) {
let mut data = self.state.stderr.lock().unwrap();
data.extend_from_slice(bytes);
self.state.stderr_cv.notify_all();
}
pub fn complete_with(&self, status: ExitStatus) {
let mut exit = self.state.exit.lock().unwrap();
if exit.is_none() {
*exit = Some(ExitOutcome {
status,
killed: false,
});
}
drop(exit);
self.state.exit_cv.notify_all();
self.state.stdout_cv.notify_all();
self.state.stderr_cv.notify_all();
}
pub fn was_killed(&self) -> bool {
self.state
.exit
.lock()
.unwrap()
.as_ref()
.map(|o| o.killed)
.unwrap_or(false)
}
pub fn stdin_written(&self) -> Vec<u8> {
self.state.stdin_written.lock().unwrap().clone()
}
}
struct MockState {
stdout: Mutex<Vec<u8>>,
stderr: Mutex<Vec<u8>>,
stdin_written: Mutex<Vec<u8>>,
exit: Mutex<Option<ExitOutcome>>,
exit_cv: Condvar,
stdout_cv: Condvar,
stderr_cv: Condvar,
force_timeout: bool,
}
#[derive(Clone, Copy, Debug)]
struct ExitOutcome {
status: ExitStatus,
killed: bool,
}
impl MockState {
fn new(config: &MockProcessConfig) -> Self {
let exit = config.exit_status.map(|status| ExitOutcome {
status,
killed: false,
});
Self {
stdout: Mutex::new(config.stdout.clone()),
stderr: Mutex::new(config.stderr.clone()),
stdin_written: Mutex::new(Vec::new()),
exit: Mutex::new(exit),
exit_cv: Condvar::new(),
stdout_cv: Condvar::new(),
stderr_cv: Condvar::new(),
force_timeout: config.force_timeout,
}
}
fn is_exited(&self) -> bool {
self.exit.lock().unwrap().is_some()
}
fn wait_for_exit(&self, timeout: Option<Duration>) -> Option<ExitOutcome> {
let mut exit = self.exit.lock().unwrap();
if let Some(timeout) = timeout {
if exit.is_none() {
let (next, result) = self.exit_cv.wait_timeout(exit, timeout).unwrap();
exit = next;
if result.timed_out() && exit.is_none() {
return None;
}
}
} else {
while exit.is_none() {
exit = self.exit_cv.wait(exit).unwrap();
}
}
*exit
}
fn record_kill(&self) {
let mut exit = self.exit.lock().unwrap();
if exit.is_none() {
*exit = Some(ExitOutcome {
status: ExitStatus::from_signal(9),
killed: true,
});
} else if let Some(outcome) = exit.as_mut() {
outcome.killed = true;
}
drop(exit);
self.exit_cv.notify_all();
self.stdout_cv.notify_all();
self.stderr_cv.notify_all();
}
}
pub struct MockProcess {
pid: u32,
pgid: Option<u32>,
killer: Arc<dyn ProcessKiller>,
state: Arc<MockState>,
stdin_taken: bool,
stdout_taken: bool,
stderr_taken: bool,
}
impl ProcessHandle for MockProcess {
fn pid(&self) -> Option<u32> {
Some(self.pid)
}
fn process_group_id(&self) -> Option<u32> {
self.pgid
}
fn killer(&self) -> Arc<dyn ProcessKiller> {
Arc::clone(&self.killer)
}
fn take_stdin(&mut self) -> Option<Box<dyn Write + Send>> {
if self.stdin_taken {
return None;
}
self.stdin_taken = true;
Some(Box::new(MockStdin {
state: Arc::clone(&self.state),
}))
}
fn take_stdout(&mut self) -> Option<Box<dyn Read + Send>> {
if self.stdout_taken {
return None;
}
self.stdout_taken = true;
Some(Box::new(MockStdoutReader {
state: Arc::clone(&self.state),
kind: PipeKind::Stdout,
}))
}
fn take_stderr(&mut self) -> Option<Box<dyn Read + Send>> {
if self.stderr_taken {
return None;
}
self.stderr_taken = true;
Some(Box::new(MockStdoutReader {
state: Arc::clone(&self.state),
kind: PipeKind::Stderr,
}))
}
fn wait_with_timeout(
&mut self,
timeout: Option<Duration>,
) -> io::Result<(Option<ExitStatus>, bool)> {
if self.state.force_timeout {
self.state.record_kill();
return Ok((None, true));
}
let Some(timeout) = timeout else {
let outcome = self
.state
.wait_for_exit(None)
.expect("wait without timeout returned None");
return Ok((Some(outcome.status), false));
};
match self.state.wait_for_exit(Some(timeout)) {
Some(outcome) => Ok((Some(outcome.status), false)),
None => {
self.state.record_kill();
Ok((None, true))
}
}
}
fn wait(&mut self) -> io::Result<ExitStatus> {
let outcome = self
.state
.wait_for_exit(None)
.expect("wait without timeout returned None");
Ok(outcome.status)
}
}
struct MockStdin {
state: Arc<MockState>,
}
impl Write for MockStdin {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.state
.stdin_written
.lock()
.unwrap()
.extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
#[derive(Clone, Copy)]
enum PipeKind {
Stdout,
Stderr,
}
struct MockStdoutReader {
state: Arc<MockState>,
kind: PipeKind,
}
impl MockStdoutReader {
fn pipe_lock(&self) -> &Mutex<Vec<u8>> {
match self.kind {
PipeKind::Stdout => &self.state.stdout,
PipeKind::Stderr => &self.state.stderr,
}
}
fn pipe_cv(&self) -> &Condvar {
match self.kind {
PipeKind::Stdout => &self.state.stdout_cv,
PipeKind::Stderr => &self.state.stderr_cv,
}
}
}
impl Read for MockStdoutReader {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
let lock = self.pipe_lock();
let cv = self.pipe_cv();
let mut data = lock.lock().unwrap();
loop {
if !data.is_empty() {
let n = data.len().min(buf.len());
buf[..n].copy_from_slice(&data[..n]);
data.drain(..n);
return Ok(n);
}
if self.state.is_exited() {
return Ok(0);
}
data = cv.wait(data).unwrap();
}
}
}
struct MockKiller {
state: Arc<MockState>,
}
impl ProcessKiller for MockKiller {
fn kill(&self) {
self.state.record_kill();
}
}