1use crate::error::{BenchError, Result};
8use serde::Deserialize;
9use std::path::Path;
10
11#[derive(Debug, Deserialize)]
13pub struct CustomConformanceConfig {
14 pub custom_checks: Vec<CustomCheck>,
16}
17
18#[derive(Debug, Deserialize)]
20pub struct CustomCheck {
21 pub name: String,
23 pub path: String,
25 pub method: String,
27 pub expected_status: u16,
29 #[serde(default)]
31 pub body: Option<String>,
32 #[serde(default)]
34 pub expected_headers: std::collections::HashMap<String, String>,
35 #[serde(default)]
37 pub expected_body_fields: Vec<ExpectedBodyField>,
38 #[serde(default)]
40 pub headers: std::collections::HashMap<String, String>,
41}
42
43#[derive(Debug, Deserialize)]
45pub struct ExpectedBodyField {
46 pub name: String,
48 #[serde(rename = "type")]
50 pub field_type: String,
51}
52
53impl CustomConformanceConfig {
54 pub fn from_file(path: &Path) -> Result<Self> {
56 let content = std::fs::read_to_string(path).map_err(|e| {
57 BenchError::Other(format!(
58 "Failed to read custom conformance file '{}': {}",
59 path.display(),
60 e
61 ))
62 })?;
63 serde_yaml::from_str(&content).map_err(|e| {
64 BenchError::Other(format!(
65 "Failed to parse custom conformance YAML '{}': {}",
66 path.display(),
67 e
68 ))
69 })
70 }
71
72 pub fn generate_k6_group(&self, base_url: &str, custom_headers: &[(String, String)]) -> String {
77 let mut script = String::with_capacity(4096);
78 script.push_str(" group('Custom', function () {\n");
79
80 for check in &self.custom_checks {
81 script.push_str(" {\n");
82
83 let mut all_headers: Vec<(String, String)> = Vec::new();
85 for (k, v) in &check.headers {
87 all_headers.push((k.clone(), v.clone()));
88 }
89 for (k, v) in custom_headers {
91 if !check.headers.contains_key(k) {
92 all_headers.push((k.clone(), v.clone()));
93 }
94 }
95 if check.body.is_some()
97 && !all_headers.iter().any(|(k, _)| k.eq_ignore_ascii_case("content-type"))
98 {
99 all_headers.push(("Content-Type".to_string(), "application/json".to_string()));
100 }
101
102 let headers_js = if all_headers.is_empty() {
103 "{}".to_string()
104 } else {
105 let entries: Vec<String> = all_headers
106 .iter()
107 .map(|(k, v)| format!("'{}': '{}'", k, v.replace('\'', "\\'")))
108 .collect();
109 format!("{{ {} }}", entries.join(", "))
110 };
111
112 let method = check.method.to_uppercase();
113 let url = format!("${{{}}}{}", base_url, check.path);
114 let escaped_name = check.name.replace('\'', "\\'");
115
116 match method.as_str() {
117 "GET" | "HEAD" | "OPTIONS" | "DELETE" => {
118 let k6_method = match method.as_str() {
119 "DELETE" => "del",
120 other => &other.to_lowercase(),
121 };
122 if all_headers.is_empty() {
123 script
124 .push_str(&format!(" let res = http.{}(`{}`);\n", k6_method, url));
125 } else {
126 script.push_str(&format!(
127 " let res = http.{}(`{}`, {{ headers: {} }});\n",
128 k6_method, url, headers_js
129 ));
130 }
131 }
132 _ => {
133 let k6_method = method.to_lowercase();
135 let body_expr = match &check.body {
136 Some(b) => format!("'{}'", b.replace('\'', "\\'")),
137 None => "null".to_string(),
138 };
139 script.push_str(&format!(
140 " let res = http.{}(`{}`, {}, {{ headers: {} }});\n",
141 k6_method, url, body_expr, headers_js
142 ));
143 }
144 }
145
146 script.push_str(&format!(
148 " {{ let ok = check(res, {{ '{}': (r) => r.status === {} }}); if (!ok) __captureFailure('{}', res, 'status === {}'); }}\n",
149 escaped_name, check.expected_status, escaped_name, check.expected_status
150 ));
151
152 for (header_name, pattern) in &check.expected_headers {
154 let header_check_name = format!("{}:header:{}", escaped_name, header_name);
155 let escaped_pattern = pattern.replace('\\', "\\\\").replace('\'', "\\'");
156 script.push_str(&format!(
157 " {{ let ok = check(res, {{ '{}': (r) => new RegExp('{}').test(r.headers['{}'] || r.headers['{}'] || '') }}); if (!ok) __captureFailure('{}', res, 'header {} matches /{}/ '); }}\n",
158 header_check_name,
159 escaped_pattern,
160 header_name,
161 header_name.to_lowercase(),
162 header_check_name,
163 header_name,
164 escaped_pattern
165 ));
166 }
167
168 for field in &check.expected_body_fields {
170 let field_check_name =
171 format!("{}:body:{}:{}", escaped_name, field.name, field.field_type);
172 let accessor = generate_field_accessor(&field.name);
175 let type_check = match field.field_type.as_str() {
176 "string" => format!("typeof ({}) === 'string'", accessor),
177 "integer" => format!("Number.isInteger({})", accessor),
178 "number" => format!("typeof ({}) === 'number'", accessor),
179 "boolean" => format!("typeof ({}) === 'boolean'", accessor),
180 "array" => format!("Array.isArray({})", accessor),
181 "object" => format!(
182 "typeof ({}) === 'object' && !Array.isArray({})",
183 accessor, accessor
184 ),
185 _ => format!("({}) !== undefined", accessor),
186 };
187 script.push_str(&format!(
188 " {{ let ok = check(res, {{ '{}': (r) => {{ try {{ return {}; }} catch(e) {{ return false; }} }} }}); if (!ok) __captureFailure('{}', res, 'body field {} is {}'); }}\n",
189 field_check_name, type_check, field_check_name, field.name, field.field_type
190 ));
191 }
192
193 script.push_str(" }\n");
194 }
195
196 script.push_str(" });\n\n");
197 script
198 }
199}
200
201fn generate_field_accessor(field_name: &str) -> String {
208 let parts: Vec<&str> = field_name.split('.').collect();
210 let mut expr = String::from("JSON.parse(r.body)");
211
212 for part in &parts {
213 if let Some(arr_name) = part.strip_suffix("[]") {
214 expr.push_str(&format!("['{}'][0]", arr_name));
216 } else {
217 expr.push_str(&format!("['{}']", part));
218 }
219 }
220
221 expr
222}
223
224#[cfg(test)]
225mod tests {
226 use super::*;
227
228 #[test]
229 fn test_parse_custom_yaml() {
230 let yaml = r#"
231custom_checks:
232 - name: "custom:pets-returns-200"
233 path: /pets
234 method: GET
235 expected_status: 200
236 - name: "custom:create-product"
237 path: /api/products
238 method: POST
239 expected_status: 201
240 body: '{"sku": "TEST-001", "name": "Test"}'
241 expected_body_fields:
242 - name: id
243 type: integer
244 expected_headers:
245 content-type: "application/json"
246"#;
247 let config: CustomConformanceConfig = serde_yaml::from_str(yaml).unwrap();
248 assert_eq!(config.custom_checks.len(), 2);
249 assert_eq!(config.custom_checks[0].name, "custom:pets-returns-200");
250 assert_eq!(config.custom_checks[0].expected_status, 200);
251 assert_eq!(config.custom_checks[1].expected_body_fields.len(), 1);
252 assert_eq!(config.custom_checks[1].expected_body_fields[0].name, "id");
253 assert_eq!(config.custom_checks[1].expected_body_fields[0].field_type, "integer");
254 }
255
256 #[test]
257 fn test_generate_k6_group_get() {
258 let config = CustomConformanceConfig {
259 custom_checks: vec![CustomCheck {
260 name: "custom:test-get".to_string(),
261 path: "/api/test".to_string(),
262 method: "GET".to_string(),
263 expected_status: 200,
264 body: None,
265 expected_headers: std::collections::HashMap::new(),
266 expected_body_fields: vec![],
267 headers: std::collections::HashMap::new(),
268 }],
269 };
270
271 let script = config.generate_k6_group("BASE_URL", &[]);
272 assert!(script.contains("group('Custom'"));
273 assert!(script.contains("http.get(`${BASE_URL}/api/test`)"));
274 assert!(script.contains("'custom:test-get': (r) => r.status === 200"));
275 }
276
277 #[test]
278 fn test_generate_k6_group_post_with_body() {
279 let config = CustomConformanceConfig {
280 custom_checks: vec![CustomCheck {
281 name: "custom:create".to_string(),
282 path: "/api/items".to_string(),
283 method: "POST".to_string(),
284 expected_status: 201,
285 body: Some(r#"{"name": "test"}"#.to_string()),
286 expected_headers: std::collections::HashMap::new(),
287 expected_body_fields: vec![ExpectedBodyField {
288 name: "id".to_string(),
289 field_type: "integer".to_string(),
290 }],
291 headers: std::collections::HashMap::new(),
292 }],
293 };
294
295 let script = config.generate_k6_group("BASE_URL", &[]);
296 assert!(script.contains("http.post("));
297 assert!(script.contains("'custom:create': (r) => r.status === 201"));
298 assert!(script.contains("custom:create:body:id:integer"));
299 assert!(script.contains("Number.isInteger"));
300 }
301
302 #[test]
303 fn test_generate_k6_group_with_header_checks() {
304 let mut expected_headers = std::collections::HashMap::new();
305 expected_headers.insert("content-type".to_string(), "application/json".to_string());
306
307 let config = CustomConformanceConfig {
308 custom_checks: vec![CustomCheck {
309 name: "custom:header-check".to_string(),
310 path: "/api/test".to_string(),
311 method: "GET".to_string(),
312 expected_status: 200,
313 body: None,
314 expected_headers,
315 expected_body_fields: vec![],
316 headers: std::collections::HashMap::new(),
317 }],
318 };
319
320 let script = config.generate_k6_group("BASE_URL", &[]);
321 assert!(script.contains("custom:header-check:header:content-type"));
322 assert!(script.contains("new RegExp('application/json')"));
323 }
324
325 #[test]
326 fn test_generate_k6_group_with_custom_headers() {
327 let config = CustomConformanceConfig {
328 custom_checks: vec![CustomCheck {
329 name: "custom:auth-test".to_string(),
330 path: "/api/secure".to_string(),
331 method: "GET".to_string(),
332 expected_status: 200,
333 body: None,
334 expected_headers: std::collections::HashMap::new(),
335 expected_body_fields: vec![],
336 headers: std::collections::HashMap::new(),
337 }],
338 };
339
340 let custom_headers = vec![("Authorization".to_string(), "Bearer token123".to_string())];
341 let script = config.generate_k6_group("BASE_URL", &custom_headers);
342 assert!(script.contains("'Authorization': 'Bearer token123'"));
343 }
344
345 #[test]
346 fn test_failure_capture_emitted() {
347 let config = CustomConformanceConfig {
348 custom_checks: vec![CustomCheck {
349 name: "custom:capture-test".to_string(),
350 path: "/api/test".to_string(),
351 method: "GET".to_string(),
352 expected_status: 200,
353 body: None,
354 expected_headers: {
355 let mut m = std::collections::HashMap::new();
356 m.insert("X-Rate-Limit".to_string(), ".*".to_string());
357 m
358 },
359 expected_body_fields: vec![ExpectedBodyField {
360 name: "id".to_string(),
361 field_type: "integer".to_string(),
362 }],
363 headers: std::collections::HashMap::new(),
364 }],
365 };
366
367 let script = config.generate_k6_group("BASE_URL", &[]);
368 assert!(
370 script.contains("__captureFailure('custom:capture-test', res, 'status === 200')"),
371 "Status check should emit __captureFailure"
372 );
373 assert!(
375 script.contains("__captureFailure('custom:capture-test:header:X-Rate-Limit'"),
376 "Header check should emit __captureFailure"
377 );
378 assert!(
380 script.contains("__captureFailure('custom:capture-test:body:id:integer'"),
381 "Body field check should emit __captureFailure"
382 );
383 }
384
385 #[test]
386 fn test_from_file_nonexistent() {
387 let result = CustomConformanceConfig::from_file(Path::new("/nonexistent/file.yaml"));
388 assert!(result.is_err());
389 let err = result.unwrap_err().to_string();
390 assert!(err.contains("Failed to read custom conformance file"));
391 }
392
393 #[test]
394 fn test_generate_k6_group_delete() {
395 let config = CustomConformanceConfig {
396 custom_checks: vec![CustomCheck {
397 name: "custom:delete-item".to_string(),
398 path: "/api/items/1".to_string(),
399 method: "DELETE".to_string(),
400 expected_status: 204,
401 body: None,
402 expected_headers: std::collections::HashMap::new(),
403 expected_body_fields: vec![],
404 headers: std::collections::HashMap::new(),
405 }],
406 };
407
408 let script = config.generate_k6_group("BASE_URL", &[]);
409 assert!(script.contains("http.del("));
410 assert!(script.contains("r.status === 204"));
411 }
412
413 #[test]
414 fn test_field_accessor_simple() {
415 assert_eq!(generate_field_accessor("name"), "JSON.parse(r.body)['name']");
416 }
417
418 #[test]
419 fn test_field_accessor_nested_dot() {
420 assert_eq!(
421 generate_field_accessor("config.enabled"),
422 "JSON.parse(r.body)['config']['enabled']"
423 );
424 }
425
426 #[test]
427 fn test_field_accessor_array_bracket() {
428 assert_eq!(generate_field_accessor("items[].id"), "JSON.parse(r.body)['items'][0]['id']");
429 }
430
431 #[test]
432 fn test_field_accessor_deep_nested() {
433 assert_eq!(generate_field_accessor("a.b.c"), "JSON.parse(r.body)['a']['b']['c']");
434 }
435
436 #[test]
437 fn test_generate_k6_nested_body_fields() {
438 let config = CustomConformanceConfig {
439 custom_checks: vec![CustomCheck {
440 name: "custom:nested".to_string(),
441 path: "/api/data".to_string(),
442 method: "GET".to_string(),
443 expected_status: 200,
444 body: None,
445 expected_headers: std::collections::HashMap::new(),
446 expected_body_fields: vec![
447 ExpectedBodyField {
448 name: "count".to_string(),
449 field_type: "integer".to_string(),
450 },
451 ExpectedBodyField {
452 name: "results[].name".to_string(),
453 field_type: "string".to_string(),
454 },
455 ],
456 headers: std::collections::HashMap::new(),
457 }],
458 };
459
460 let script = config.generate_k6_group("BASE_URL", &[]);
461 assert!(script.contains("JSON.parse(r.body)['count']"));
463 assert!(script.contains("JSON.parse(r.body)['results'][0]['name']"));
465 }
466}