use crate::c::{ForkReturn, PipeFds, RealSystemFunctions, SystemFunctions};
use libc::c_int;
use std::cell::{Cell, RefCell};
use std::collections::VecDeque;
use std::io;
use std::thread_local;
#[derive(Clone, Debug)]
enum MockResult<T> {
Ok(T),
Err(i32), }
impl<T> MockResult<T> {
#[allow(clippy::wrong_self_convention)]
fn to_result(self) -> Result<T, io::Error> {
match self {
MockResult::Ok(val) => Ok(val),
MockResult::Err(code) => Err(io::Error::from_raw_os_error(code)),
}
}
fn from_result(result: Result<T, io::Error>) -> Self {
match result {
Ok(val) => MockResult::Ok(val),
Err(err) => {
if let Some(code) = err.raw_os_error() {
MockResult::Err(code)
} else {
MockResult::Err(libc::EIO) }
}
}
}
}
#[derive(Clone, Debug)]
enum CallImplementation<T> {
Real, Mock(MockResult<T>), }
pub enum CallBehavior<T> {
Real, Mock(Result<T, io::Error>), }
#[derive(Clone, Debug)]
struct CallQueue<T> {
queue: RefCell<VecDeque<CallImplementation<T>>>,
name: &'static str, }
impl<T: Clone> CallQueue<T> {
fn new(name: &'static str) -> Self {
Self {
queue: RefCell::new(VecDeque::new()),
name,
}
}
fn push(&self, behavior: CallBehavior<T>) {
let mut queue = self.queue.borrow_mut();
match behavior {
CallBehavior::Real => queue.push_back(CallImplementation::Real),
CallBehavior::Mock(result) => {
queue.push_back(CallImplementation::Mock(MockResult::from_result(result)));
}
}
}
fn execute_next<F>(&self, real_impl: F, fallback_enabled: bool) -> Result<T, io::Error>
where
F: FnOnce() -> Result<T, io::Error>,
{
let mut queue = self.queue.borrow_mut();
match queue.pop_front() {
Some(CallImplementation::Real) => real_impl(),
Some(CallImplementation::Mock(result)) => result.to_result(),
None if fallback_enabled => real_impl(),
None => panic!(
"No mock behavior configured for {}() and fallback is disabled",
self.name
),
}
}
}
#[derive(Clone)]
pub struct MockableSystemFunctions {
real_impl: RealSystemFunctions,
fallback_enabled: Cell<bool>,
fork_queue: CallQueue<ForkReturn>,
pipe_queue: CallQueue<PipeFds>,
close_queue: CallQueue<()>,
waitpid_queue: CallQueue<c_int>,
}
impl std::fmt::Debug for MockableSystemFunctions {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "MockableSystemFunctions")
}
}
impl Default for MockableSystemFunctions {
fn default() -> Self {
Self {
real_impl: RealSystemFunctions,
fallback_enabled: Cell::new(true), fork_queue: CallQueue::new("fork"),
pipe_queue: CallQueue::new("pipe"),
close_queue: CallQueue::new("close"),
waitpid_queue: CallQueue::new("waitpid"),
}
}
}
impl MockableSystemFunctions {
pub fn with_fallback() -> Self {
let mock = Self::default();
mock.enable_fallback();
mock
}
pub fn strict() -> Self {
let mock = Self::default();
mock.disable_fallback();
mock
}
pub fn enable_fallback(&self) -> &Self {
self.fallback_enabled.set(true);
self
}
pub fn disable_fallback(&self) -> &Self {
self.fallback_enabled.set(false);
self
}
pub fn is_fallback_enabled(&self) -> bool {
self.fallback_enabled.get()
}
pub fn expect_fork(&self, behavior: CallBehavior<ForkReturn>) -> &Self {
self.fork_queue.push(behavior);
self
}
pub fn expect_pipe(&self, behavior: CallBehavior<PipeFds>) -> &Self {
self.pipe_queue.push(behavior);
self
}
pub fn expect_close(&self, behavior: CallBehavior<()>) -> &Self {
self.close_queue.push(behavior);
self
}
pub fn expect_waitpid(&self, behavior: CallBehavior<c_int>) -> &Self {
self.waitpid_queue.push(behavior);
self
}
}
impl SystemFunctions for MockableSystemFunctions {
fn fork(&self) -> Result<ForkReturn, io::Error> {
self.fork_queue
.execute_next(|| self.real_impl.fork(), self.is_fallback_enabled())
}
fn pipe(&self) -> Result<PipeFds, io::Error> {
self.pipe_queue
.execute_next(|| self.real_impl.pipe(), self.is_fallback_enabled())
}
fn close(&self, fd: c_int) -> Result<(), io::Error> {
self.close_queue
.execute_next(|| self.real_impl.close(fd), self.is_fallback_enabled())
}
fn waitpid(&self, pid: c_int) -> Result<c_int, io::Error> {
self.waitpid_queue
.execute_next(|| self.real_impl.waitpid(pid), self.is_fallback_enabled())
}
fn _exit(&self, status: c_int) -> ! {
if is_mocking_enabled() {
panic!("_exit({status}) called in mock context");
} else {
#[allow(clippy::used_underscore_items)]
self.real_impl._exit(status)
}
}
}
thread_local! {
static CURRENT_MOCK: RefCell<Option<MockableSystemFunctions>> = const { RefCell::new(None) };
static IS_MOCKING_ENABLED: Cell<bool> = const { Cell::new(false) };
}
pub fn enable_mocking(mock: &MockableSystemFunctions) {
CURRENT_MOCK.with(|m| {
*m.borrow_mut() = Some(mock.clone());
});
IS_MOCKING_ENABLED.with(|e| e.set(true));
}
pub fn disable_mocking() {
CURRENT_MOCK.with(|m| {
*m.borrow_mut() = None;
});
IS_MOCKING_ENABLED.with(|e| e.set(false));
}
pub fn is_mocking_enabled() -> bool {
IS_MOCKING_ENABLED.with(std::cell::Cell::get)
}
pub fn get_current_mock() -> MockableSystemFunctions {
CURRENT_MOCK.with(|m| {
m.borrow()
.clone()
.expect("No mock available in thread-local storage")
})
}
pub enum MockConfig {
Fallback,
Strict,
ConfiguredWithFallback(Box<dyn FnOnce(&MockableSystemFunctions)>),
ConfiguredStrict(Box<dyn FnOnce(&MockableSystemFunctions)>),
}
pub fn with_mock_system<R>(
config: MockConfig,
test_fn: impl FnOnce(&MockableSystemFunctions) -> R,
) -> R {
let mock = MockableSystemFunctions::with_fallback();
match config {
MockConfig::Fallback => {}
MockConfig::Strict => {
mock.disable_fallback();
}
MockConfig::ConfiguredWithFallback(configure_fn) => {
configure_fn(&mock);
}
MockConfig::ConfiguredStrict(configure_fn) => {
mock.disable_fallback();
configure_fn(&mock);
}
}
enable_mocking(&mock);
let result = test_fn(&mock);
disable_mocking();
result
}
pub fn configured_with_fallback<F>(configure_fn: F) -> MockConfig
where
F: FnOnce(&MockableSystemFunctions) + 'static,
{
MockConfig::ConfiguredWithFallback(Box::new(configure_fn))
}
pub fn configured_strict<F>(configure_fn: F) -> MockConfig
where
F: FnOnce(&MockableSystemFunctions) + 'static,
{
MockConfig::ConfiguredStrict(Box::new(configure_fn))
}