1use crate::assertion::Assertion;
2use std::ops::RangeInclusive;
3
4use crate::agent::{AgentConfig, AgentOutput};
5
6pub 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
20pub 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
41pub 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 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 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 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 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}