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(self) -> Result<StepOutput, EngineError> {
99 if self.exit_code != 0 {
100 return Err(EngineError::Operation(OperationError::Shell {
101 exit_code: self.exit_code,
102 stderr: self.stderr,
103 }));
104 }
105
106 Ok(StepOutput {
107 output: json!({
108 "stdout": self.stdout,
109 "stderr": self.stderr,
110 "exit_code": self.exit_code,
111 }),
112 duration_ms: 0,
113 cost_usd: Decimal::ZERO,
114 input_tokens: None,
115 cache_read_input_tokens: None,
116 cache_creation_input_tokens: None,
117 output_tokens: None,
118 model: None,
119 debug_messages: None,
120 artifacts: StepArtifacts::default(),
121 })
122 }
123}
124
125#[derive(Debug, Clone, PartialEq, Eq)]
144pub struct MockHttpResponse {
145 pub status: u16,
147 pub headers: Vec<(String, String)>,
149 pub body: String,
151}
152
153impl Default for MockHttpResponse {
154 fn default() -> Self {
155 Self {
156 status: 200,
157 headers: Vec::new(),
158 body: String::new(),
159 }
160 }
161}
162
163impl MockHttpResponse {
164 pub fn ok(body: &Value) -> Self {
177 Self::json(200, body)
178 }
179
180 pub fn json(status: u16, body: &Value) -> Self {
192 Self {
193 status,
194 headers: Vec::new(),
195 body: to_string(body).unwrap_or_else(|_| body.to_string()),
197 }
198 }
199
200 pub fn text(status: u16, body: &str) -> Self {
211 Self {
212 status,
213 headers: Vec::new(),
214 body: body.to_string(),
215 }
216 }
217
218 pub fn header(mut self, name: &str, value: &str) -> Self {
229 self.headers.push((name.to_string(), value.to_string()));
230 self
231 }
232
233 pub(crate) fn into_step_output(self) -> StepOutput {
236 let headers: BTreeMap<String, String> = self.headers.into_iter().collect();
237 StepOutput {
238 output: json!({
239 "status": self.status,
240 "headers": headers,
241 "body": self.body,
242 }),
243 duration_ms: 0,
244 cost_usd: Decimal::ZERO,
245 input_tokens: None,
246 cache_read_input_tokens: None,
247 cache_creation_input_tokens: None,
248 output_tokens: None,
249 model: None,
250 debug_messages: None,
251 artifacts: StepArtifacts::default(),
252 }
253 }
254}
255
256pub type ShellMock =
258 Arc<dyn Fn(&ShellConfig) -> Result<MockShellOutput, OperationError> + Send + Sync>;
259
260pub type HttpMock =
262 Arc<dyn Fn(&HttpConfig) -> Result<MockHttpResponse, OperationError> + Send + Sync>;
263
264pub type HumanInputMock = Arc<dyn Fn(&str, &HumanInputConfig) -> HumanInputOutcome + Send + Sync>;
266
267pub type AgentMock = Arc<dyn Fn(&AgentConfig) -> Result<AgentOutput, AgentError> + Send + Sync>;
269
270#[derive(Clone, Default)]
292pub struct MockInterceptor {
293 shell: Option<ShellMock>,
294 http: Option<HttpMock>,
295 approval: Option<ApprovalOutcome>,
296 human_input: Option<HumanInputMock>,
297}
298
299impl MockInterceptor {
300 pub fn new() -> Self {
314 Self::default()
315 }
316
317 pub fn shell(
329 mut self,
330 f: impl Fn(&ShellConfig) -> Result<MockShellOutput, OperationError> + Send + Sync + 'static,
331 ) -> Self {
332 self.shell = Some(Arc::new(f));
333 self
334 }
335
336 pub fn http(
349 mut self,
350 f: impl Fn(&HttpConfig) -> Result<MockHttpResponse, OperationError> + Send + Sync + 'static,
351 ) -> Self {
352 self.http = Some(Arc::new(f));
353 self
354 }
355
356 pub fn approval(mut self, outcome: ApprovalOutcome) -> Self {
367 self.approval = Some(outcome);
368 self
369 }
370
371 pub fn human_input(
384 mut self,
385 f: impl Fn(&str, &HumanInputConfig) -> HumanInputOutcome + Send + Sync + 'static,
386 ) -> Self {
387 self.human_input = Some(Arc::new(f));
388 self
389 }
390}
391
392impl fmt::Debug for MockInterceptor {
393 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
394 f.debug_struct("MockInterceptor")
396 .field("shell", &self.shell.is_some())
397 .field("http", &self.http.is_some())
398 .field("approval", &self.approval)
399 .field("human_input", &self.human_input.is_some())
400 .finish()
401 }
402}
403
404impl StepInterceptor for MockInterceptor {
405 fn intercept(&self, config: &StepConfig) -> Option<Result<StepOutput, EngineError>> {
406 match config {
407 StepConfig::Shell(cfg) => {
408 let mock = self.shell.as_ref()?;
409 Some(match mock(cfg) {
410 Ok(out) => out.into_step_result(),
411 Err(err) => Err(EngineError::Operation(err)),
412 })
413 }
414 StepConfig::Http(cfg) => {
415 let mock = self.http.as_ref()?;
416 Some(match mock(cfg) {
417 Ok(res) => Ok(res.into_step_output()),
418 Err(err) => Err(EngineError::Operation(err)),
419 })
420 }
421 _ => None,
423 }
424 }
425
426 fn intercept_approval(&self, _name: &str, _config: &ApprovalConfig) -> Option<ApprovalOutcome> {
427 self.approval.clone()
428 }
429
430 fn intercept_human_input(
431 &self,
432 name: &str,
433 config: &HumanInputConfig,
434 _schema: &Value,
435 ) -> Option<HumanInputOutcome> {
436 self.human_input.as_ref().map(|f| f(name, config))
437 }
438}
439
440pub struct MockAgentProvider {
456 f: AgentMock,
457}
458
459impl MockAgentProvider {
460 pub fn new(
473 f: impl Fn(&AgentConfig) -> Result<AgentOutput, AgentError> + Send + Sync + 'static,
474 ) -> Self {
475 Self { f: Arc::new(f) }
476 }
477}
478
479impl fmt::Debug for MockAgentProvider {
480 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
481 f.debug_struct("MockAgentProvider").finish_non_exhaustive()
482 }
483}
484
485impl AgentProvider for MockAgentProvider {
486 fn invoke<'a>(&'a self, config: &'a AgentConfig) -> InvokeFuture<'a> {
487 let result = (self.f)(config);
488 Box::pin(async move { result })
489 }
490}
491
492#[derive(Debug, Default, Clone, Copy)]
510pub struct MissingAgentProvider;
511
512impl AgentProvider for MissingAgentProvider {
513 fn invoke<'a>(&'a self, _config: &'a AgentConfig) -> InvokeFuture<'a> {
514 Box::pin(async move {
515 Err(AgentError::ProcessFailed {
516 exit_code: -1,
517 stderr: MISSING_AGENT_PROVIDER.to_string(),
518 })
519 })
520 }
521}
522
523#[cfg(test)]
524mod tests {
525 use super::*;
526
527 use crate::config::AgentStepConfig;
528
529 #[test]
530 fn shell_ok_maps_to_the_real_executor_output_shape() {
531 let output = MockShellOutput::ok("hello\n")
532 .into_step_result()
533 .expect("exit code 0 succeeds");
534
535 assert_eq!(output.output["stdout"], "hello\n");
536 assert_eq!(output.output["stderr"], "");
537 assert_eq!(output.output["exit_code"], 0);
538 assert_eq!(output.cost_usd, Decimal::ZERO);
539 }
540
541 #[test]
542 fn shell_default_is_an_empty_success() {
543 let default = MockShellOutput::default();
544 assert_eq!(default.exit_code, 0);
545 assert!(default.stdout.is_empty());
546 assert!(default.stderr.is_empty());
547 }
548
549 #[test]
550 fn shell_non_zero_exit_is_an_operation_error() {
551 let err = MockShellOutput::failed(2, "x")
552 .into_step_result()
553 .expect_err("a non-zero exit code fails the step");
554
555 match err {
556 EngineError::Operation(OperationError::Shell { exit_code, stderr }) => {
557 assert_eq!(exit_code, 2);
558 assert_eq!(stderr, "x");
559 }
560 other => panic!("expected a shell operation error, got {other}"),
561 }
562 }
563
564 #[test]
565 fn http_json_carries_status_body_and_headers() {
566 let output = MockHttpResponse::json(201, &json!({"id": 7}))
567 .header("location", "/things/7")
568 .into_step_output();
569
570 assert_eq!(output.output["status"], 201);
571 assert_eq!(output.output["body"], r#"{"id":7}"#);
572 assert_eq!(output.output["headers"]["location"], "/things/7");
573 }
574
575 #[test]
576 fn http_default_is_an_empty_200() {
577 let default = MockHttpResponse::default();
578 assert_eq!(default.status, 200);
579 assert!(default.body.is_empty());
580 assert!(default.headers.is_empty());
581 }
582
583 #[test]
584 fn http_non_2xx_is_still_an_output() {
585 let output = MockHttpResponse::text(500, "boom").into_step_output();
586 assert_eq!(output.status(), Some(500));
587 assert_eq!(output.body(), "boom");
588 }
589
590 #[test]
591 fn intercept_declines_agent_steps() {
592 let interceptor = MockInterceptor::new().shell(|_| Ok(MockShellOutput::ok("x")));
593 let config = StepConfig::Agent(AgentStepConfig::new("review this"));
594
595 assert!(interceptor.intercept(&config).is_none());
596 }
597
598 #[test]
599 fn intercept_declines_shell_steps_without_a_shell_mock() {
600 let interceptor = MockInterceptor::new();
601 let config = StepConfig::Shell(ShellConfig::new("echo hi"));
602
603 assert!(interceptor.intercept(&config).is_none());
604 }
605
606 #[test]
607 fn intercept_approval_returns_the_configured_outcome() {
608 let interceptor = MockInterceptor::new().approval(ApprovalOutcome::reject("nope"));
609 let config = ApprovalConfig::new("Approve?");
610
611 assert_eq!(
612 interceptor.intercept_approval("gate", &config),
613 Some(ApprovalOutcome::reject("nope"))
614 );
615 assert_eq!(
616 MockInterceptor::new().intercept_approval("gate", &config),
617 None
618 );
619 }
620
621 #[test]
622 fn intercept_human_input_returns_the_mocked_answer() {
623 let interceptor = MockInterceptor::new().human_input(|name, cfg| {
624 HumanInputOutcome::Provided(json!({"step": name, "message": cfg.message()}))
625 });
626 let config = HumanInputConfig::new("Answer?");
627 let answer = json!({"step": "clarify", "message": "Answer?"});
628
629 assert_eq!(
630 interceptor.intercept_human_input("clarify", &config, &json!({})),
631 Some(HumanInputOutcome::Provided(answer))
632 );
633 }
634
635 #[test]
636 fn intercept_human_input_declines_without_a_mock() {
637 let config = HumanInputConfig::new("Answer?");
638
639 assert_eq!(
640 MockInterceptor::new().intercept_human_input("clarify", &config, &json!({})),
641 None
642 );
643 }
644
645 #[test]
646 fn debug_reports_which_seams_are_mocked() {
647 let interceptor = MockInterceptor::new().http(|_| Ok(MockHttpResponse::default()));
648 let rendered = format!("{interceptor:?}");
649
650 assert!(rendered.contains("shell: false"));
651 assert!(rendered.contains("http: true"));
652 }
653
654 #[tokio::test]
655 async fn missing_agent_provider_names_the_three_constructors() {
656 let config = AgentConfig::new("anything");
657 let err = MissingAgentProvider
658 .invoke(&config)
659 .await
660 .expect_err("no agent backend is configured");
661
662 let message = err.to_string();
663 assert!(message.contains("with_mock_agent"));
664 assert!(message.contains("with_recorded_agent"));
665 assert!(message.contains("with_agent_provider"));
666 }
667
668 #[tokio::test]
669 async fn mock_agent_provider_runs_the_closure() {
670 let provider = MockAgentProvider::new(|cfg| {
671 let echoed = json!({"echoed": cfg.prompt.clone()});
672 Ok(AgentOutput::new(echoed))
673 });
674 let config = AgentConfig::new("say hi");
675
676 let output = provider.invoke(&config).await.expect("the mock succeeded");
677
678 assert_eq!(output.value["echoed"], "say hi");
679 }
680}