use std::collections::BTreeMap;
use std::fmt;
use std::sync::Arc;
use rust_decimal::Decimal;
use serde_json::{Value, json, to_string};
use ironflow_core::error::{AgentError, OperationError};
use ironflow_core::provider::{AgentConfig, AgentOutput, AgentProvider, InvokeFuture};
use crate::config::{ApprovalConfig, HttpConfig, HumanInputConfig, ShellConfig, StepConfig};
use crate::error::EngineError;
use crate::executor::{
ApprovalOutcome, HumanInputOutcome, StepArtifacts, StepInterceptor, StepOutput,
};
const MISSING_AGENT_PROVIDER: &str = "TestEngine has no agent provider: call with_mock_agent(...), with_recorded_agent(...) or \
with_agent_provider(...)";
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct MockShellOutput {
pub stdout: String,
pub stderr: String,
pub exit_code: i32,
}
impl MockShellOutput {
pub fn ok(stdout: &str) -> Self {
Self {
stdout: stdout.to_string(),
..Self::default()
}
}
pub fn failed(exit_code: i32, stderr: &str) -> Self {
Self {
stdout: String::new(),
stderr: stderr.to_string(),
exit_code,
}
}
pub(crate) fn into_step_result(self) -> Result<StepOutput, EngineError> {
if self.exit_code != 0 {
return Err(EngineError::Operation(OperationError::Shell {
exit_code: self.exit_code,
stderr: self.stderr,
}));
}
Ok(StepOutput {
output: json!({
"stdout": self.stdout,
"stderr": self.stderr,
"exit_code": self.exit_code,
}),
duration_ms: 0,
cost_usd: Decimal::ZERO,
input_tokens: None,
cache_read_input_tokens: None,
cache_creation_input_tokens: None,
output_tokens: None,
model: None,
debug_messages: None,
artifacts: StepArtifacts::default(),
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MockHttpResponse {
pub status: u16,
pub headers: Vec<(String, String)>,
pub body: String,
}
impl Default for MockHttpResponse {
fn default() -> Self {
Self {
status: 200,
headers: Vec::new(),
body: String::new(),
}
}
}
impl MockHttpResponse {
pub fn ok(body: &Value) -> Self {
Self::json(200, body)
}
pub fn json(status: u16, body: &Value) -> Self {
Self {
status,
headers: Vec::new(),
body: to_string(body).unwrap_or_else(|_| body.to_string()),
}
}
pub fn text(status: u16, body: &str) -> Self {
Self {
status,
headers: Vec::new(),
body: body.to_string(),
}
}
pub fn header(mut self, name: &str, value: &str) -> Self {
self.headers.push((name.to_string(), value.to_string()));
self
}
pub(crate) fn into_step_output(self) -> StepOutput {
let headers: BTreeMap<String, String> = self.headers.into_iter().collect();
StepOutput {
output: json!({
"status": self.status,
"headers": headers,
"body": self.body,
}),
duration_ms: 0,
cost_usd: Decimal::ZERO,
input_tokens: None,
cache_read_input_tokens: None,
cache_creation_input_tokens: None,
output_tokens: None,
model: None,
debug_messages: None,
artifacts: StepArtifacts::default(),
}
}
}
pub type ShellMock =
Arc<dyn Fn(&ShellConfig) -> Result<MockShellOutput, OperationError> + Send + Sync>;
pub type HttpMock =
Arc<dyn Fn(&HttpConfig) -> Result<MockHttpResponse, OperationError> + Send + Sync>;
pub type HumanInputMock = Arc<dyn Fn(&str, &HumanInputConfig) -> HumanInputOutcome + Send + Sync>;
pub type AgentMock = Arc<dyn Fn(&AgentConfig) -> Result<AgentOutput, AgentError> + Send + Sync>;
#[derive(Clone, Default)]
pub struct MockInterceptor {
shell: Option<ShellMock>,
http: Option<HttpMock>,
approval: Option<ApprovalOutcome>,
human_input: Option<HumanInputMock>,
}
impl MockInterceptor {
pub fn new() -> Self {
Self::default()
}
pub fn shell(
mut self,
f: impl Fn(&ShellConfig) -> Result<MockShellOutput, OperationError> + Send + Sync + 'static,
) -> Self {
self.shell = Some(Arc::new(f));
self
}
pub fn http(
mut self,
f: impl Fn(&HttpConfig) -> Result<MockHttpResponse, OperationError> + Send + Sync + 'static,
) -> Self {
self.http = Some(Arc::new(f));
self
}
pub fn approval(mut self, outcome: ApprovalOutcome) -> Self {
self.approval = Some(outcome);
self
}
pub fn human_input(
mut self,
f: impl Fn(&str, &HumanInputConfig) -> HumanInputOutcome + Send + Sync + 'static,
) -> Self {
self.human_input = Some(Arc::new(f));
self
}
}
impl fmt::Debug for MockInterceptor {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("MockInterceptor")
.field("shell", &self.shell.is_some())
.field("http", &self.http.is_some())
.field("approval", &self.approval)
.field("human_input", &self.human_input.is_some())
.finish()
}
}
impl StepInterceptor for MockInterceptor {
fn intercept(&self, config: &StepConfig) -> Option<Result<StepOutput, EngineError>> {
match config {
StepConfig::Shell(cfg) => {
let mock = self.shell.as_ref()?;
Some(match mock(cfg) {
Ok(out) => out.into_step_result(),
Err(err) => Err(EngineError::Operation(err)),
})
}
StepConfig::Http(cfg) => {
let mock = self.http.as_ref()?;
Some(match mock(cfg) {
Ok(res) => Ok(res.into_step_output()),
Err(err) => Err(EngineError::Operation(err)),
})
}
_ => None,
}
}
fn intercept_approval(&self, _name: &str, _config: &ApprovalConfig) -> Option<ApprovalOutcome> {
self.approval.clone()
}
fn intercept_human_input(
&self,
name: &str,
config: &HumanInputConfig,
_schema: &Value,
) -> Option<HumanInputOutcome> {
self.human_input.as_ref().map(|f| f(name, config))
}
}
pub struct MockAgentProvider {
f: AgentMock,
}
impl MockAgentProvider {
pub fn new(
f: impl Fn(&AgentConfig) -> Result<AgentOutput, AgentError> + Send + Sync + 'static,
) -> Self {
Self { f: Arc::new(f) }
}
}
impl fmt::Debug for MockAgentProvider {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("MockAgentProvider").finish_non_exhaustive()
}
}
impl AgentProvider for MockAgentProvider {
fn invoke<'a>(&'a self, config: &'a AgentConfig) -> InvokeFuture<'a> {
let result = (self.f)(config);
Box::pin(async move { result })
}
}
#[derive(Debug, Default, Clone, Copy)]
pub struct MissingAgentProvider;
impl AgentProvider for MissingAgentProvider {
fn invoke<'a>(&'a self, _config: &'a AgentConfig) -> InvokeFuture<'a> {
Box::pin(async move {
Err(AgentError::ProcessFailed {
exit_code: -1,
stderr: MISSING_AGENT_PROVIDER.to_string(),
})
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::AgentStepConfig;
#[test]
fn shell_ok_maps_to_the_real_executor_output_shape() {
let output = MockShellOutput::ok("hello\n")
.into_step_result()
.expect("exit code 0 succeeds");
assert_eq!(output.output["stdout"], "hello\n");
assert_eq!(output.output["stderr"], "");
assert_eq!(output.output["exit_code"], 0);
assert_eq!(output.cost_usd, Decimal::ZERO);
}
#[test]
fn shell_default_is_an_empty_success() {
let default = MockShellOutput::default();
assert_eq!(default.exit_code, 0);
assert!(default.stdout.is_empty());
assert!(default.stderr.is_empty());
}
#[test]
fn shell_non_zero_exit_is_an_operation_error() {
let err = MockShellOutput::failed(2, "x")
.into_step_result()
.expect_err("a non-zero exit code fails the step");
match err {
EngineError::Operation(OperationError::Shell { exit_code, stderr }) => {
assert_eq!(exit_code, 2);
assert_eq!(stderr, "x");
}
other => panic!("expected a shell operation error, got {other}"),
}
}
#[test]
fn http_json_carries_status_body_and_headers() {
let output = MockHttpResponse::json(201, &json!({"id": 7}))
.header("location", "/things/7")
.into_step_output();
assert_eq!(output.output["status"], 201);
assert_eq!(output.output["body"], r#"{"id":7}"#);
assert_eq!(output.output["headers"]["location"], "/things/7");
}
#[test]
fn http_default_is_an_empty_200() {
let default = MockHttpResponse::default();
assert_eq!(default.status, 200);
assert!(default.body.is_empty());
assert!(default.headers.is_empty());
}
#[test]
fn http_non_2xx_is_still_an_output() {
let output = MockHttpResponse::text(500, "boom").into_step_output();
assert_eq!(output.status(), Some(500));
assert_eq!(output.body(), "boom");
}
#[test]
fn intercept_declines_agent_steps() {
let interceptor = MockInterceptor::new().shell(|_| Ok(MockShellOutput::ok("x")));
let config = StepConfig::Agent(AgentStepConfig::new("review this"));
assert!(interceptor.intercept(&config).is_none());
}
#[test]
fn intercept_declines_shell_steps_without_a_shell_mock() {
let interceptor = MockInterceptor::new();
let config = StepConfig::Shell(ShellConfig::new("echo hi"));
assert!(interceptor.intercept(&config).is_none());
}
#[test]
fn intercept_approval_returns_the_configured_outcome() {
let interceptor = MockInterceptor::new().approval(ApprovalOutcome::reject("nope"));
let config = ApprovalConfig::new("Approve?");
assert_eq!(
interceptor.intercept_approval("gate", &config),
Some(ApprovalOutcome::reject("nope"))
);
assert_eq!(
MockInterceptor::new().intercept_approval("gate", &config),
None
);
}
#[test]
fn intercept_human_input_returns_the_mocked_answer() {
let interceptor = MockInterceptor::new().human_input(|name, cfg| {
HumanInputOutcome::Provided(json!({"step": name, "message": cfg.message()}))
});
let config = HumanInputConfig::new("Answer?");
let answer = json!({"step": "clarify", "message": "Answer?"});
assert_eq!(
interceptor.intercept_human_input("clarify", &config, &json!({})),
Some(HumanInputOutcome::Provided(answer))
);
}
#[test]
fn intercept_human_input_declines_without_a_mock() {
let config = HumanInputConfig::new("Answer?");
assert_eq!(
MockInterceptor::new().intercept_human_input("clarify", &config, &json!({})),
None
);
}
#[test]
fn debug_reports_which_seams_are_mocked() {
let interceptor = MockInterceptor::new().http(|_| Ok(MockHttpResponse::default()));
let rendered = format!("{interceptor:?}");
assert!(rendered.contains("shell: false"));
assert!(rendered.contains("http: true"));
}
#[tokio::test]
async fn missing_agent_provider_names_the_three_constructors() {
let config = AgentConfig::new("anything");
let err = MissingAgentProvider
.invoke(&config)
.await
.expect_err("no agent backend is configured");
let message = err.to_string();
assert!(message.contains("with_mock_agent"));
assert!(message.contains("with_recorded_agent"));
assert!(message.contains("with_agent_provider"));
}
#[tokio::test]
async fn mock_agent_provider_runs_the_closure() {
let provider = MockAgentProvider::new(|cfg| {
let echoed = json!({"echoed": cfg.prompt.clone()});
Ok(AgentOutput::new(echoed))
});
let config = AgentConfig::new("say hi");
let output = provider.invoke(&config).await.expect("the mock succeeded");
assert_eq!(output.value["echoed"], "say hi");
}
}