mockforge_bench/
data_driven.rs

1//! Data-driven testing support for load testing
2//!
3//! This module provides functionality to load test data from CSV or JSON files
4//! and generate k6 scripts that use SharedArray for memory-efficient data distribution.
5
6use crate::error::{BenchError, Result};
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9use std::path::Path;
10
11/// Strategy for distributing data across VUs and iterations
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
13#[serde(rename_all = "kebab-case")]
14pub enum DataDistribution {
15    /// Each VU gets a unique row (VU 1 gets row 0, VU 2 gets row 1, etc.)
16    UniquePerVu,
17    /// Each iteration gets a unique row (wraps around when data is exhausted)
18    UniquePerIteration,
19    /// Random row selection on each iteration
20    Random,
21    /// Sequential iteration through all rows (same for all VUs)
22    Sequential,
23}
24
25impl Default for DataDistribution {
26    fn default() -> Self {
27        Self::UniquePerVu
28    }
29}
30
31impl std::str::FromStr for DataDistribution {
32    type Err = BenchError;
33
34    fn from_str(s: &str) -> Result<Self> {
35        match s.to_lowercase().replace('_', "-").as_str() {
36            "unique-per-vu" | "uniquepervu" => Ok(Self::UniquePerVu),
37            "unique-per-iteration" | "uniqueperiteration" => Ok(Self::UniquePerIteration),
38            "random" => Ok(Self::Random),
39            "sequential" => Ok(Self::Sequential),
40            _ => Err(BenchError::Other(format!(
41                "Invalid data distribution: '{}'. Valid options: unique-per-vu, unique-per-iteration, random, sequential",
42                s
43            ))),
44        }
45    }
46}
47
48/// Mapping of data columns to request fields
49#[derive(Debug, Clone, Serialize, Deserialize)]
50pub struct DataMapping {
51    /// Source column name in the data file
52    pub column: String,
53    /// Target field path in the request (e.g., "body.name", "path.id", "header.X-Custom")
54    pub target: String,
55}
56
57impl DataMapping {
58    /// Create a new data mapping
59    pub fn new(column: String, target: String) -> Self {
60        Self { column, target }
61    }
62
63    /// Parse mappings from a comma-separated string
64    /// Format: "column1:target1,column2:target2"
65    pub fn parse_mappings(s: &str) -> Result<Vec<Self>> {
66        if s.is_empty() {
67            return Ok(Vec::new());
68        }
69
70        s.split(',')
71            .map(|pair| {
72                let parts: Vec<&str> = pair.trim().splitn(2, ':').collect();
73                if parts.len() != 2 {
74                    return Err(BenchError::Other(format!(
75                        "Invalid mapping format: '{}'. Expected 'column:target'",
76                        pair
77                    )));
78                }
79                Ok(DataMapping::new(parts[0].trim().to_string(), parts[1].trim().to_string()))
80            })
81            .collect()
82    }
83}
84
85/// Configuration for data-driven testing
86#[derive(Debug, Clone, Serialize, Deserialize)]
87pub struct DataDrivenConfig {
88    /// Path to the data file (CSV or JSON)
89    pub file_path: String,
90    /// Data distribution strategy
91    #[serde(default)]
92    pub distribution: DataDistribution,
93    /// Column to field mappings
94    #[serde(default)]
95    pub mappings: Vec<DataMapping>,
96    /// Whether the CSV has a header row
97    #[serde(default = "default_true")]
98    pub csv_has_header: bool,
99}
100
101fn default_true() -> bool {
102    true
103}
104
105impl DataDrivenConfig {
106    /// Create a new data-driven config
107    pub fn new(file_path: String) -> Self {
108        Self {
109            file_path,
110            distribution: DataDistribution::default(),
111            mappings: Vec::new(),
112            csv_has_header: true,
113        }
114    }
115
116    /// Set the distribution strategy
117    pub fn with_distribution(mut self, distribution: DataDistribution) -> Self {
118        self.distribution = distribution;
119        self
120    }
121
122    /// Add mappings
123    pub fn with_mappings(mut self, mappings: Vec<DataMapping>) -> Self {
124        self.mappings = mappings;
125        self
126    }
127
128    /// Detect file type from extension
129    pub fn file_type(&self) -> DataFileType {
130        if self.file_path.ends_with(".csv") {
131            DataFileType::Csv
132        } else if self.file_path.ends_with(".json") {
133            DataFileType::Json
134        } else {
135            // Default to CSV
136            DataFileType::Csv
137        }
138    }
139}
140
141/// Type of data file
142#[derive(Debug, Clone, Copy, PartialEq, Eq)]
143pub enum DataFileType {
144    Csv,
145    Json,
146}
147
148/// Generates k6 JavaScript code for data-driven testing
149pub struct DataDrivenGenerator;
150
151impl DataDrivenGenerator {
152    /// Generate k6 imports for data-driven testing
153    pub fn generate_imports(file_type: DataFileType) -> String {
154        let mut imports = String::new();
155
156        imports.push_str("import { SharedArray } from 'k6/data';\n");
157
158        if file_type == DataFileType::Csv {
159            imports.push_str(
160                "import papaparse from 'https://jslib.k6.io/papaparse/5.1.1/index.js';\n",
161            );
162        }
163
164        imports
165    }
166
167    /// Generate k6 code to load the data file
168    pub fn generate_data_loading(config: &DataDrivenConfig) -> String {
169        let mut code = String::new();
170
171        code.push_str("// Load test data using SharedArray for memory efficiency\n");
172        code.push_str("const testData = new SharedArray('test data', function() {\n");
173
174        match config.file_type() {
175            DataFileType::Csv => {
176                code.push_str(&format!("  const csvData = open('{}');\n", config.file_path));
177                if config.csv_has_header {
178                    code.push_str("  return papaparse.parse(csvData, { header: true }).data;\n");
179                } else {
180                    code.push_str("  return papaparse.parse(csvData, { header: false }).data;\n");
181                }
182            }
183            DataFileType::Json => {
184                code.push_str(&format!("  return JSON.parse(open('{}'));\n", config.file_path));
185            }
186        }
187
188        code.push_str("});\n\n");
189
190        code
191    }
192
193    /// Generate k6 code for row selection based on distribution strategy
194    pub fn generate_row_selection(distribution: DataDistribution) -> String {
195        match distribution {
196            DataDistribution::UniquePerVu => {
197                "// Unique row per VU (wraps if more VUs than data rows)\n\
198                 const rowIndex = (__VU - 1) % testData.length;\n\
199                 const row = testData[rowIndex];\n"
200                    .to_string()
201            }
202            DataDistribution::UniquePerIteration => {
203                "// Unique row per iteration (cycles through data)\n\
204                 const rowIndex = __ITER % testData.length;\n\
205                 const row = testData[rowIndex];\n"
206                    .to_string()
207            }
208            DataDistribution::Random => "// Random row selection\n\
209                 const rowIndex = Math.floor(Math.random() * testData.length);\n\
210                 const row = testData[rowIndex];\n"
211                .to_string(),
212            DataDistribution::Sequential => {
213                "// Sequential iteration (same for all VUs, based on iteration)\n\
214                 const rowIndex = __ITER % testData.length;\n\
215                 const row = testData[rowIndex];\n"
216                    .to_string()
217            }
218        }
219    }
220
221    /// Generate k6 code to apply mappings from row data
222    pub fn generate_apply_mappings(mappings: &[DataMapping]) -> String {
223        if mappings.is_empty() {
224            return "// No explicit mappings - row data available as 'row' object\n".to_string();
225        }
226
227        let mut code = String::new();
228        code.push_str("// Apply data mappings\n");
229
230        for mapping in mappings {
231            let target_parts: Vec<&str> = mapping.target.splitn(2, '.').collect();
232            if target_parts.len() == 2 {
233                let target_type = target_parts[0];
234                let field_name = target_parts[1];
235
236                match target_type {
237                    "body" => {
238                        code.push_str(&format!(
239                            "requestBody['{}'] = row['{}'];\n",
240                            field_name, mapping.column
241                        ));
242                    }
243                    "path" => {
244                        code.push_str(&format!(
245                            "pathParams['{}'] = row['{}'];\n",
246                            field_name, mapping.column
247                        ));
248                    }
249                    "query" => {
250                        code.push_str(&format!(
251                            "queryParams['{}'] = row['{}'];\n",
252                            field_name, mapping.column
253                        ));
254                    }
255                    "header" => {
256                        code.push_str(&format!(
257                            "requestHeaders['{}'] = row['{}'];\n",
258                            field_name, mapping.column
259                        ));
260                    }
261                    _ => {
262                        code.push_str(&format!(
263                            "// Unknown target type '{}' for column '{}'\n",
264                            target_type, mapping.column
265                        ));
266                    }
267                }
268            } else {
269                // Simple mapping without type prefix - assume body
270                code.push_str(&format!(
271                    "requestBody['{}'] = row['{}'];\n",
272                    mapping.target, mapping.column
273                ));
274            }
275        }
276
277        code
278    }
279
280    /// Generate complete data-driven test setup code
281    pub fn generate_setup(config: &DataDrivenConfig) -> String {
282        let mut code = String::new();
283
284        code.push_str(&Self::generate_imports(config.file_type()));
285        code.push('\n');
286        code.push_str(&Self::generate_data_loading(config));
287
288        code
289    }
290
291    /// Generate code for within the default function
292    pub fn generate_iteration_code(config: &DataDrivenConfig) -> String {
293        let mut code = String::new();
294
295        code.push_str(&Self::generate_row_selection(config.distribution));
296        code.push('\n');
297        code.push_str(&Self::generate_apply_mappings(&config.mappings));
298
299        code
300    }
301}
302
303/// Validate a data file exists and has the expected format
304pub fn validate_data_file(path: &Path) -> Result<DataFileInfo> {
305    if !path.exists() {
306        return Err(BenchError::Other(format!("Data file not found: {}", path.display())));
307    }
308
309    let content = std::fs::read_to_string(path)
310        .map_err(|e| BenchError::Other(format!("Failed to read data file: {}", e)))?;
311
312    let extension = path.extension().and_then(|e| e.to_str()).unwrap_or("");
313
314    match extension {
315        "csv" => validate_csv(&content),
316        "json" => validate_json(&content),
317        _ => Err(BenchError::Other(format!(
318            "Unsupported data file format: .{}. Use .csv or .json",
319            extension
320        ))),
321    }
322}
323
324/// Information about a validated data file
325#[derive(Debug, Clone)]
326pub struct DataFileInfo {
327    /// Number of rows in the file
328    pub row_count: usize,
329    /// Column names (if available)
330    pub columns: Vec<String>,
331    /// File type
332    pub file_type: DataFileType,
333}
334
335fn validate_csv(content: &str) -> Result<DataFileInfo> {
336    let lines: Vec<&str> = content.lines().collect();
337    if lines.is_empty() {
338        return Err(BenchError::Other("CSV file is empty".to_string()));
339    }
340
341    // Assume first line is header
342    let header = lines[0];
343    let columns: Vec<String> = header.split(',').map(|s| s.trim().to_string()).collect();
344    let row_count = lines.len() - 1; // Exclude header
345
346    Ok(DataFileInfo {
347        row_count,
348        columns,
349        file_type: DataFileType::Csv,
350    })
351}
352
353fn validate_json(content: &str) -> Result<DataFileInfo> {
354    let value: serde_json::Value = serde_json::from_str(content)
355        .map_err(|e| BenchError::Other(format!("Invalid JSON: {}", e)))?;
356
357    match value {
358        serde_json::Value::Array(arr) => {
359            if arr.is_empty() {
360                return Err(BenchError::Other("JSON array is empty".to_string()));
361            }
362
363            // Get columns from first object
364            let columns = if let Some(serde_json::Value::Object(obj)) = arr.first() {
365                obj.keys().cloned().collect()
366            } else {
367                Vec::new()
368            };
369
370            Ok(DataFileInfo {
371                row_count: arr.len(),
372                columns,
373                file_type: DataFileType::Json,
374            })
375        }
376        _ => Err(BenchError::Other("JSON data must be an array of objects".to_string())),
377    }
378}
379
380#[cfg(test)]
381mod tests {
382    use super::*;
383    use std::str::FromStr;
384
385    #[test]
386    fn test_data_distribution_default() {
387        assert_eq!(DataDistribution::default(), DataDistribution::UniquePerVu);
388    }
389
390    #[test]
391    fn test_data_distribution_from_str() {
392        assert_eq!(
393            DataDistribution::from_str("unique-per-vu").unwrap(),
394            DataDistribution::UniquePerVu
395        );
396        assert_eq!(
397            DataDistribution::from_str("unique-per-iteration").unwrap(),
398            DataDistribution::UniquePerIteration
399        );
400        assert_eq!(DataDistribution::from_str("random").unwrap(), DataDistribution::Random);
401        assert_eq!(DataDistribution::from_str("sequential").unwrap(), DataDistribution::Sequential);
402    }
403
404    #[test]
405    fn test_data_distribution_from_str_variants() {
406        // Test with underscores
407        assert_eq!(
408            DataDistribution::from_str("unique_per_vu").unwrap(),
409            DataDistribution::UniquePerVu
410        );
411
412        // Test camelCase-ish
413        assert_eq!(
414            DataDistribution::from_str("uniquePerVu").unwrap(),
415            DataDistribution::UniquePerVu
416        );
417    }
418
419    #[test]
420    fn test_data_distribution_from_str_invalid() {
421        assert!(DataDistribution::from_str("invalid").is_err());
422    }
423
424    #[test]
425    fn test_data_mapping_parse() {
426        let mappings = DataMapping::parse_mappings("name:body.username,id:path.userId").unwrap();
427        assert_eq!(mappings.len(), 2);
428        assert_eq!(mappings[0].column, "name");
429        assert_eq!(mappings[0].target, "body.username");
430        assert_eq!(mappings[1].column, "id");
431        assert_eq!(mappings[1].target, "path.userId");
432    }
433
434    #[test]
435    fn test_data_mapping_parse_empty() {
436        let mappings = DataMapping::parse_mappings("").unwrap();
437        assert!(mappings.is_empty());
438    }
439
440    #[test]
441    fn test_data_mapping_parse_invalid() {
442        assert!(DataMapping::parse_mappings("invalid").is_err());
443    }
444
445    #[test]
446    fn test_data_driven_config_file_type() {
447        let csv_config = DataDrivenConfig::new("data.csv".to_string());
448        assert_eq!(csv_config.file_type(), DataFileType::Csv);
449
450        let json_config = DataDrivenConfig::new("data.json".to_string());
451        assert_eq!(json_config.file_type(), DataFileType::Json);
452
453        let unknown_config = DataDrivenConfig::new("data.txt".to_string());
454        assert_eq!(unknown_config.file_type(), DataFileType::Csv); // Default to CSV
455    }
456
457    #[test]
458    fn test_generate_imports_csv() {
459        let imports = DataDrivenGenerator::generate_imports(DataFileType::Csv);
460        assert!(imports.contains("SharedArray"));
461        assert!(imports.contains("papaparse"));
462    }
463
464    #[test]
465    fn test_generate_imports_json() {
466        let imports = DataDrivenGenerator::generate_imports(DataFileType::Json);
467        assert!(imports.contains("SharedArray"));
468        assert!(!imports.contains("papaparse"));
469    }
470
471    #[test]
472    fn test_generate_data_loading_csv() {
473        let config = DataDrivenConfig::new("test.csv".to_string());
474        let code = DataDrivenGenerator::generate_data_loading(&config);
475
476        assert!(code.contains("SharedArray"));
477        assert!(code.contains("open('test.csv')"));
478        assert!(code.contains("papaparse.parse"));
479        assert!(code.contains("header: true"));
480    }
481
482    #[test]
483    fn test_generate_data_loading_json() {
484        let config = DataDrivenConfig::new("test.json".to_string());
485        let code = DataDrivenGenerator::generate_data_loading(&config);
486
487        assert!(code.contains("SharedArray"));
488        assert!(code.contains("open('test.json')"));
489        assert!(code.contains("JSON.parse"));
490    }
491
492    #[test]
493    fn test_generate_row_selection_unique_per_vu() {
494        let code = DataDrivenGenerator::generate_row_selection(DataDistribution::UniquePerVu);
495        assert!(code.contains("__VU - 1"));
496        assert!(code.contains("testData.length"));
497    }
498
499    #[test]
500    fn test_generate_row_selection_unique_per_iteration() {
501        let code =
502            DataDrivenGenerator::generate_row_selection(DataDistribution::UniquePerIteration);
503        assert!(code.contains("__ITER"));
504        assert!(code.contains("testData.length"));
505    }
506
507    #[test]
508    fn test_generate_row_selection_random() {
509        let code = DataDrivenGenerator::generate_row_selection(DataDistribution::Random);
510        assert!(code.contains("Math.random()"));
511        assert!(code.contains("testData.length"));
512    }
513
514    #[test]
515    fn test_generate_apply_mappings() {
516        let mappings = vec![
517            DataMapping::new("name".to_string(), "body.username".to_string()),
518            DataMapping::new("id".to_string(), "path.userId".to_string()),
519            DataMapping::new("token".to_string(), "header.Authorization".to_string()),
520        ];
521
522        let code = DataDrivenGenerator::generate_apply_mappings(&mappings);
523
524        assert!(code.contains("requestBody['username'] = row['name']"));
525        assert!(code.contains("pathParams['userId'] = row['id']"));
526        assert!(code.contains("requestHeaders['Authorization'] = row['token']"));
527    }
528
529    #[test]
530    fn test_generate_apply_mappings_empty() {
531        let code = DataDrivenGenerator::generate_apply_mappings(&[]);
532        assert!(code.contains("No explicit mappings"));
533    }
534
535    #[test]
536    fn test_validate_csv() {
537        let content = "name,email,age\nAlice,alice@test.com,30\nBob,bob@test.com,25";
538        let info = validate_csv(content).unwrap();
539
540        assert_eq!(info.row_count, 2);
541        assert_eq!(info.columns, vec!["name", "email", "age"]);
542        assert_eq!(info.file_type, DataFileType::Csv);
543    }
544
545    #[test]
546    fn test_validate_csv_empty() {
547        let content = "";
548        assert!(validate_csv(content).is_err());
549    }
550
551    #[test]
552    fn test_validate_json() {
553        let content = r#"[{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}]"#;
554        let info = validate_json(content).unwrap();
555
556        assert_eq!(info.row_count, 2);
557        assert!(info.columns.contains(&"name".to_string()));
558        assert!(info.columns.contains(&"age".to_string()));
559        assert_eq!(info.file_type, DataFileType::Json);
560    }
561
562    #[test]
563    fn test_validate_json_empty_array() {
564        let content = "[]";
565        assert!(validate_json(content).is_err());
566    }
567
568    #[test]
569    fn test_validate_json_not_array() {
570        let content = r#"{"name": "Alice"}"#;
571        assert!(validate_json(content).is_err());
572    }
573
574    #[test]
575    fn test_generate_setup() {
576        let config = DataDrivenConfig::new("users.csv".to_string())
577            .with_distribution(DataDistribution::Random);
578
579        let code = DataDrivenGenerator::generate_setup(&config);
580
581        assert!(code.contains("SharedArray"));
582        assert!(code.contains("papaparse"));
583        assert!(code.contains("users.csv"));
584    }
585
586    #[test]
587    fn test_generate_iteration_code() {
588        let config = DataDrivenConfig::new("data.csv".to_string())
589            .with_distribution(DataDistribution::UniquePerVu)
590            .with_mappings(vec![DataMapping::new(
591                "email".to_string(),
592                "body.email".to_string(),
593            )]);
594
595        let code = DataDrivenGenerator::generate_iteration_code(&config);
596
597        assert!(code.contains("__VU - 1"));
598        assert!(code.contains("requestBody['email'] = row['email']"));
599    }
600}
mockforge_bench/data_driven.rs

mockforge_bench/
data_driven.rs
