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    /// Enable per-URI control mode (each row specifies method, uri, body, etc.)
100    #[serde(default)]
101    pub per_uri_control: bool,
102    /// Per-URI control column configuration
103    #[serde(default)]
104    pub per_uri_columns: PerUriColumns,
105}
106
107/// Column names for per-URI control mode
108#[derive(Debug, Clone, Serialize, Deserialize)]
109pub struct PerUriColumns {
110    /// Column name for HTTP method (default: "method")
111    #[serde(default = "default_method_column")]
112    pub method: String,
113    /// Column name for URI/path (default: "uri")
114    #[serde(default = "default_uri_column")]
115    pub uri: String,
116    /// Column name for request body (default: "body")
117    #[serde(default = "default_body_column")]
118    pub body: String,
119    /// Column name for query parameters (default: "query_params")
120    #[serde(default = "default_query_params_column")]
121    pub query_params: String,
122    /// Column name for headers (default: "headers")
123    #[serde(default = "default_headers_column")]
124    pub headers: String,
125    /// Column name for attack/security type (default: "attack_type")
126    #[serde(default = "default_attack_type_column")]
127    pub attack_type: String,
128    /// Column name for expected status code (default: "expected_status")
129    #[serde(default = "default_expected_status_column")]
130    pub expected_status: String,
131}
132
133fn default_method_column() -> String {
134    "method".to_string()
135}
136
137fn default_uri_column() -> String {
138    "uri".to_string()
139}
140
141fn default_body_column() -> String {
142    "body".to_string()
143}
144
145fn default_query_params_column() -> String {
146    "query_params".to_string()
147}
148
149fn default_headers_column() -> String {
150    "headers".to_string()
151}
152
153fn default_attack_type_column() -> String {
154    "attack_type".to_string()
155}
156
157fn default_expected_status_column() -> String {
158    "expected_status".to_string()
159}
160
161impl Default for PerUriColumns {
162    fn default() -> Self {
163        Self {
164            method: default_method_column(),
165            uri: default_uri_column(),
166            body: default_body_column(),
167            query_params: default_query_params_column(),
168            headers: default_headers_column(),
169            attack_type: default_attack_type_column(),
170            expected_status: default_expected_status_column(),
171        }
172    }
173}
174
175fn default_true() -> bool {
176    true
177}
178
179impl DataDrivenConfig {
180    /// Create a new data-driven config
181    pub fn new(file_path: String) -> Self {
182        Self {
183            file_path,
184            distribution: DataDistribution::default(),
185            mappings: Vec::new(),
186            csv_has_header: true,
187            per_uri_control: false,
188            per_uri_columns: PerUriColumns::default(),
189        }
190    }
191
192    /// Set the distribution strategy
193    pub fn with_distribution(mut self, distribution: DataDistribution) -> Self {
194        self.distribution = distribution;
195        self
196    }
197
198    /// Add mappings
199    pub fn with_mappings(mut self, mappings: Vec<DataMapping>) -> Self {
200        self.mappings = mappings;
201        self
202    }
203
204    /// Enable per-URI control mode
205    pub fn with_per_uri_control(mut self, enabled: bool) -> Self {
206        self.per_uri_control = enabled;
207        self
208    }
209
210    /// Set custom per-URI column names
211    pub fn with_per_uri_columns(mut self, columns: PerUriColumns) -> Self {
212        self.per_uri_columns = columns;
213        self
214    }
215
216    /// Detect file type from extension
217    pub fn file_type(&self) -> DataFileType {
218        if self.file_path.ends_with(".csv") {
219            DataFileType::Csv
220        } else if self.file_path.ends_with(".json") {
221            DataFileType::Json
222        } else {
223            // Default to CSV
224            DataFileType::Csv
225        }
226    }
227}
228
229/// Type of data file
230#[derive(Debug, Clone, Copy, PartialEq, Eq)]
231pub enum DataFileType {
232    Csv,
233    Json,
234}
235
236/// Generates k6 JavaScript code for data-driven testing
237pub struct DataDrivenGenerator;
238
239impl DataDrivenGenerator {
240    /// Generate k6 imports for data-driven testing
241    pub fn generate_imports(file_type: DataFileType) -> String {
242        let mut imports = String::new();
243
244        imports.push_str("import { SharedArray } from 'k6/data';\n");
245
246        if file_type == DataFileType::Csv {
247            imports.push_str(
248                "import papaparse from 'https://jslib.k6.io/papaparse/5.1.1/index.js';\n",
249            );
250        }
251
252        imports
253    }
254
255    /// Generate k6 code to load the data file
256    pub fn generate_data_loading(config: &DataDrivenConfig) -> String {
257        let mut code = String::new();
258
259        code.push_str("// Load test data using SharedArray for memory efficiency\n");
260        code.push_str("const testData = new SharedArray('test data', function() {\n");
261
262        match config.file_type() {
263            DataFileType::Csv => {
264                code.push_str(&format!("  const csvData = open('{}');\n", config.file_path));
265                if config.csv_has_header {
266                    code.push_str("  return papaparse.parse(csvData, { header: true }).data;\n");
267                } else {
268                    code.push_str("  return papaparse.parse(csvData, { header: false }).data;\n");
269                }
270            }
271            DataFileType::Json => {
272                code.push_str(&format!("  return JSON.parse(open('{}'));\n", config.file_path));
273            }
274        }
275
276        code.push_str("});\n\n");
277
278        code
279    }
280
281    /// Generate k6 code for row selection based on distribution strategy
282    pub fn generate_row_selection(distribution: DataDistribution) -> String {
283        match distribution {
284            DataDistribution::UniquePerVu => {
285                "// Unique row per VU (wraps if more VUs than data rows)\n\
286                 const rowIndex = (__VU - 1) % testData.length;\n\
287                 const row = testData[rowIndex];\n"
288                    .to_string()
289            }
290            DataDistribution::UniquePerIteration => {
291                "// Unique row per iteration (cycles through data)\n\
292                 const rowIndex = __ITER % testData.length;\n\
293                 const row = testData[rowIndex];\n"
294                    .to_string()
295            }
296            DataDistribution::Random => "// Random row selection\n\
297                 const rowIndex = Math.floor(Math.random() * testData.length);\n\
298                 const row = testData[rowIndex];\n"
299                .to_string(),
300            DataDistribution::Sequential => {
301                "// Sequential iteration (same for all VUs, based on iteration)\n\
302                 const rowIndex = __ITER % testData.length;\n\
303                 const row = testData[rowIndex];\n"
304                    .to_string()
305            }
306        }
307    }
308
309    /// Generate k6 code to apply mappings from row data
310    pub fn generate_apply_mappings(mappings: &[DataMapping]) -> String {
311        if mappings.is_empty() {
312            return "// No explicit mappings - row data available as 'row' object\n".to_string();
313        }
314
315        let mut code = String::new();
316        code.push_str("// Apply data mappings\n");
317
318        for mapping in mappings {
319            let target_parts: Vec<&str> = mapping.target.splitn(2, '.').collect();
320            if target_parts.len() == 2 {
321                let target_type = target_parts[0];
322                let field_name = target_parts[1];
323
324                match target_type {
325                    "body" => {
326                        code.push_str(&format!(
327                            "requestBody['{}'] = row['{}'];\n",
328                            field_name, mapping.column
329                        ));
330                    }
331                    "path" => {
332                        code.push_str(&format!(
333                            "pathParams['{}'] = row['{}'];\n",
334                            field_name, mapping.column
335                        ));
336                    }
337                    "query" => {
338                        code.push_str(&format!(
339                            "queryParams['{}'] = row['{}'];\n",
340                            field_name, mapping.column
341                        ));
342                    }
343                    "header" => {
344                        code.push_str(&format!(
345                            "requestHeaders['{}'] = row['{}'];\n",
346                            field_name, mapping.column
347                        ));
348                    }
349                    _ => {
350                        code.push_str(&format!(
351                            "// Unknown target type '{}' for column '{}'\n",
352                            target_type, mapping.column
353                        ));
354                    }
355                }
356            } else {
357                // Simple mapping without type prefix - assume body
358                code.push_str(&format!(
359                    "requestBody['{}'] = row['{}'];\n",
360                    mapping.target, mapping.column
361                ));
362            }
363        }
364
365        code
366    }
367
368    /// Generate complete data-driven test setup code
369    pub fn generate_setup(config: &DataDrivenConfig) -> String {
370        let mut code = String::new();
371
372        code.push_str(&Self::generate_imports(config.file_type()));
373        code.push('\n');
374        code.push_str(&Self::generate_data_loading(config));
375
376        code
377    }
378
379    /// Generate code for within the default function
380    pub fn generate_iteration_code(config: &DataDrivenConfig) -> String {
381        let mut code = String::new();
382
383        code.push_str(&Self::generate_row_selection(config.distribution));
384        code.push('\n');
385
386        if config.per_uri_control {
387            code.push_str(&Self::generate_per_uri_control_code(config));
388        } else {
389            code.push_str(&Self::generate_apply_mappings(&config.mappings));
390        }
391
392        code
393    }
394
395    /// Generate k6 code for per-URI control mode
396    ///
397    /// This mode allows each row in the CSV/JSON to specify:
398    /// - HTTP method (GET, POST, PUT, PATCH, DELETE)
399    /// - URI/path to call
400    /// - Request body (JSON string)
401    /// - Query parameters (JSON string or key=value&key=value format)
402    /// - Additional headers (JSON string)
403    /// - Attack type for security testing
404    /// - Expected status code for validation
405    ///
406    /// Example CSV:
407    /// ```csv
408    /// method,uri,body,query_params,headers,attack_type,expected_status
409    /// GET,/virtualservice,,include_name=true,,sqli,200
410    /// POST,/virtualservice,"{\"name\":\"test-vs\",\"port\":80}",,,,201
411    /// PUT,/virtualservice/{uuid},"{\"name\":\"updated-vs\"}",,,,200
412    /// DELETE,/virtualservice/{uuid},,,,xss,204
413    /// ```
414    pub fn generate_per_uri_control_code(config: &DataDrivenConfig) -> String {
415        let cols = &config.per_uri_columns;
416        let mut code = String::new();
417
418        code.push_str("// Per-URI control mode: each row specifies method, URI, body, etc.\n");
419        code.push_str(&format!(
420            "const method = (row['{}'] || 'GET').toUpperCase();\n",
421            cols.method
422        ));
423        code.push_str(&format!("const uri = row['{}'] || '/';\n", cols.uri));
424        code.push_str(&format!("const bodyStr = row['{}'] || '';\n", cols.body));
425        code.push_str(&format!("const queryParamsStr = row['{}'] || '';\n", cols.query_params));
426        code.push_str(&format!("const extraHeadersStr = row['{}'] || '';\n", cols.headers));
427        code.push_str(&format!("const attackType = row['{}'] || '';\n", cols.attack_type));
428        code.push_str(&format!(
429            "const expectedStatus = row['{}'] ? parseInt(row['{}']) : null;\n",
430            cols.expected_status, cols.expected_status
431        ));
432
433        code.push_str("\n// Parse body if present\n");
434        code.push_str("let requestBody = null;\n");
435        code.push_str("if (bodyStr && bodyStr.trim()) {\n");
436        code.push_str("  try {\n");
437        code.push_str("    requestBody = JSON.parse(bodyStr);\n");
438        code.push_str("  } catch (e) {\n");
439        code.push_str("    // If not valid JSON, use as string (for form data or plain text)\n");
440        code.push_str("    requestBody = bodyStr;\n");
441        code.push_str("  }\n");
442        code.push_str("}\n\n");
443
444        code.push_str("// Parse query parameters if present\n");
445        code.push_str("let queryString = '';\n");
446        code.push_str("if (queryParamsStr && queryParamsStr.trim()) {\n");
447        code.push_str("  try {\n");
448        code.push_str("    // Try parsing as JSON first\n");
449        code.push_str("    const qp = JSON.parse(queryParamsStr);\n");
450        code.push_str("    queryString = '?' + Object.entries(qp).map(([k,v]) => `${encodeURIComponent(k)}=${encodeURIComponent(v)}`).join('&');\n");
451        code.push_str("  } catch (e) {\n");
452        code.push_str("    // Assume it's already in key=value&key=value format\n");
453        code.push_str("    queryString = queryParamsStr.startsWith('?') ? queryParamsStr : '?' + queryParamsStr;\n");
454        code.push_str("  }\n");
455        code.push_str("}\n\n");
456
457        code.push_str("// Parse extra headers if present\n");
458        code.push_str("let extraHeaders = {};\n");
459        code.push_str("if (extraHeadersStr && extraHeadersStr.trim()) {\n");
460        code.push_str("  try {\n");
461        code.push_str("    extraHeaders = JSON.parse(extraHeadersStr);\n");
462        code.push_str("  } catch (e) {\n");
463        code.push_str("    console.warn('Failed to parse extra headers:', e.message);\n");
464        code.push_str("  }\n");
465        code.push_str("}\n\n");
466
467        code.push_str("// Merge headers\n");
468        code.push_str("const mergedHeaders = Object.assign({}, headers, extraHeaders);\n\n");
469
470        code.push_str("// Build full URL with query string\n");
471        code.push_str("const fullUrl = `${BASE_URL}${uri}${queryString}`;\n\n");
472
473        code.push_str("// Make the request based on method\n");
474        code.push_str("let res;\n");
475        code.push_str("switch (method) {\n");
476        code.push_str("  case 'GET':\n");
477        code.push_str("    res = http.get(fullUrl, { headers: mergedHeaders });\n");
478        code.push_str("    break;\n");
479        code.push_str("  case 'POST':\n");
480        code.push_str("    res = http.post(fullUrl, requestBody ? JSON.stringify(requestBody) : null, { headers: mergedHeaders });\n");
481        code.push_str("    break;\n");
482        code.push_str("  case 'PUT':\n");
483        code.push_str("    res = http.put(fullUrl, requestBody ? JSON.stringify(requestBody) : null, { headers: mergedHeaders });\n");
484        code.push_str("    break;\n");
485        code.push_str("  case 'PATCH':\n");
486        code.push_str("    res = http.patch(fullUrl, requestBody ? JSON.stringify(requestBody) : null, { headers: mergedHeaders });\n");
487        code.push_str("    break;\n");
488        code.push_str("  case 'DELETE':\n");
489        code.push_str("    res = http.del(fullUrl, requestBody ? JSON.stringify(requestBody) : null, { headers: mergedHeaders });\n");
490        code.push_str("    break;\n");
491        code.push_str("  default:\n");
492        code.push_str("    console.error(`Unsupported HTTP method: ${method}`);\n");
493        code.push_str("    return;\n");
494        code.push_str("}\n\n");
495
496        code.push_str("// Validate response status if expected status is specified\n");
497        code.push_str("if (expectedStatus !== null) {\n");
498        code.push_str("  check(res, {\n");
499        code.push_str("    [`${method} ${uri}: status is ${expectedStatus}`]: (r) => r.status === expectedStatus,\n");
500        code.push_str("  });\n");
501        code.push_str("} else {\n");
502        code.push_str("  check(res, {\n");
503        code.push_str(
504            "    [`${method} ${uri}: status is 2xx`]: (r) => r.status >= 200 && r.status < 300,\n",
505        );
506        code.push_str("  });\n");
507        code.push_str("}\n\n");
508
509        code.push_str("// Record metrics with operation name\n");
510        code.push_str(
511            "const opName = `${method.toLowerCase()}_${uri.replace(/[^a-zA-Z0-9]/g, '_')}`;\n",
512        );
513        code.push_str("if (typeof perUriLatency !== 'undefined' && perUriLatency[opName]) {\n");
514        code.push_str("  perUriLatency[opName].add(res.timings.duration);\n");
515        code.push_str("}\n\n");
516
517        code.push_str("// Log attack type if security testing\n");
518        code.push_str("if (attackType) {\n");
519        code.push_str(
520            "  console.log(`[Security Test] ${attackType}: ${method} ${uri} => ${res.status}`);\n",
521        );
522        code.push_str("}\n");
523
524        code
525    }
526
527    /// Generate metrics declarations for per-URI control mode
528    pub fn generate_per_uri_metrics(operations: &[(String, String)]) -> String {
529        let mut code = String::new();
530
531        code.push_str("// Per-URI latency metrics\n");
532        code.push_str("const perUriLatency = {\n");
533
534        for (method, uri) in operations {
535            let op_name = format!(
536                "{}_{}",
537                method.to_lowercase(),
538                uri.replace(|c: char| !c.is_alphanumeric(), "_")
539            );
540            code.push_str(&format!("  '{}': new Trend('{}_latency'),\n", op_name, op_name));
541        }
542
543        code.push_str("};\n\n");
544
545        code
546    }
547
548    /// Generate a complete per-URI control mode script
549    pub fn generate_per_uri_control_script(
550        config: &DataDrivenConfig,
551        target_url: &str,
552        custom_headers: &std::collections::HashMap<String, String>,
553        skip_tls_verify: bool,
554    ) -> String {
555        let mut script = String::new();
556
557        // Imports
558        script.push_str("import http from 'k6/http';\n");
559        script.push_str("import { check, sleep } from 'k6';\n");
560        script.push_str("import { Trend, Rate } from 'k6/metrics';\n");
561        script.push_str(&Self::generate_imports(config.file_type()));
562        script.push('\n');
563
564        // Data loading
565        script.push_str(&Self::generate_data_loading(config));
566
567        // Custom metrics
568        script.push_str("// Custom metrics\n");
569        script.push_str("const requestLatency = new Trend('request_latency');\n");
570        script.push_str("const requestErrors = new Rate('request_errors');\n\n");
571
572        // Options
573        script.push_str("export const options = {\n");
574        if skip_tls_verify {
575            script.push_str("  insecureSkipTLSVerify: true,\n");
576        }
577        script.push_str("  scenarios: {\n");
578        script.push_str("    per_uri_control: {\n");
579        script.push_str("      executor: 'shared-iterations',\n");
580        script.push_str("      vus: 10,\n");
581        script.push_str("      iterations: testData.length,\n");
582        script.push_str("      maxDuration: '5m',\n");
583        script.push_str("    },\n");
584        script.push_str("  },\n");
585        script.push_str("  thresholds: {\n");
586        script.push_str("    'http_req_duration': ['p(95)<500'],\n");
587        script.push_str("    'http_req_failed': ['rate<0.1'],\n");
588        script.push_str("  },\n");
589        script.push_str("};\n\n");
590
591        // Base URL and headers
592        script.push_str(&format!("const BASE_URL = '{}';\n\n", target_url));
593        let headers_json =
594            serde_json::to_string(custom_headers).unwrap_or_else(|_| "{}".to_string());
595        script.push_str(&format!("const headers = {};\n\n", headers_json));
596
597        // Default function
598        script.push_str("export default function () {\n");
599        script.push_str("  ");
600        script.push_str(
601            &Self::generate_iteration_code(config).lines().collect::<Vec<_>>().join("\n  "),
602        );
603        script.push_str("\n\n  // Record overall latency\n");
604        script.push_str("  if (res) {\n");
605        script.push_str("    requestLatency.add(res.timings.duration);\n");
606        script.push_str("    requestErrors.add(res.status >= 400);\n");
607        script.push_str("  }\n\n");
608        script.push_str("  sleep(0.1);\n");
609        script.push_str("}\n");
610
611        script
612    }
613}
614
615/// Validate a data file exists and has the expected format
616pub fn validate_data_file(path: &Path) -> Result<DataFileInfo> {
617    if !path.exists() {
618        return Err(BenchError::Other(format!("Data file not found: {}", path.display())));
619    }
620
621    let content = std::fs::read_to_string(path)
622        .map_err(|e| BenchError::Other(format!("Failed to read data file: {}", e)))?;
623
624    let extension = path.extension().and_then(|e| e.to_str()).unwrap_or("");
625
626    match extension {
627        "csv" => validate_csv(&content),
628        "json" => validate_json(&content),
629        _ => Err(BenchError::Other(format!(
630            "Unsupported data file format: .{}. Use .csv or .json",
631            extension
632        ))),
633    }
634}
635
636/// Information about a validated data file
637#[derive(Debug, Clone)]
638pub struct DataFileInfo {
639    /// Number of rows in the file
640    pub row_count: usize,
641    /// Column names (if available)
642    pub columns: Vec<String>,
643    /// File type
644    pub file_type: DataFileType,
645}
646
647fn validate_csv(content: &str) -> Result<DataFileInfo> {
648    let lines: Vec<&str> = content.lines().collect();
649    if lines.is_empty() {
650        return Err(BenchError::Other("CSV file is empty".to_string()));
651    }
652
653    // Assume first line is header
654    let header = lines[0];
655    let columns: Vec<String> = header.split(',').map(|s| s.trim().to_string()).collect();
656    let row_count = lines.len() - 1; // Exclude header
657
658    Ok(DataFileInfo {
659        row_count,
660        columns,
661        file_type: DataFileType::Csv,
662    })
663}
664
665fn validate_json(content: &str) -> Result<DataFileInfo> {
666    let value: serde_json::Value = serde_json::from_str(content)
667        .map_err(|e| BenchError::Other(format!("Invalid JSON: {}", e)))?;
668
669    match value {
670        serde_json::Value::Array(arr) => {
671            if arr.is_empty() {
672                return Err(BenchError::Other("JSON array is empty".to_string()));
673            }
674
675            // Get columns from first object
676            let columns = if let Some(serde_json::Value::Object(obj)) = arr.first() {
677                obj.keys().cloned().collect()
678            } else {
679                Vec::new()
680            };
681
682            Ok(DataFileInfo {
683                row_count: arr.len(),
684                columns,
685                file_type: DataFileType::Json,
686            })
687        }
688        _ => Err(BenchError::Other("JSON data must be an array of objects".to_string())),
689    }
690}
691
692#[cfg(test)]
693mod tests {
694    use super::*;
695    use std::str::FromStr;
696
697    #[test]
698    fn test_data_distribution_default() {
699        assert_eq!(DataDistribution::default(), DataDistribution::UniquePerVu);
700    }
701
702    #[test]
703    fn test_data_distribution_from_str() {
704        assert_eq!(
705            DataDistribution::from_str("unique-per-vu").unwrap(),
706            DataDistribution::UniquePerVu
707        );
708        assert_eq!(
709            DataDistribution::from_str("unique-per-iteration").unwrap(),
710            DataDistribution::UniquePerIteration
711        );
712        assert_eq!(DataDistribution::from_str("random").unwrap(), DataDistribution::Random);
713        assert_eq!(DataDistribution::from_str("sequential").unwrap(), DataDistribution::Sequential);
714    }
715
716    #[test]
717    fn test_data_distribution_from_str_variants() {
718        // Test with underscores
719        assert_eq!(
720            DataDistribution::from_str("unique_per_vu").unwrap(),
721            DataDistribution::UniquePerVu
722        );
723
724        // Test camelCase-ish
725        assert_eq!(
726            DataDistribution::from_str("uniquePerVu").unwrap(),
727            DataDistribution::UniquePerVu
728        );
729    }
730
731    #[test]
732    fn test_data_distribution_from_str_invalid() {
733        assert!(DataDistribution::from_str("invalid").is_err());
734    }
735
736    #[test]
737    fn test_data_mapping_parse() {
738        let mappings = DataMapping::parse_mappings("name:body.username,id:path.userId").unwrap();
739        assert_eq!(mappings.len(), 2);
740        assert_eq!(mappings[0].column, "name");
741        assert_eq!(mappings[0].target, "body.username");
742        assert_eq!(mappings[1].column, "id");
743        assert_eq!(mappings[1].target, "path.userId");
744    }
745
746    #[test]
747    fn test_data_mapping_parse_empty() {
748        let mappings = DataMapping::parse_mappings("").unwrap();
749        assert!(mappings.is_empty());
750    }
751
752    #[test]
753    fn test_data_mapping_parse_invalid() {
754        assert!(DataMapping::parse_mappings("invalid").is_err());
755    }
756
757    #[test]
758    fn test_data_driven_config_file_type() {
759        let csv_config = DataDrivenConfig::new("data.csv".to_string());
760        assert_eq!(csv_config.file_type(), DataFileType::Csv);
761
762        let json_config = DataDrivenConfig::new("data.json".to_string());
763        assert_eq!(json_config.file_type(), DataFileType::Json);
764
765        let unknown_config = DataDrivenConfig::new("data.txt".to_string());
766        assert_eq!(unknown_config.file_type(), DataFileType::Csv); // Default to CSV
767    }
768
769    #[test]
770    fn test_generate_imports_csv() {
771        let imports = DataDrivenGenerator::generate_imports(DataFileType::Csv);
772        assert!(imports.contains("SharedArray"));
773        assert!(imports.contains("papaparse"));
774    }
775
776    #[test]
777    fn test_generate_imports_json() {
778        let imports = DataDrivenGenerator::generate_imports(DataFileType::Json);
779        assert!(imports.contains("SharedArray"));
780        assert!(!imports.contains("papaparse"));
781    }
782
783    #[test]
784    fn test_generate_data_loading_csv() {
785        let config = DataDrivenConfig::new("test.csv".to_string());
786        let code = DataDrivenGenerator::generate_data_loading(&config);
787
788        assert!(code.contains("SharedArray"));
789        assert!(code.contains("open('test.csv')"));
790        assert!(code.contains("papaparse.parse"));
791        assert!(code.contains("header: true"));
792    }
793
794    #[test]
795    fn test_generate_data_loading_json() {
796        let config = DataDrivenConfig::new("test.json".to_string());
797        let code = DataDrivenGenerator::generate_data_loading(&config);
798
799        assert!(code.contains("SharedArray"));
800        assert!(code.contains("open('test.json')"));
801        assert!(code.contains("JSON.parse"));
802    }
803
804    #[test]
805    fn test_generate_row_selection_unique_per_vu() {
806        let code = DataDrivenGenerator::generate_row_selection(DataDistribution::UniquePerVu);
807        assert!(code.contains("__VU - 1"));
808        assert!(code.contains("testData.length"));
809    }
810
811    #[test]
812    fn test_generate_row_selection_unique_per_iteration() {
813        let code =
814            DataDrivenGenerator::generate_row_selection(DataDistribution::UniquePerIteration);
815        assert!(code.contains("__ITER"));
816        assert!(code.contains("testData.length"));
817    }
818
819    #[test]
820    fn test_generate_row_selection_random() {
821        let code = DataDrivenGenerator::generate_row_selection(DataDistribution::Random);
822        assert!(code.contains("Math.random()"));
823        assert!(code.contains("testData.length"));
824    }
825
826    #[test]
827    fn test_generate_apply_mappings() {
828        let mappings = vec![
829            DataMapping::new("name".to_string(), "body.username".to_string()),
830            DataMapping::new("id".to_string(), "path.userId".to_string()),
831            DataMapping::new("token".to_string(), "header.Authorization".to_string()),
832        ];
833
834        let code = DataDrivenGenerator::generate_apply_mappings(&mappings);
835
836        assert!(code.contains("requestBody['username'] = row['name']"));
837        assert!(code.contains("pathParams['userId'] = row['id']"));
838        assert!(code.contains("requestHeaders['Authorization'] = row['token']"));
839    }
840
841    #[test]
842    fn test_generate_apply_mappings_empty() {
843        let code = DataDrivenGenerator::generate_apply_mappings(&[]);
844        assert!(code.contains("No explicit mappings"));
845    }
846
847    #[test]
848    fn test_validate_csv() {
849        let content = "name,email,age\nAlice,alice@test.com,30\nBob,bob@test.com,25";
850        let info = validate_csv(content).unwrap();
851
852        assert_eq!(info.row_count, 2);
853        assert_eq!(info.columns, vec!["name", "email", "age"]);
854        assert_eq!(info.file_type, DataFileType::Csv);
855    }
856
857    #[test]
858    fn test_validate_csv_empty() {
859        let content = "";
860        assert!(validate_csv(content).is_err());
861    }
862
863    #[test]
864    fn test_validate_json() {
865        let content = r#"[{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}]"#;
866        let info = validate_json(content).unwrap();
867
868        assert_eq!(info.row_count, 2);
869        assert!(info.columns.contains(&"name".to_string()));
870        assert!(info.columns.contains(&"age".to_string()));
871        assert_eq!(info.file_type, DataFileType::Json);
872    }
873
874    #[test]
875    fn test_validate_json_empty_array() {
876        let content = "[]";
877        assert!(validate_json(content).is_err());
878    }
879
880    #[test]
881    fn test_validate_json_not_array() {
882        let content = r#"{"name": "Alice"}"#;
883        assert!(validate_json(content).is_err());
884    }
885
886    #[test]
887    fn test_generate_setup() {
888        let config = DataDrivenConfig::new("users.csv".to_string())
889            .with_distribution(DataDistribution::Random);
890
891        let code = DataDrivenGenerator::generate_setup(&config);
892
893        assert!(code.contains("SharedArray"));
894        assert!(code.contains("papaparse"));
895        assert!(code.contains("users.csv"));
896    }
897
898    #[test]
899    fn test_generate_iteration_code() {
900        let config = DataDrivenConfig::new("data.csv".to_string())
901            .with_distribution(DataDistribution::UniquePerVu)
902            .with_mappings(vec![DataMapping::new(
903                "email".to_string(),
904                "body.email".to_string(),
905            )]);
906
907        let code = DataDrivenGenerator::generate_iteration_code(&config);
908
909        assert!(code.contains("__VU - 1"));
910        assert!(code.contains("requestBody['email'] = row['email']"));
911    }
912
913    #[test]
914    fn test_per_uri_columns_default() {
915        let cols = PerUriColumns::default();
916        assert_eq!(cols.method, "method");
917        assert_eq!(cols.uri, "uri");
918        assert_eq!(cols.body, "body");
919        assert_eq!(cols.query_params, "query_params");
920        assert_eq!(cols.headers, "headers");
921        assert_eq!(cols.attack_type, "attack_type");
922        assert_eq!(cols.expected_status, "expected_status");
923    }
924
925    #[test]
926    fn test_data_driven_config_per_uri_control() {
927        let config = DataDrivenConfig::new("test.csv".to_string()).with_per_uri_control(true);
928
929        assert!(config.per_uri_control);
930    }
931
932    #[test]
933    fn test_generate_per_uri_control_code() {
934        let config = DataDrivenConfig::new("test.csv".to_string()).with_per_uri_control(true);
935
936        let code = DataDrivenGenerator::generate_per_uri_control_code(&config);
937
938        // Check for key elements of per-URI control code
939        assert!(code.contains("const method = (row['method']"));
940        assert!(code.contains("const uri = row['uri']"));
941        assert!(code.contains("const bodyStr = row['body']"));
942        assert!(code.contains("const queryParamsStr = row['query_params']"));
943        assert!(code.contains("const attackType = row['attack_type']"));
944        assert!(code.contains("const expectedStatus = row['expected_status']"));
945
946        // Check for HTTP method switch
947        assert!(code.contains("switch (method)"));
948        assert!(code.contains("case 'GET':"));
949        assert!(code.contains("case 'POST':"));
950        assert!(code.contains("case 'PUT':"));
951        assert!(code.contains("case 'PATCH':"));
952        assert!(code.contains("case 'DELETE':"));
953
954        // Check for validation
955        assert!(code.contains("if (expectedStatus !== null)"));
956        assert!(code.contains("check(res"));
957    }
958
959    #[test]
960    fn test_generate_iteration_code_with_per_uri_control() {
961        let config = DataDrivenConfig::new("test.csv".to_string())
962            .with_distribution(DataDistribution::UniquePerIteration)
963            .with_per_uri_control(true);
964
965        let code = DataDrivenGenerator::generate_iteration_code(&config);
966
967        // Should use per-URI control code, not mappings
968        assert!(code.contains("Per-URI control mode"));
969        assert!(code.contains("switch (method)"));
970        assert!(!code.contains("requestBody['"));
971    }
972
973    #[test]
974    fn test_generate_per_uri_metrics() {
975        let operations = vec![
976            ("GET".to_string(), "/users".to_string()),
977            ("POST".to_string(), "/users".to_string()),
978            ("GET".to_string(), "/users/{id}".to_string()),
979        ];
980
981        let code = DataDrivenGenerator::generate_per_uri_metrics(&operations);
982
983        assert!(code.contains("get__users"));
984        assert!(code.contains("post__users"));
985        assert!(code.contains("get__users__id_"));
986        assert!(code.contains("new Trend"));
987    }
988
989    #[test]
990    fn test_generate_per_uri_control_script() {
991        let config = DataDrivenConfig::new("requests.csv".to_string())
992            .with_per_uri_control(true)
993            .with_distribution(DataDistribution::Sequential);
994
995        let headers = std::collections::HashMap::from([(
996            "Content-Type".to_string(),
997            "application/json".to_string(),
998        )]);
999
1000        let script = DataDrivenGenerator::generate_per_uri_control_script(
1001            &config,
1002            "https://api.example.com",
1003            &headers,
1004            true,
1005        );
1006
1007        // Check imports
1008        assert!(script.contains("import http from 'k6/http'"));
1009        assert!(script.contains("import { check, sleep }"));
1010        assert!(script.contains("SharedArray"));
1011
1012        // Check data loading
1013        assert!(script.contains("requests.csv"));
1014
1015        // Check options
1016        assert!(script.contains("insecureSkipTLSVerify: true"));
1017        assert!(script.contains("per_uri_control:"));
1018
1019        // Check base URL
1020        assert!(script.contains("const BASE_URL = 'https://api.example.com'"));
1021
1022        // Check headers
1023        assert!(script.contains("Content-Type"));
1024
1025        // Check default function
1026        assert!(script.contains("export default function"));
1027        assert!(script.contains("switch (method)"));
1028    }
1029}
mockforge_bench/data_driven.rs

mockforge_bench/
data_driven.rs
