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

mockforge_bench/
data_driven.rs
