Skip to main content

spice_framework/
test_case.rs

1use crate::assertion::Assertion;
2use std::ops::RangeInclusive;
3
4use crate::agent::{AgentConfig, AgentOutput};
5
6/// A single test case for an agent.
7pub struct TestCase {
8    pub id: String,
9    pub name: Option<String>,
10    pub user_message: String,
11    pub config: AgentConfig,
12    pub assertions: Vec<Assertion>,
13    pub tags: Vec<String>,
14    pub retries: usize,
15    pub consensus_runs: Option<usize>,
16    pub consensus_required: Option<usize>,
17    pub timeout: Option<std::time::Duration>,
18}
19
20/// A collection of test cases.
21pub struct TestSuite {
22    pub name: String,
23    pub tests: Vec<TestCase>,
24    pub default_config: AgentConfig,
25    pub default_retries: usize,
26    pub default_timeout: Option<std::time::Duration>,
27}
28
29impl Default for TestSuite {
30    fn default() -> Self {
31        Self {
32            name: "Test Suite".into(),
33            tests: vec![],
34            default_config: AgentConfig::empty(),
35            default_retries: 0,
36            default_timeout: None,
37        }
38    }
39}
40
41/// Builder for constructing test cases fluently.
42pub struct TestCaseBuilder {
43    id: String,
44    user_message: String,
45    name: Option<String>,
46    config: Option<AgentConfig>,
47    assertions: Vec<Assertion>,
48    tags: Vec<String>,
49    retries: usize,
50    consensus_runs: Option<usize>,
51    consensus_required: Option<usize>,
52    timeout: Option<std::time::Duration>,
53}
54
55impl TestCaseBuilder {
56    pub fn new(id: impl Into<String>, user_message: impl Into<String>) -> Self {
57        Self {
58            id: id.into(),
59            user_message: user_message.into(),
60            name: None,
61            config: None,
62            assertions: vec![],
63            tags: vec![],
64            retries: 0,
65            consensus_runs: None,
66            consensus_required: None,
67            timeout: None,
68        }
69    }
70
71    pub fn name(mut self, name: impl Into<String>) -> Self {
72        self.name = Some(name.into());
73        self
74    }
75
76    pub fn tag(mut self, tag: impl Into<String>) -> Self {
77        self.tags.push(tag.into());
78        self
79    }
80
81    pub fn tags(mut self, tags: &[&str]) -> Self {
82        self.tags.extend(tags.iter().map(|s| s.to_string()));
83        self
84    }
85
86    pub fn config(mut self, config: AgentConfig) -> Self {
87        self.config = Some(config);
88        self
89    }
90
91    pub fn config_json(mut self, data: serde_json::Value) -> Self {
92        self.config = Some(AgentConfig::new(data));
93        self
94    }
95
96    pub fn retries(mut self, n: usize) -> Self {
97        self.retries = n;
98        self
99    }
100
101    pub fn consensus(mut self, runs: usize, required: usize) -> Self {
102        self.consensus_runs = Some(runs);
103        self.consensus_required = Some(required);
104        self
105    }
106
107    pub fn timeout(mut self, duration: std::time::Duration) -> Self {
108        self.timeout = Some(duration);
109        self
110    }
111
112    // --- Assertion builders ---
113
114    pub fn expect_tools(mut self, tools: &[&str]) -> Self {
115        self.assertions.push(Assertion::ExpectTools(
116            tools.iter().map(|s| s.to_string()).collect(),
117        ));
118        self
119    }
120
121    pub fn forbid_tools(mut self, tools: &[&str]) -> Self {
122        self.assertions.push(Assertion::ForbidTools(
123            tools.iter().map(|s| s.to_string()).collect(),
124        ));
125        self
126    }
127
128    pub fn expect_any_tool(mut self) -> Self {
129        self.assertions.push(Assertion::ExpectAnyTool);
130        self
131    }
132
133    pub fn expect_no_tools(mut self) -> Self {
134        self.assertions.push(Assertion::ExpectNoTools);
135        self
136    }
137
138    pub fn expect_text_contains(mut self, s: impl Into<String>) -> Self {
139        self.assertions
140            .push(Assertion::ExpectTextContains(s.into()));
141        self
142    }
143
144    pub fn expect_text_not_contains(mut self, s: impl Into<String>) -> Self {
145        self.assertions
146            .push(Assertion::ExpectTextNotContains(s.into()));
147        self
148    }
149
150    pub fn expect_turns(mut self, range: RangeInclusive<usize>) -> Self {
151        self.assertions.push(Assertion::ExpectTurns(range));
152        self
153    }
154
155    pub fn expect_tools_within_allowlist(mut self) -> Self {
156        self.assertions.push(Assertion::ExpectToolsWithinAllowlist);
157        self
158    }
159
160    pub fn expect_no_error(mut self) -> Self {
161        self.assertions.push(Assertion::ExpectNoError);
162        self
163    }
164
165    pub fn expect_tool_args(
166        mut self,
167        tool: impl Into<String>,
168        args: serde_json::Value,
169    ) -> Self {
170        self.assertions
171            .push(Assertion::ExpectToolArgs(tool.into(), args));
172        self
173    }
174
175    pub fn expect_tool_args_contain(
176        mut self,
177        tool: impl Into<String>,
178        partial: serde_json::Value,
179    ) -> Self {
180        self.assertions
181            .push(Assertion::ExpectToolArgsContain(tool.into(), partial));
182        self
183    }
184
185    pub fn expect_tool_arg(
186        mut self,
187        tool: impl Into<String>,
188        param: impl Into<String>,
189        value: serde_json::Value,
190    ) -> Self {
191        self.assertions
192            .push(Assertion::ExpectToolArg(tool.into(), param.into(), value));
193        self
194    }
195
196    pub fn expect_tool_arg_exists(
197        mut self,
198        tool: impl Into<String>,
199        param: impl Into<String>,
200    ) -> Self {
201        self.assertions
202            .push(Assertion::ExpectToolArgExists(tool.into(), param.into()));
203        self
204    }
205
206    pub fn expect_tool_call_count(mut self, tool: impl Into<String>, count: usize) -> Self {
207        self.assertions
208            .push(Assertion::ExpectToolCallCount(tool.into(), count));
209        self
210    }
211
212    pub fn expect_tool_call_order(mut self, order: &[&str]) -> Self {
213        self.assertions.push(Assertion::ExpectToolCallOrder(
214            order.iter().map(|s| s.to_string()).collect(),
215        ));
216        self
217    }
218
219    pub fn expect_tool_on_turn(mut self, turn: usize, tool: impl Into<String>) -> Self {
220        self.assertions
221            .push(Assertion::ExpectToolOnTurn(turn, tool.into()));
222        self
223    }
224
225    pub fn expect<F>(mut self, f: F) -> Self
226    where
227        F: Fn(&AgentOutput) -> Result<(), String> + Send + Sync + 'static,
228    {
229        self.assertions.push(Assertion::Custom(Box::new(f)));
230        self
231    }
232
233    // --- Multi-turn assertion builders ---
234
235    pub fn with_role(self, role: &str) -> Self {
236        self.config_json(serde_json::json!({"role": role}))
237    }
238
239    pub fn expect_tools_in_turn_range(
240        mut self,
241        range: RangeInclusive<usize>,
242        tools: &[&str],
243    ) -> Self {
244        self.assertions.push(Assertion::ExpectToolsInTurnRange(
245            range,
246            tools.iter().map(|s| s.to_string()).collect(),
247        ));
248        self
249    }
250
251    pub fn forbid_tools_in_turn_range(
252        mut self,
253        range: RangeInclusive<usize>,
254        tools: &[&str],
255    ) -> Self {
256        self.assertions.push(Assertion::ForbidToolsInTurnRange(
257            range,
258            tools.iter().map(|s| s.to_string()).collect(),
259        ));
260        self
261    }
262
263    pub fn expect_final_tool(mut self, tool: &str) -> Self {
264        self.assertions
265            .push(Assertion::ExpectFinalTool(tool.to_string()));
266        self
267    }
268
269    pub fn expect_final_tool_arg(
270        mut self,
271        tool: &str,
272        param: &str,
273        value: serde_json::Value,
274    ) -> Self {
275        self.assertions.push(Assertion::ExpectFinalToolArg(
276            tool.to_string(),
277            param.to_string(),
278            value,
279        ));
280        self
281    }
282
283    /// Shorthand: gathering tools before default action tools (say_to_user, task_fully_completed).
284    pub fn expect_gathering_phase(mut self, gather_tools: &[&str]) -> Self {
285        self.assertions.push(Assertion::ExpectGatheringBeforeAction(
286            gather_tools.iter().map(|s| s.to_string()).collect(),
287            vec!["say_to_user".to_string(), "task_fully_completed".to_string()],
288        ));
289        self
290    }
291
292    /// Explicit: gathering tools before specified action tools.
293    pub fn expect_gathering_before_action(
294        mut self,
295        gather_tools: &[&str],
296        action_tools: &[&str],
297    ) -> Self {
298        self.assertions.push(Assertion::ExpectGatheringBeforeAction(
299            gather_tools.iter().map(|s| s.to_string()).collect(),
300            action_tools.iter().map(|s| s.to_string()).collect(),
301        ));
302        self
303    }
304
305    pub fn expect_tool_only_on_final_turn(mut self, tool: &str) -> Self {
306        self.assertions
307            .push(Assertion::ExpectToolOnlyOnFinalTurn(tool.to_string()));
308        self
309    }
310
311    pub fn build(self) -> TestCase {
312        TestCase {
313            id: self.id,
314            name: self.name,
315            user_message: self.user_message,
316            config: self.config.unwrap_or_else(AgentConfig::empty),
317            assertions: self.assertions,
318            tags: self.tags,
319            retries: self.retries,
320            consensus_runs: self.consensus_runs,
321            consensus_required: self.consensus_required,
322            timeout: self.timeout,
323        }
324    }
325}