1use super::{Tool, Result, ToolError, common_options, parse_output_format, OutputFormat};
2use clap::{Arg, ArgMatches, Command};
3use colored::*;
4use std::path::Path;
5use std::fs;
6use syn::{parse_file, FnArg, Pat, ReturnType, visit::Visit};
7use quote::ToTokens;
8#[derive(Debug, Clone)]
9pub struct TestGenTool;
10#[derive(Debug)]
11struct FunctionInfo {
12 name: String,
13 params: Vec<ParamInfo>,
14 return_type: Option<String>,
15 is_async: bool,
16 visibility: String,
17}
18#[derive(Debug)]
19struct ParamInfo {
20 name: String,
21 ty: String,
22 is_reference: bool,
23}
24#[derive(Debug, serde::Serialize)]
25struct GeneratedTest {
26 function_name: String,
27 test_name: String,
28 test_code: String,
29}
30impl TestGenTool {
31 pub fn new() -> Self {
32 Self
33 }
34 fn parse_rust_file(&self, file_path: &Path) -> Result<Vec<FunctionInfo>> {
35 let content = fs::read_to_string(file_path)
36 .map_err(|e| ToolError::ExecutionFailed(
37 format!("Failed to read file: {}", e),
38 ))?;
39 let syntax = parse_file(&content)
40 .map_err(|e| ToolError::ExecutionFailed(
41 format!("Failed to parse Rust code: {}", e),
42 ))?;
43 let mut visitor = FunctionVisitor::new();
44 visitor.visit_file(&syntax);
45 Ok(visitor.functions)
46 }
47 fn generate_test_for_function(
48 &self,
49 func: &FunctionInfo,
50 test_type: &str,
51 ) -> GeneratedTest {
52 let test_name = format!("test_{}_{}", func.name, test_type);
53 let test_code = match test_type {
54 "unit" => self.generate_unit_test(func),
55 "integration" => self.generate_integration_test(func),
56 "property" => self.generate_property_test(func),
57 _ => self.generate_unit_test(func),
58 };
59 GeneratedTest {
60 function_name: func.name.clone(),
61 test_name,
62 test_code,
63 }
64 }
65 fn generate_unit_test(&self, func: &FunctionInfo) -> String {
66 let mut code = format!("/// Unit test for `{}`\n", func.name);
67 code.push_str("#[test]\n");
68 if func.is_async {
69 code.push_str("#[tokio::test]\n");
70 }
71 code.push_str(
72 &format!(
73 "fn {}() {{\n", self.snake_to_pascal(& format!("test_{}", func.name))
74 ),
75 );
76 for param in &func.params {
77 let mock_value = self.generate_mock_value(¶m.ty);
78 code.push_str(&format!(" let {} = {};\n", param.name, mock_value));
79 }
80 let params_str = func
81 .params
82 .iter()
83 .map(|p| {
84 if p.is_reference { format!("&{}", p.name) } else { p.name.clone() }
85 })
86 .collect::<Vec<_>>()
87 .join(", ");
88 if func.is_async {
89 code.push_str(
90 &format!(" let result = {}({}).await;\n", func.name, params_str),
91 );
92 } else {
93 code.push_str(&format!(" let result = {}({});\n", func.name, params_str));
94 }
95 if let Some(return_type) = &func.return_type {
96 if return_type.contains("Result") {
97 code.push_str(" assert!(result.is_ok());\n");
98 } else if return_type.contains("Option") {
99 code.push_str(" assert!(result.is_some());\n");
100 } else if return_type == "bool" {
101 code.push_str(" assert!(result);\n");
102 } else if return_type.contains("Vec") || return_type.contains("HashMap") {
103 code.push_str(" assert!(!result.is_empty());\n");
104 } else {
105 code.push_str(" // Add your assertions here\n");
106 code.push_str(" assert!(true); // Placeholder assertion\n");
107 }
108 } else {
109 code.push_str(
110 " // Function returns nothing - add your test logic here\n",
111 );
112 code.push_str(" assert!(true); // Placeholder assertion\n");
113 }
114 code.push_str("}\n");
115 code
116 }
117 fn generate_integration_test(&self, func: &FunctionInfo) -> String {
118 let mut code = format!("/// Integration test for `{}`\n", func.name);
119 code.push_str("#[test]\n");
120 if func.is_async {
121 code.push_str("#[tokio::test]\n");
122 }
123 code.push_str(
124 &format!(
125 "fn {}() {{\n", self.snake_to_pascal(& format!("integration_test_{}",
126 func.name))
127 ),
128 );
129 code.push_str(" // Setup test environment\n");
130 for param in &func.params {
131 let mock_value = self.generate_integration_mock_value(¶m.ty);
132 code.push_str(&format!(" let {} = {};\n", param.name, mock_value));
133 }
134 let params_str = func
135 .params
136 .iter()
137 .map(|p| {
138 if p.is_reference { format!("&{}", p.name) } else { p.name.clone() }
139 })
140 .collect::<Vec<_>>()
141 .join(", ");
142 if func.is_async {
143 code.push_str(
144 &format!(" let result = {}({}).await;\n", func.name, params_str),
145 );
146 } else {
147 code.push_str(&format!(" let result = {}({});\n", func.name, params_str));
148 }
149 code.push_str(" // Verify the result\n");
150 if let Some(return_type) = &func.return_type {
151 if return_type.contains("Result") {
152 code.push_str(" match result {\n");
153 code.push_str(" Ok(value) => {\n");
154 code.push_str(" // Add your success case assertions here\n");
155 code.push_str(" assert!(true); // Placeholder\n");
156 code.push_str(" }\n");
157 code.push_str(" Err(e) => {\n");
158 code.push_str(
159 " panic!(\"Integration test failed: {}\", e);\n",
160 );
161 code.push_str(" }\n");
162 code.push_str(" }\n");
163 } else {
164 code.push_str(" // Add your integration test assertions here\n");
165 code.push_str(" assert!(true); // Placeholder assertion\n");
166 }
167 }
168 code.push_str("}\n");
169 code
170 }
171 fn generate_property_test(&self, func: &FunctionInfo) -> String {
172 let mut code = format!("/// Property-based test for `{}`\n", func.name);
173 code.push_str("#[cfg(test)]\n");
174 code.push_str("mod property_tests {\n");
175 code.push_str(" use super::*;\n");
176 code.push_str(" use proptest::prelude::*;\n\n");
177 let has_numeric_params = func
178 .params
179 .iter()
180 .any(|p| {
181 let ty = p.ty.to_lowercase();
182 ty.contains("i32") || ty.contains("i64") || ty.contains("u32")
183 || ty.contains("u64") || ty.contains("f32") || ty.contains("f64")
184 || ty.contains("usize") || ty.contains("isize")
185 });
186 if has_numeric_params {
187 code.push_str(&format!(" proptest! {{\n"));
188 code.push_str(&format!(" #[test]\n"));
189 code.push_str(
190 &format!(
191 " fn {}({}) {{\n", func.name, self.generate_proptest_params(&
192 func.params)
193 ),
194 );
195 let params_str = func
196 .params
197 .iter()
198 .map(|p| {
199 if p.is_reference { format!("&{}", p.name) } else { p.name.clone() }
200 })
201 .collect::<Vec<_>>()
202 .join(", ");
203 if func.is_async {
204 code.push_str(
205 &format!(
206 " let result = {}({}).await;\n", func.name, params_str
207 ),
208 );
209 } else {
210 code.push_str(
211 &format!(" let result = {}({});\n", func.name, params_str),
212 );
213 }
214 if let Some(return_type) = &func.return_type {
215 if return_type.contains("Result") {
216 code.push_str(" prop_assert!(result.is_ok());\n");
217 } else if return_type.contains("bool") {
218 code.push_str(
219 " // Add property-based assertions for boolean results\n",
220 );
221 code.push_str(" prop_assert!(true); // Placeholder\n");
222 } else {
223 code.push_str(
224 " // Add your property-based assertions here\n",
225 );
226 code.push_str(" prop_assert!(true); // Placeholder\n");
227 }
228 }
229 code.push_str(" }\n");
230 code.push_str(" }\n");
231 } else {
232 code.push_str(
233 &format!(
234 " // Property-based testing not applicable for this function\n"
235 ),
236 );
237 code.push_str(&format!(" // Function doesn't have numeric parameters\n"));
238 }
239 code.push_str("}\n");
240 code
241 }
242 fn generate_mock_value(&self, ty: &str) -> String {
243 let ty_lower = ty.to_lowercase();
244 if ty_lower.contains("string") {
245 "\"test_value\".to_string()".to_string()
246 } else if ty_lower.contains("i32") {
247 "42".to_string()
248 } else if ty_lower.contains("i64") {
249 "42i64".to_string()
250 } else if ty_lower.contains("u32") {
251 "42u32".to_string()
252 } else if ty_lower.contains("u64") {
253 "42u64".to_string()
254 } else if ty_lower.contains("f32") {
255 "3.14f32".to_string()
256 } else if ty_lower.contains("f64") {
257 "3.14".to_string()
258 } else if ty_lower.contains("bool") {
259 "true".to_string()
260 } else if ty_lower.contains("vec") {
261 "vec![1, 2, 3]".to_string()
262 } else if ty_lower.contains("hashmap") {
263 "HashMap::new()".to_string()
264 } else if ty_lower.contains("option") {
265 "Some(\"test\".to_string())".to_string()
266 } else if ty_lower.contains("result") {
267 "Ok(\"success\".to_string())".to_string()
268 } else {
269 format!("{}::default()", ty)
270 }
271 }
272 fn generate_integration_mock_value(&self, ty: &str) -> String {
273 let ty_lower = ty.to_lowercase();
274 if ty_lower.contains("string") {
275 "\"integration_test_value\".to_string()".to_string()
276 } else if ty_lower.contains("i32") {
277 "100".to_string()
278 } else if ty_lower.contains("i64") {
279 "1000i64".to_string()
280 } else if ty_lower.contains("u32") {
281 "100u32".to_string()
282 } else if ty_lower.contains("u64") {
283 "1000u64".to_string()
284 } else if ty_lower.contains("f32") {
285 "1.414f32".to_string()
286 } else if ty_lower.contains("f64") {
287 "2.718".to_string()
288 } else if ty_lower.contains("bool") {
289 "false".to_string()
290 } else if ty_lower.contains("vec") {
291 "vec![10, 20, 30, 40, 50]".to_string()
292 } else if ty_lower.contains("hashmap") {
293 "{\n let mut map = HashMap::new();\n map.insert(\"key1\".to_string(), \"value1\".to_string());\n map\n }"
294 .to_string()
295 } else if ty_lower.contains("option") {
296 "None".to_string()
297 } else if ty_lower.contains("result") {
298 "Err(\"integration test error\".to_string())".to_string()
299 } else {
300 format!("{}::default()", ty)
301 }
302 }
303 fn generate_proptest_params(&self, params: &[ParamInfo]) -> String {
304 params
305 .iter()
306 .filter_map(|p| {
307 let ty_lower = p.ty.to_lowercase();
308 if ty_lower.contains("i32") {
309 Some(format!("{} in 0..1000i32", p.name))
310 } else if ty_lower.contains("i64") {
311 Some(format!("{} in 0..10000i64", p.name))
312 } else if ty_lower.contains("u32") {
313 Some(format!("{} in 0..1000u32", p.name))
314 } else if ty_lower.contains("u64") {
315 Some(format!("{} in 0..10000u64", p.name))
316 } else if ty_lower.contains("f32") {
317 Some(format!("{} in 0.0..1000.0f32", p.name))
318 } else if ty_lower.contains("f64") {
319 Some(format!("{} in 0.0..1000.0", p.name))
320 } else {
321 None
322 }
323 })
324 .collect::<Vec<_>>()
325 .join(", ")
326 }
327 fn snake_to_pascal(&self, snake_case: &str) -> String {
328 snake_case
329 .split('_')
330 .map(|word| {
331 let mut chars = word.chars();
332 match chars.next() {
333 None => String::new(),
334 Some(first) => {
335 first.to_uppercase().chain(chars.as_str().chars()).collect()
336 }
337 }
338 })
339 .collect()
340 }
341 fn display_generated_tests(&self, tests: &[GeneratedTest], format: OutputFormat) {
342 match format {
343 OutputFormat::Json => {
344 println!("{}", serde_json::to_string_pretty(tests).unwrap());
345 }
346 OutputFormat::Table => {
347 println!("{:<25} {:<30}", "Function", "Generated Test");
348 println!("{}", "─".repeat(55));
349 for test in tests {
350 println!("{:<25} {:<30}", test.function_name, test.test_name);
351 }
352 }
353 OutputFormat::Human => {
354 println!("{}", "🧪 Generated Tests".bold().blue());
355 println!("{}", "═".repeat(50).blue());
356 for test in tests {
357 println!(
358 "📝 {} -> {}", test.function_name.green(), test.test_name
359 .cyan()
360 );
361 println!("```rust");
362 println!("{}", test.test_code.trim());
363 println!("```");
364 println!();
365 }
366 println!("💡 {} tests generated successfully!", tests.len());
367 println!(
368 "📍 Add these tests to your test module or create a new test file"
369 );
370 }
371 }
372 }
373}
374struct FunctionVisitor {
375 functions: Vec<FunctionInfo>,
376}
377impl FunctionVisitor {
378 fn new() -> Self {
379 Self { functions: Vec::new() }
380 }
381}
382impl<'ast> Visit<'ast> for FunctionVisitor {
383 fn visit_item_fn(&mut self, node: &'ast syn::ItemFn) {
384 let fn_name = node.sig.ident.to_string();
385 let is_async = node.sig.asyncness.is_some();
386 let params = node
387 .sig
388 .inputs
389 .iter()
390 .filter_map(|arg| {
391 match arg {
392 FnArg::Receiver(_) => None,
393 FnArg::Typed(pat_type) => {
394 if let Pat::Ident(pat_ident) = &*pat_type.pat {
395 let param_name = pat_ident.ident.to_string();
396 let param_type = pat_type.ty.to_token_stream().to_string();
397 let is_reference = param_type.contains('&');
398 Some(ParamInfo {
399 name: param_name,
400 ty: param_type
401 .replace('&', "")
402 .replace("mut", "")
403 .trim()
404 .to_string(),
405 is_reference,
406 })
407 } else {
408 None
409 }
410 }
411 }
412 })
413 .collect();
414 let return_type = match &node.sig.output {
415 ReturnType::Default => None,
416 ReturnType::Type(_, ty) => Some(ty.to_token_stream().to_string()),
417 };
418 let visibility = match &node.vis {
419 syn::Visibility::Public(_) => "public".to_string(),
420 _ => "private".to_string(),
421 };
422 self.functions
423 .push(FunctionInfo {
424 name: fn_name,
425 params,
426 return_type,
427 is_async,
428 visibility,
429 });
430 }
431}
432impl Tool for TestGenTool {
433 fn name(&self) -> &'static str {
434 "test-gen"
435 }
436 fn description(&self) -> &'static str {
437 "Generate test boilerplate from functions"
438 }
439 fn command(&self) -> Command {
440 Command::new(self.name())
441 .about(self.description())
442 .long_about(
443 "Parse Rust AST to find functions and generate test templates with edge cases",
444 )
445 .args(
446 &[
447 Arg::new("file")
448 .long("file")
449 .short('f')
450 .help("Path to Rust source file")
451 .required(true),
452 Arg::new("module")
453 .long("module")
454 .short('m')
455 .help("Module name for generated tests"),
456 Arg::new("type")
457 .long("type")
458 .short('t')
459 .help("Test type to generate")
460 .value_parser(["unit", "integration", "property"])
461 .default_value("unit"),
462 Arg::new("output")
463 .long("output")
464 .short('o')
465 .help("Output file path for generated tests"),
466 ],
467 )
468 .args(&common_options())
469 }
470 fn execute(&self, matches: &ArgMatches) -> Result<()> {
471 let file_path = matches.get_one::<String>("file").unwrap();
472 let test_type = matches.get_one::<String>("type").unwrap();
473 let module_name = matches.get_one::<String>("module");
474 let output_file = matches.get_one::<String>("output");
475 let output_format = parse_output_format(matches);
476 let verbose = matches.get_flag("verbose");
477 let path = Path::new(file_path);
478 if !path.exists() {
479 return Err(
480 ToolError::ExecutionFailed(format!("File not found: {}", file_path)),
481 );
482 }
483 println!("🧪 {} - Generating tests", "CargoMate TestGen".bold().blue());
484 println!(
485 " File: {} | Type: {} | Format: {:?}", file_path, test_type, output_format
486 );
487 let functions = self.parse_rust_file(path)?;
488 if functions.is_empty() {
489 println!("⚠️ No functions found in {}", file_path);
490 return Ok(());
491 }
492 if verbose {
493 println!("📊 Found {} functions", functions.len());
494 }
495 let mut generated_tests = Vec::new();
496 for func in &functions {
497 if verbose {
498 println!(" Generating {} test for: {}", test_type, func.name);
499 }
500 let test = self.generate_test_for_function(func, test_type);
501 generated_tests.push(test);
502 }
503 if let Some(output_path) = output_file {
504 let mut output_content = String::new();
505 if let Some(mod_name) = module_name {
506 output_content
507 .push_str(
508 &format!(
509 "#[cfg(test)]\nmod {} {{\n use super::*;\n\n", mod_name
510 ),
511 );
512 } else {
513 output_content
514 .push_str(
515 "#[cfg(test)]\nmod generated_tests {\n use super::*;\n\n",
516 );
517 }
518 for test in &generated_tests {
519 output_content.push_str(&test.test_code);
520 output_content.push_str("\n");
521 }
522 output_content.push_str("}\n");
523 fs::write(output_path, output_content)
524 .map_err(|e| ToolError::ExecutionFailed(
525 format!("Failed to write output file: {}", e),
526 ))?;
527 println!("💾 Tests saved to: {}", output_path.cyan());
528 } else {
529 self.display_generated_tests(&generated_tests, output_format);
530 }
531 Ok(())
532 }
533}
534impl Default for TestGenTool {
535 fn default() -> Self {
536 Self::new()
537 }
538}