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

mockforge_bench/
data_driven.rs
