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