luft_core/
testing.rs

1//! Test utilities for resume and integration testing.
2//!
3//! Provides instrumented backends that can simulate crashes, blocking,
4//! and call recording — useful for testing crash-and-resume scenarios.
5
6use crate::contract::backend::{
7    AgentBackend, AgentCapabilities, AgentResult, AgentStatus, AgentTask, BackendError, LogRef,
8    RunContext,
9};
10use crate::contract::ids::TokenUsage;
11use serde_json::Value;
12use std::path::Path;
13use std::sync::atomic::{AtomicU64, Ordering};
14use std::sync::{Arc, Mutex};
15use std::time::Duration;
16
17#[derive(Debug, Clone)]
18pub struct CallRecord {
19    pub seq: u64,
20    pub agent_name: Option<String>,
21    pub thread_id: Option<String>,
22    pub prompt: String,
23}
24
25/// Backend that calls `std::process::exit(1)` after N calls.
26pub struct CrashBackend {
27    canned: Value,
28    crash_after: u64,
29    count: AtomicU64,
30}
31
32impl CrashBackend {
33    pub fn new(canned: Value, crash_after: u64) -> Self {
34        Self { canned, crash_after, count: AtomicU64::new(0) }
35    }
36}
37
38#[async_trait::async_trait]
39impl AgentBackend for CrashBackend {
40    fn id(&self) -> &'static str { "crash" }
41    fn capabilities(&self) -> AgentCapabilities {
42        AgentCapabilities { streaming: true, mcp_injection: false, structured_output: false, models: vec![] }
43    }
44    async fn run(&self, task: AgentTask, _ctx: RunContext) -> Result<AgentResult, BackendError> {
45        let n = self.count.fetch_add(1, Ordering::SeqCst) + 1;
46        eprintln!("[crash-backend] agent #{n}: {}", task.name.as_deref().unwrap_or("?"));
47        if n >= self.crash_after {
48            eprintln!("[crash-backend] exiting after {n} calls");
49            std::process::exit(1);
50        }
51        Ok(AgentResult {
52            agent_id: task.agent_id, status: AgentStatus::Ok, output: self.canned.clone(),
53            thread_id: task.thread_id.clone(), findings: vec![], tokens_used: TokenUsage::default(),
54            artifacts: vec![], logs: LogRef::default(),
55        })
56    }
57    fn as_any(&self) -> &dyn std::any::Any { self }
58}
59
60/// Backend that records all dispatched agent names.
61#[derive(Clone)]
62pub struct CountingBackend {
63    canned: Value,
64    calls: Arc<Mutex<Vec<String>>>,
65}
66
67impl CountingBackend {
68    pub fn new(canned: Value) -> Self {
69        Self { canned, calls: Arc::new(Mutex::new(Vec::new())) }
70    }
71    pub fn dispatched_names(&self) -> Vec<String> { self.calls.lock().unwrap().clone() }
72    pub fn total_calls(&self) -> usize { self.calls.lock().unwrap().len() }
73}
74
75#[async_trait::async_trait]
76impl AgentBackend for CountingBackend {
77    fn id(&self) -> &'static str { "counting" }
78    fn capabilities(&self) -> AgentCapabilities {
79        AgentCapabilities { streaming: true, mcp_injection: false, structured_output: false, models: vec![] }
80    }
81    async fn run(&self, task: AgentTask, _ctx: RunContext) -> Result<AgentResult, BackendError> {
82        self.calls.lock().unwrap().push(task.name.clone().unwrap_or_default());
83        Ok(AgentResult {
84            agent_id: task.agent_id, status: AgentStatus::Ok, output: self.canned.clone(),
85            thread_id: task.thread_id.clone(), findings: vec![], tokens_used: TokenUsage::default(),
86            artifacts: vec![], logs: LogRef::default(),
87        })
88    }
89    fn as_any(&self) -> &dyn std::any::Any { self }
90}
91
92/// Backend with shared state via Arc for resume tests.
93#[derive(Clone)]
94pub struct SharedBackend {
95    canned: Value,
96    call_count: Arc<AtomicU64>,
97    pub block_on: Arc<Mutex<Option<u64>>>,
98    pub fail_on: Arc<Mutex<Option<u64>>>,
99    calls: Arc<Mutex<Vec<CallRecord>>>,
100}
101
102impl SharedBackend {
103    pub fn new(canned: Value) -> Self {
104        Self {
105            canned, call_count: Arc::new(AtomicU64::new(0)),
106            block_on: Arc::new(Mutex::new(None)), fail_on: Arc::new(Mutex::new(None)),
107            calls: Arc::new(Mutex::new(Vec::new())),
108        }
109    }
110    pub fn with_block_on(self, n: u64) -> Self { *self.block_on.lock().unwrap() = Some(n); self }
111    pub fn with_fail_on(self, n: u64) -> Self { *self.fail_on.lock().unwrap() = Some(n); self }
112    pub fn total_calls(&self) -> usize { self.calls.lock().unwrap().len() }
113    pub fn calls_snapshot(&self) -> Vec<CallRecord> { self.calls.lock().unwrap().clone() }
114    pub fn dispatched_names(&self) -> Vec<String> {
115        self.calls.lock().unwrap().iter().map(|c| c.agent_name.clone().unwrap_or_default()).collect()
116    }
117    pub fn mirror(&self) -> Self {
118        Self {
119            canned: self.canned.clone(), call_count: self.call_count.clone(),
120            block_on: self.block_on.clone(), fail_on: self.fail_on.clone(), calls: self.calls.clone(),
121        }
122    }
123}
124
125#[async_trait::async_trait]
126impl AgentBackend for SharedBackend {
127    fn id(&self) -> &'static str { "shared" }
128    fn capabilities(&self) -> AgentCapabilities {
129        AgentCapabilities { streaming: true, mcp_injection: false, structured_output: false, models: vec![] }
130    }
131    async fn run(&self, task: AgentTask, ctx: RunContext) -> Result<AgentResult, BackendError> {
132        let seq = self.call_count.fetch_add(1, Ordering::SeqCst) + 1;
133        self.calls.lock().unwrap().push(CallRecord {
134            seq, agent_name: task.name.clone(), thread_id: task.thread_id.clone(), prompt: task.prompt.clone(),
135        });
136        if self.fail_on.lock().unwrap().map(|n| n == seq).unwrap_or(false) {
137            return Err(BackendError::Execution("simulated failure".into()));
138        }
139        if self.block_on.lock().unwrap().map(|n| n == seq).unwrap_or(false) {
140            ctx.cancel.cancelled().await;
141            return Err(BackendError::Cancelled);
142        }
143        Ok(AgentResult {
144            agent_id: task.agent_id, status: AgentStatus::Ok, output: self.canned.clone(),
145            thread_id: task.thread_id.clone(), findings: vec![], tokens_used: TokenUsage::default(),
146            artifacts: vec![], logs: LogRef::default(),
147        })
148    }
149    fn as_any(&self) -> &dyn std::any::Any { self }
150}
151
152pub async fn wait_for_calls(backend: &SharedBackend, n: usize, timeout_ms: u64) {
153    let deadline = tokio::time::sleep(Duration::from_millis(timeout_ms));
154    tokio::pin!(deadline);
155    loop {
156        if backend.total_calls() >= n { return; }
157        tokio::select! {
158            _ = &mut deadline => panic!("timeout waiting for {n} agent calls (got {})", backend.total_calls()),
159            _ = tokio::time::sleep(Duration::from_millis(25)) => {}
160        }
161    }
162}
163
164pub async fn read_checkpoint(base: &Path, run_dir: &str) -> Value {
165    let path = base.join(run_dir).join("checkpoint.json");
166    match tokio::fs::read(&path).await {
167        Ok(bytes) => serde_json::from_slice(&bytes).unwrap_or(Value::Null),
168        Err(_) => Value::Null,
169    }
170}
171
172pub fn completed_span_names(cp: &Value) -> Vec<String> {
173    cp.get("phase_state")
174        .and_then(|ps| ps.get("completed_spans"))
175        .and_then(|cs| cs.as_array())
176        .map(|arr| arr.iter().filter_map(|v| v.get("name").and_then(|n| n.as_str()).map(String::from)).collect())
177        .unwrap_or_default()
178}
luft_core/testing.rs

luft_core/
testing.rs
