1use std::collections::BTreeMap;
14use std::fmt;
15use std::sync::Arc;
16
17use rust_decimal::Decimal;
18use serde_json::{Value, json, to_string};
19
20use ironflow_core::error::{AgentError, OperationError};
21use ironflow_core::provider::{AgentConfig, AgentOutput, AgentProvider, InvokeFuture};
22
23use crate::config::{ApprovalConfig, HttpConfig, HumanInputConfig, ShellConfig, StepConfig};
24use crate::error::EngineError;
25use crate::executor::{
26 ApprovalOutcome, HumanInputOutcome, StepArtifacts, StepInterceptor, StepOutput,
27};
28
29const MISSING_AGENT_PROVIDER: &str = "TestEngine has no agent provider: call with_mock_agent(...), with_recorded_agent(...) or \
31 with_agent_provider(...)";
32
33#[derive(Debug, Clone, PartialEq, Eq, Default)]
47pub struct MockShellOutput {
48 pub stdout: String,
50 pub stderr: String,
52 pub exit_code: i32,
54}
55
56impl MockShellOutput {
57 pub fn ok(stdout: &str) -> Self {
69 Self {
70 stdout: stdout.to_string(),
71 ..Self::default()
72 }
73 }
74
75 pub fn failed(exit_code: i32, stderr: &str) -> Self {
86 Self {
87 stdout: String::new(),
88 stderr: stderr.to_string(),
89 exit_code,
90 }
91 }
92
93 pub(crate) fn into_step_result(
100 self,
101 exit_code_as_output: bool,
102 ) -> Result<StepOutput, EngineError> {
103 if self.exit_code != 0 && !exit_code_as_output {
104 return Err(EngineError::Operation(OperationError::Shell {
105 exit_code: self.exit_code,
106 stderr: self.stderr,
107 }));
108 }
109
110 Ok(StepOutput {
111 output: json!({
112 "stdout": self.stdout,
113 "stderr": self.stderr,
114 "exit_code": self.exit_code,
115 }),
116 duration_ms: 0,
117 cost_usd: Decimal::ZERO,
118 input_tokens: None,
119 cache_read_input_tokens: None,
120 cache_creation_input_tokens: None,
121 output_tokens: None,
122 model: None,
123 debug_messages: None,
124 artifacts: StepArtifacts::default(),
125 })
126 }
127}
128
129#[derive(Debug, Clone, PartialEq, Eq)]
148pub struct MockHttpResponse {
149 pub status: u16,
151 pub headers: Vec<(String, String)>,
153 pub body: String,
155}
156
157impl Default for MockHttpResponse {
158 fn default() -> Self {
159 Self {
160 status: 200,
161 headers: Vec::new(),
162 body: String::new(),
163 }
164 }
165}
166
167impl MockHttpResponse {
168 pub fn ok(body: &Value) -> Self {
181 Self::json(200, body)
182 }
183
184 pub fn json(status: u16, body: &Value) -> Self {
196 Self {
197 status,
198 headers: Vec::new(),
199 body: to_string(body).unwrap_or_else(|_| body.to_string()),
201 }
202 }
203
204 pub fn text(status: u16, body: &str) -> Self {
215 Self {
216 status,
217 headers: Vec::new(),
218 body: body.to_string(),
219 }
220 }
221
222 pub fn header(mut self, name: &str, value: &str) -> Self {
233 self.headers.push((name.to_string(), value.to_string()));
234 self
235 }
236
237 pub(crate) fn into_step_output(self) -> StepOutput {
240 let headers: BTreeMap<String, String> = self.headers.into_iter().collect();
241 StepOutput {
242 output: json!({
243 "status": self.status,
244 "headers": headers,
245 "body": self.body,
246 }),
247 duration_ms: 0,
248 cost_usd: Decimal::ZERO,
249 input_tokens: None,
250 cache_read_input_tokens: None,
251 cache_creation_input_tokens: None,
252 output_tokens: None,
253 model: None,
254 debug_messages: None,
255 artifacts: StepArtifacts::default(),
256 }
257 }
258}
259
260pub type ShellMock =
262 Arc<dyn Fn(&ShellConfig) -> Result<MockShellOutput, OperationError> + Send + Sync>;
263
264pub type HttpMock =
266 Arc<dyn Fn(&HttpConfig) -> Result<MockHttpResponse, OperationError> + Send + Sync>;
267
268pub type HumanInputMock = Arc<dyn Fn(&str, &HumanInputConfig) -> HumanInputOutcome + Send + Sync>;
270
271pub type AgentMock = Arc<dyn Fn(&AgentConfig) -> Result<AgentOutput, AgentError> + Send + Sync>;
273
274#[derive(Clone, Default)]
296pub struct MockInterceptor {
297 shell: Option<ShellMock>,
298 http: Option<HttpMock>,
299 approval: Option<ApprovalOutcome>,
300 human_input: Option<HumanInputMock>,
301}
302
303impl MockInterceptor {
304 pub fn new() -> Self {
318 Self::default()
319 }
320
321 pub fn shell(
333 mut self,
334 f: impl Fn(&ShellConfig) -> Result<MockShellOutput, OperationError> + Send + Sync + 'static,
335 ) -> Self {
336 self.shell = Some(Arc::new(f));
337 self
338 }
339
340 pub fn http(
353 mut self,
354 f: impl Fn(&HttpConfig) -> Result<MockHttpResponse, OperationError> + Send + Sync + 'static,
355 ) -> Self {
356 self.http = Some(Arc::new(f));
357 self
358 }
359
360 pub fn approval(mut self, outcome: ApprovalOutcome) -> Self {
371 self.approval = Some(outcome);
372 self
373 }
374
375 pub fn human_input(
388 mut self,
389 f: impl Fn(&str, &HumanInputConfig) -> HumanInputOutcome + Send + Sync + 'static,
390 ) -> Self {
391 self.human_input = Some(Arc::new(f));
392 self
393 }
394}
395
396impl fmt::Debug for MockInterceptor {
397 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
398 f.debug_struct("MockInterceptor")
400 .field("shell", &self.shell.is_some())
401 .field("http", &self.http.is_some())
402 .field("approval", &self.approval)
403 .field("human_input", &self.human_input.is_some())
404 .finish()
405 }
406}
407
408impl StepInterceptor for MockInterceptor {
409 fn intercept(&self, config: &StepConfig) -> Option<Result<StepOutput, EngineError>> {
410 match config {
411 StepConfig::Shell(cfg) => {
412 let mock = self.shell.as_ref()?;
413 Some(match mock(cfg) {
414 Ok(out) => out.into_step_result(cfg.exit_code_as_output),
415 Err(err) => Err(EngineError::Operation(err)),
416 })
417 }
418 StepConfig::Http(cfg) => {
419 let mock = self.http.as_ref()?;
420 Some(match mock(cfg) {
421 Ok(res) => Ok(res.into_step_output()),
422 Err(err) => Err(EngineError::Operation(err)),
423 })
424 }
425 _ => None,
427 }
428 }
429
430 fn intercept_approval(&self, _name: &str, _config: &ApprovalConfig) -> Option<ApprovalOutcome> {
431 self.approval.clone()
432 }
433
434 fn intercept_human_input(
435 &self,
436 name: &str,
437 config: &HumanInputConfig,
438 _schema: &Value,
439 ) -> Option<HumanInputOutcome> {
440 self.human_input.as_ref().map(|f| f(name, config))
441 }
442}
443
444pub struct MockAgentProvider {
460 f: AgentMock,
461}
462
463impl MockAgentProvider {
464 pub fn new(
477 f: impl Fn(&AgentConfig) -> Result<AgentOutput, AgentError> + Send + Sync + 'static,
478 ) -> Self {
479 Self { f: Arc::new(f) }
480 }
481}
482
483impl fmt::Debug for MockAgentProvider {
484 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
485 f.debug_struct("MockAgentProvider").finish_non_exhaustive()
486 }
487}
488
489impl AgentProvider for MockAgentProvider {
490 fn invoke<'a>(&'a self, config: &'a AgentConfig) -> InvokeFuture<'a> {
491 let result = (self.f)(config);
492 Box::pin(async move { result })
493 }
494}
495
496#[derive(Debug, Default, Clone, Copy)]
514pub struct MissingAgentProvider;
515
516impl AgentProvider for MissingAgentProvider {
517 fn invoke<'a>(&'a self, _config: &'a AgentConfig) -> InvokeFuture<'a> {
518 Box::pin(async move {
519 Err(AgentError::ProcessFailed {
520 exit_code: -1,
521 stderr: MISSING_AGENT_PROVIDER.to_string(),
522 })
523 })
524 }
525}
526
527#[cfg(test)]
528mod tests {
529 use super::*;
530
531 use crate::config::AgentStepConfig;
532
533 #[test]
534 fn shell_ok_maps_to_the_real_executor_output_shape() {
535 let output = MockShellOutput::ok("hello\n")
536 .into_step_result(false)
537 .expect("exit code 0 succeeds");
538
539 assert_eq!(output.output["stdout"], "hello\n");
540 assert_eq!(output.output["stderr"], "");
541 assert_eq!(output.output["exit_code"], 0);
542 assert_eq!(output.cost_usd, Decimal::ZERO);
543 }
544
545 #[test]
546 fn shell_default_is_an_empty_success() {
547 let default = MockShellOutput::default();
548 assert_eq!(default.exit_code, 0);
549 assert!(default.stdout.is_empty());
550 assert!(default.stderr.is_empty());
551 }
552
553 #[test]
554 fn shell_non_zero_exit_is_an_operation_error() {
555 let err = MockShellOutput::failed(2, "x")
556 .into_step_result(false)
557 .expect_err("a non-zero exit code fails the step");
558
559 match err {
560 EngineError::Operation(OperationError::Shell { exit_code, stderr }) => {
561 assert_eq!(exit_code, 2);
562 assert_eq!(stderr, "x");
563 }
564 other => panic!("expected a shell operation error, got {other}"),
565 }
566 }
567
568 #[test]
569 fn shell_non_zero_exit_with_option_is_an_output() {
570 let output = MockShellOutput::failed(2, "x")
571 .into_step_result(true)
572 .expect("the option turns a non-zero exit into an output");
573
574 assert_eq!(output.output["exit_code"], 2);
575 assert_eq!(output.output["stderr"], "x");
576 assert!(!output.is_success());
577 }
578
579 #[test]
580 fn intercept_shell_non_zero_with_option_completes() {
581 let interceptor = MockInterceptor::new().shell(|_| Ok(MockShellOutput::failed(2, "x")));
582 let config = StepConfig::Shell(ShellConfig::new("x").exit_code_as_output());
583
584 let output = interceptor
585 .intercept(&config)
586 .expect("the shell mock answers")
587 .expect("the step completes");
588 assert_eq!(output.exit_code(), Some(2));
589
590 let plain = StepConfig::Shell(ShellConfig::new("x"));
591 assert!(
592 interceptor
593 .intercept(&plain)
594 .expect("the shell mock answers")
595 .is_err()
596 );
597 }
598
599 #[test]
600 fn http_json_carries_status_body_and_headers() {
601 let output = MockHttpResponse::json(201, &json!({"id": 7}))
602 .header("location", "/things/7")
603 .into_step_output();
604
605 assert_eq!(output.output["status"], 201);
606 assert_eq!(output.output["body"], r#"{"id":7}"#);
607 assert_eq!(output.output["headers"]["location"], "/things/7");
608 }
609
610 #[test]
611 fn http_default_is_an_empty_200() {
612 let default = MockHttpResponse::default();
613 assert_eq!(default.status, 200);
614 assert!(default.body.is_empty());
615 assert!(default.headers.is_empty());
616 }
617
618 #[test]
619 fn http_non_2xx_is_still_an_output() {
620 let output = MockHttpResponse::text(500, "boom").into_step_output();
621 assert_eq!(output.status(), Some(500));
622 assert_eq!(output.body(), "boom");
623 }
624
625 #[test]
626 fn intercept_declines_agent_steps() {
627 let interceptor = MockInterceptor::new().shell(|_| Ok(MockShellOutput::ok("x")));
628 let config = StepConfig::Agent(AgentStepConfig::new("review this"));
629
630 assert!(interceptor.intercept(&config).is_none());
631 }
632
633 #[test]
634 fn intercept_declines_shell_steps_without_a_shell_mock() {
635 let interceptor = MockInterceptor::new();
636 let config = StepConfig::Shell(ShellConfig::new("echo hi"));
637
638 assert!(interceptor.intercept(&config).is_none());
639 }
640
641 #[test]
642 fn intercept_approval_returns_the_configured_outcome() {
643 let interceptor = MockInterceptor::new().approval(ApprovalOutcome::reject("nope"));
644 let config = ApprovalConfig::new("Approve?");
645
646 assert_eq!(
647 interceptor.intercept_approval("gate", &config),
648 Some(ApprovalOutcome::reject("nope"))
649 );
650 assert_eq!(
651 MockInterceptor::new().intercept_approval("gate", &config),
652 None
653 );
654 }
655
656 #[test]
657 fn intercept_human_input_returns_the_mocked_answer() {
658 let interceptor = MockInterceptor::new().human_input(|name, cfg| {
659 HumanInputOutcome::Provided(json!({"step": name, "message": cfg.message()}))
660 });
661 let config = HumanInputConfig::new("Answer?");
662 let answer = json!({"step": "clarify", "message": "Answer?"});
663
664 assert_eq!(
665 interceptor.intercept_human_input("clarify", &config, &json!({})),
666 Some(HumanInputOutcome::Provided(answer))
667 );
668 }
669
670 #[test]
671 fn intercept_human_input_declines_without_a_mock() {
672 let config = HumanInputConfig::new("Answer?");
673
674 assert_eq!(
675 MockInterceptor::new().intercept_human_input("clarify", &config, &json!({})),
676 None
677 );
678 }
679
680 #[test]
681 fn debug_reports_which_seams_are_mocked() {
682 let interceptor = MockInterceptor::new().http(|_| Ok(MockHttpResponse::default()));
683 let rendered = format!("{interceptor:?}");
684
685 assert!(rendered.contains("shell: false"));
686 assert!(rendered.contains("http: true"));
687 }
688
689 #[tokio::test]
690 async fn missing_agent_provider_names_the_three_constructors() {
691 let config = AgentConfig::new("anything");
692 let err = MissingAgentProvider
693 .invoke(&config)
694 .await
695 .expect_err("no agent backend is configured");
696
697 let message = err.to_string();
698 assert!(message.contains("with_mock_agent"));
699 assert!(message.contains("with_recorded_agent"));
700 assert!(message.contains("with_agent_provider"));
701 }
702
703 #[tokio::test]
704 async fn mock_agent_provider_runs_the_closure() {
705 let provider = MockAgentProvider::new(|cfg| {
706 let echoed = json!({"echoed": cfg.prompt.clone()});
707 Ok(AgentOutput::new(echoed))
708 });
709 let config = AgentConfig::new("say hi");
710
711 let output = provider.invoke(&config).await.expect("the mock succeeded");
712
713 assert_eq!(output.value["echoed"], "say hi");
714 }
715}