1use crate::chunked_bench::{
23 build_json_body, is_json_content_type, run, ChunkedBenchConfig, ChunkedBenchResult,
24};
25use crate::parallel_executor::ParallelExecutor;
26use crate::request_gen::RequestGenerator;
27use crate::spec_parser::SpecParser;
28use crate::target_parser::{parse_targets_file, TargetConfig};
29use anyhow::{bail, Context};
30use futures::StreamExt;
31use std::collections::HashMap;
32use std::path::{Path, PathBuf};
33use std::sync::Arc;
34use std::time::Duration;
35
36#[derive(Debug, Clone)]
38pub struct ChunkedCommand {
39 pub target: Option<String>,
41 pub targets_file: Option<PathBuf>,
43 pub max_concurrency: u32,
45 pub spec: Option<PathBuf>,
46 pub base_path: Option<String>,
47 pub operation_id: Option<String>,
48 pub method: String,
49 pub concurrency: u32,
50 pub duration: Duration,
51 pub chunk_size_bytes: usize,
52 pub total_size_bytes: usize,
53 pub raw_body: bool,
55 pub chunk_interval_ms: u64,
56 pub headers: HashMap<String, String>,
57 pub insecure: bool,
58 pub validate_requests: bool,
59 pub export_requests: bool,
60 pub output: PathBuf,
61 pub rps: Option<u32>,
63 pub cps: bool,
65 pub rounds: Option<u32>,
67 pub repeat_until: Option<Duration>,
69 pub keep_rounds: Option<u32>,
71}
72
73#[derive(Debug, Clone)]
75struct TargetRun {
76 prefix: String,
78 url: String,
79 headers: HashMap<String, String>,
80 spec: Option<PathBuf>,
81 output: PathBuf,
82}
83
84#[derive(Debug, Clone, Default)]
86struct TargetOutcome {
87 total_requests: u64,
88 successful: u64,
89 failed: u64,
90 bytes_sent: u64,
91 elapsed: Duration,
92 status_counts: HashMap<u16, u64>,
93 run_failed: bool,
95}
96
97impl TargetOutcome {
98 fn absorb(&mut self, r: &ChunkedBenchResult) {
99 self.total_requests += r.total_requests;
100 self.successful += r.successful;
101 self.failed += r.failed;
102 self.bytes_sent += r.bytes_sent;
103 self.elapsed += r.elapsed;
104 for (code, n) in &r.status_counts {
105 *self.status_counts.entry(*code).or_insert(0) += n;
106 }
107 }
108}
109
110impl ChunkedCommand {
111 pub async fn execute(&self) -> anyhow::Result<bool> {
114 match (&self.target, &self.targets_file) {
115 (Some(_), Some(_)) => bail!("--target and --targets-file are mutually exclusive"),
116 (None, None) => bail!("either --target or --targets-file is required"),
117 _ => {}
118 }
119 if self.rounds == Some(0) {
120 bail!("--rounds must be >= 1 (or omit it for a single pass / --repeat-until)");
121 }
122 println!(
123 "Request-start limit per target: {}",
124 self.rps
125 .map(|n| format!("{n} req/s"))
126 .unwrap_or_else(|| "none (--rps not set)".into())
127 );
128 println!(
129 "Connections: {}",
130 if self.cps {
131 "new TCP/TLS connection per request (--cps)"
132 } else {
133 "pooled (--cps not set)"
134 }
135 );
136 if self.raw_body {
137 println!("Request bodies: raw X filler; declared Content-Type preserved (--raw-body)");
138 }
139 let max_rounds = self.rounds.unwrap_or(if self.repeat_until.is_some() {
140 u32::MAX
141 } else {
142 1
143 });
144 let looping = max_rounds > 1 || self.repeat_until.is_some();
145 if !looping {
146 if self.keep_rounds.is_some() {
147 eprintln!(
148 "Note: --keep-rounds only applies to campaign runs (--repeat-until / --rounds > 1); ignoring"
149 );
150 }
151 return Ok(self.execute_pass(&self.output).await?.0);
152 }
153
154 match self.repeat_until {
155 Some(until) => println!(
156 "Campaign: re-running every pass until {}s of wall clock (or --rounds {})",
157 until.as_secs(),
158 self.rounds.map(|n| n.to_string()).unwrap_or_else(|| "unlimited".into())
159 ),
160 None => println!("Campaign: {max_rounds} round(s)"),
161 }
162 if let Some(keep) = self.keep_rounds {
163 println!(
164 "Round pruning: keeping newest {keep} round_* dir(s); per-round stats persist in campaign.jsonl + round-summaries/"
165 );
166 }
167
168 let campaign_start = std::time::Instant::now();
169 let mut all_ok = true;
170 let mut round: u32 = 0;
171 while round < max_rounds {
172 if let Some(until) = self.repeat_until {
173 if round > 0 && campaign_start.elapsed() >= until {
174 println!("Reached --repeat-until; stopping after {round} round(s)");
175 break;
176 }
177 }
178 round += 1;
179 let round_output = self.output.join(format!("round_{round}"));
180 println!(
181 "\n=== Round {round} (campaign elapsed {}s) → {} ===",
182 campaign_start.elapsed().as_secs(),
183 round_output.display()
184 );
185 let round_start = std::time::Instant::now();
186 let (ok, summary) = self.execute_pass(&round_output).await?;
187 all_ok &= ok;
188 write_round_record(
189 &self.output,
190 round,
191 campaign_start.elapsed(),
192 round_start.elapsed(),
193 ok,
194 summary,
195 )?;
196 if let Some(keep) = self.keep_rounds {
197 ParallelExecutor::prune_old_rounds(&self.output, keep);
198 }
199 }
200 Ok(all_ok)
201 }
202
203 async fn execute_pass(&self, output: &Path) -> anyhow::Result<(bool, serde_json::Value)> {
206 match (&self.target, &self.targets_file) {
207 (Some(target), _) => {
208 let run = TargetRun {
209 prefix: String::new(),
210 url: target.clone(),
211 headers: self.headers.clone(),
212 spec: self.spec.clone(),
213 output: output.to_path_buf(),
214 };
215 let o = self.run_target(&run).await?;
216 let summary = serde_json::json!({
217 "url": target,
218 "total_requests": o.total_requests,
219 "successful": o.successful,
220 "failed": o.failed,
221 "bytes_sent": o.bytes_sent,
222 "status_counts": o.status_counts,
223 "run_failed": o.run_failed,
224 "elapsed_seconds": o.elapsed.as_secs_f64(),
225 "measured_rps": measured_rps(o.total_requests, o.elapsed),
226 "rps_per_target": self.rps,
227 "rate_limit_enabled": self.rps.is_some(),
228 "new_connection_per_request": self.cps,
229 "raw_body": self.raw_body,
230 });
231 Ok((!o.run_failed, summary))
232 }
233 (None, Some(file)) => self.execute_multi_target(file, output).await,
234 (None, None) => bail!("either --target or --targets-file is required"),
235 }
236 }
237
238 async fn execute_multi_target(
239 &self,
240 file: &Path,
241 output: &Path,
242 ) -> anyhow::Result<(bool, serde_json::Value)> {
243 let targets = parse_targets_file(file)
244 .map_err(|e| anyhow::anyhow!("{e}"))
245 .with_context(|| format!("failed to read targets file {}", file.display()))?;
246 if targets.is_empty() {
247 bail!("no targets found in {}", file.display());
248 }
249 let runs = plan_target_runs(targets, &self.headers, self.spec.as_ref(), output);
250 let max_concurrency = self.max_concurrency.max(1) as usize;
251 println!(
252 "→ Chunked bench across {} targets from {:?} (max {} in parallel)",
253 runs.len(),
254 file,
255 max_concurrency
256 );
257
258 let mut outcomes: Vec<(usize, TargetRun, anyhow::Result<TargetOutcome>)> =
259 futures::stream::iter(runs.into_iter().enumerate())
260 .map(|(i, run)| async move {
261 let res = self.run_target(&run).await;
262 (i, run, res)
263 })
264 .buffer_unordered(max_concurrency)
265 .collect()
266 .await;
267 outcomes.sort_by_key(|(i, _, _)| *i);
268
269 println!();
270 println!("Multi-target chunked summary:");
271 let mut all_ok = true;
272 let mut rows = Vec::with_capacity(outcomes.len());
273 for (i, run, res) in &outcomes {
274 let name = format!("target_{}", i + 1);
275 match res {
276 Ok(o) => {
277 if o.run_failed {
278 all_ok = false;
279 }
280 println!(
281 " {:<10} {} requests={} ok={} failed={} bytes={}{}",
282 name,
283 run.url,
284 o.total_requests,
285 o.successful,
286 o.failed,
287 o.bytes_sent,
288 if o.run_failed {
289 " (some operations failed to run)"
290 } else {
291 ""
292 }
293 );
294 rows.push(serde_json::json!({
295 "target": name,
296 "url": run.url,
297 "spec": run.spec.as_ref().map(|p| p.display().to_string()),
298 "output_dir": run.spec.as_ref().map(|_| run.output.display().to_string()),
300 "total_requests": o.total_requests,
301 "successful": o.successful,
302 "failed": o.failed,
303 "bytes_sent": o.bytes_sent,
304 "status_counts": o.status_counts,
305 "run_failed": o.run_failed,
306 "elapsed_seconds": o.elapsed.as_secs_f64(),
307 "measured_rps": measured_rps(o.total_requests, o.elapsed),
308 }));
309 }
310 Err(e) => {
311 all_ok = false;
312 println!(" {:<10} {} ERROR: {:#}", name, run.url, e);
313 rows.push(serde_json::json!({
314 "target": name,
315 "url": run.url,
316 "error": format!("{e:#}"),
317 }));
318 }
319 }
320 }
321
322 std::fs::create_dir_all(output)
323 .with_context(|| format!("failed to create output directory {:?}", output))?;
324 let path = output.join("chunked-multi-target-summary.json");
325 let payload = serde_json::json!({
326 "targets_file": file.display().to_string(),
327 "concurrency_per_target": self.concurrency,
328 "max_concurrency": max_concurrency,
329 "duration_secs": self.duration.as_secs(),
330 "chunk_size_bytes": self.chunk_size_bytes,
331 "total_size_bytes": self.total_size_bytes,
332 "chunk_interval_ms": self.chunk_interval_ms,
333 "rps_per_target": self.rps,
334 "rate_limit_enabled": self.rps.is_some(),
335 "new_connection_per_request": self.cps,
336 "raw_body": self.raw_body,
337 "targets": rows,
338 });
339 std::fs::write(&path, serde_json::to_string_pretty(&payload)?)
340 .with_context(|| format!("failed to write {:?}", path))?;
341 println!("📝 Wrote {:?}", path);
342 Ok((all_ok, payload))
343 }
344
345 fn bench_config(
346 &self,
347 url: String,
348 method: reqwest::Method,
349 headers: &HashMap<String, String>,
350 json_seed: Option<&serde_json::Value>,
351 ) -> ChunkedBenchConfig {
352 let body = (!self.raw_body && is_json_content_type(headers))
355 .then(|| Arc::new(build_json_body(json_seed, self.total_size_bytes)));
356 if let Some(b) = &body {
357 if b.len() > self.total_size_bytes {
358 eprintln!(
359 "Note: --total-size-bytes {} is below the smallest valid JSON body; sending {} bytes",
360 self.total_size_bytes,
361 b.len()
362 );
363 }
364 }
365 ChunkedBenchConfig {
366 target_url: url,
367 method,
368 concurrency: self.concurrency,
369 duration: self.duration,
370 chunk_size_bytes: self.chunk_size_bytes,
371 total_size_bytes: self.total_size_bytes,
372 chunk_interval_ms: self.chunk_interval_ms,
373 headers: headers.clone(),
374 skip_tls_verify: self.insecure,
375 rps: self.rps,
376 no_keep_alive: self.cps,
377 body,
378 }
379 }
380
381 async fn run_target(&self, t: &TargetRun) -> anyhow::Result<TargetOutcome> {
382 match &t.spec {
383 Some(spec) => self.run_spec(t, spec).await,
384 None => self.run_single(t).await,
385 }
386 }
387
388 async fn run_single(&self, t: &TargetRun) -> anyhow::Result<TargetOutcome> {
389 let p = &t.prefix;
390 if p.is_empty() {
392 if self.base_path.is_some() {
393 eprintln!(
394 "Note: --base-path has no effect without --spec; \
395 single-target mode uses --target verbatim"
396 );
397 }
398 if self.validate_requests {
399 eprintln!("Note: --validate-requests requires --spec (skipped)");
400 }
401 if self.export_requests {
402 eprintln!("Note: --export-requests requires --spec (skipped)");
403 }
404 }
405 let method = reqwest::Method::from_bytes(self.method.to_uppercase().as_bytes())
406 .map_err(|_| anyhow::anyhow!("Invalid HTTP method: {}", self.method))?;
407 if !p.is_empty() {
408 println!("{p}→ {} {}", self.method.to_uppercase(), t.url);
409 }
410 let r = run(self.bench_config(t.url.clone(), method, &t.headers, None))
411 .await
412 .context("chunked bench failed")?;
413 print_result(
414 p,
415 if p.is_empty() {
416 "single-target"
417 } else {
418 &t.url
419 },
420 &r,
421 );
422 let mut outcome = TargetOutcome::default();
423 outcome.absorb(&r);
424 Ok(outcome)
425 }
426
427 async fn run_spec(&self, t: &TargetRun, spec_path: &Path) -> anyhow::Result<TargetOutcome> {
428 let p = &t.prefix;
429 let parser = SpecParser::from_file(spec_path)
430 .await
431 .map_err(|e| anyhow::anyhow!("Failed to parse spec {:?}: {}", spec_path, e))?;
432 let base_url = t.url.trim_end_matches('/').to_string();
433 let effective_base_path: Option<String> = match &self.base_path {
436 Some(bp) if bp.is_empty() => None,
437 Some(bp) => Some(bp.clone()),
438 None => parser.get_base_path(),
439 };
440 let want_methods = ["POST", "PUT", "PATCH"];
441 let ops: Vec<_> = parser
442 .get_operations()
443 .into_iter()
444 .filter(|op| want_methods.contains(&op.method.to_uppercase().as_str()))
445 .filter(|op| match &self.operation_id {
446 Some(id) => op.operation_id.as_deref() == Some(id.as_str()),
447 None => true,
448 })
449 .collect();
450 if ops.is_empty() {
451 bail!(
452 "No POST/PUT/PATCH operations matched in {:?}{}",
453 spec_path,
454 self.operation_id
455 .as_deref()
456 .map(|id| format!(" (operation-id={id})"))
457 .unwrap_or_default()
458 );
459 }
460
461 let mut violations: Vec<serde_json::Value> = Vec::new();
464 let mut planned: Vec<(String, String, String, String, Option<serde_json::Value>)> =
466 Vec::with_capacity(ops.len());
467 for op in &ops {
468 let label = op.display_name();
469 match RequestGenerator::generate_template(op) {
470 Ok(template) => {
471 let path_with_base = match &effective_base_path {
472 Some(bp) if !bp.is_empty() => {
473 format!("{}{}", bp, template.generate_path())
474 }
475 _ => template.generate_path().to_string(),
476 };
477 let url = format!("{}{}", base_url, path_with_base);
478 if reqwest::Method::from_bytes(op.method.to_uppercase().as_bytes()).is_err() {
479 violations.push(serde_json::json!({
480 "operation": label,
481 "method": op.method,
482 "path": op.path,
483 "kind": "invalid_method",
484 "detail": format!("`{}` is not a valid HTTP method", op.method),
485 }));
486 } else {
487 planned.push((
488 label,
489 op.method.to_uppercase(),
490 url,
491 op.path.clone(),
492 template.body.clone(),
493 ));
494 }
495 }
496 Err(e) => {
497 violations.push(serde_json::json!({
498 "operation": label,
499 "method": op.method,
500 "path": op.path,
501 "kind": "template_build_failure",
502 "detail": e.to_string(),
503 }));
504 }
505 }
506 }
507
508 if self.validate_requests {
509 if !violations.is_empty() {
510 std::fs::create_dir_all(&t.output)
511 .with_context(|| format!("Failed to create output directory {:?}", t.output))?;
512 let path = t.output.join("chunked-request-violations.json");
513 let payload = serde_json::json!({
514 "spec": spec_path.display().to_string(),
515 "base_url": base_url,
516 "base_path": effective_base_path,
517 "violations": violations,
518 });
519 std::fs::write(&path, serde_json::to_string_pretty(&payload)?)
520 .with_context(|| format!("Failed to write {:?}", path))?;
521 bail!("--validate-requests: {} violation(s); see {:?}", violations.len(), path);
522 }
523 println!("{p}✅ --validate-requests: all {} operations OK", planned.len());
524 }
525
526 println!(
527 "{p}→ {} chunked bench {} from {:?} (base {}{})",
528 planned.len(),
529 if planned.len() == 1 {
530 "operation"
531 } else {
532 "operations"
533 },
534 spec_path,
535 base_url,
536 effective_base_path
537 .as_deref()
538 .map(|bp| format!(", base-path {}", bp))
539 .unwrap_or_default()
540 );
541
542 let mut outcome = TargetOutcome::default();
543 let mut export_records: Vec<serde_json::Value> = Vec::new();
544 for (label, method_str, url, op_path, json_seed) in &planned {
545 let method = match reqwest::Method::from_bytes(method_str.as_bytes()) {
546 Ok(m) => m,
547 Err(_) => {
548 eprintln!("{p}✗ {}: invalid method `{}`", label, method_str);
549 outcome.run_failed = true;
550 continue;
551 }
552 };
553 println!();
554 println!("{p}→ {} {} ({})", method_str, op_path, url);
555 let cfg = self.bench_config(url.clone(), method, &t.headers, json_seed.as_ref());
556 match run(cfg).await {
557 Ok(r) => {
558 print_result(p, label, &r);
559 outcome.absorb(&r);
560 if self.export_requests {
561 export_records.push(export_record(
562 self, label, method_str, url, op_path, &t.headers, &r,
563 ));
564 }
565 }
566 Err(e) => {
567 eprintln!("{p}✗ {}: chunked bench failed: {}", label, e);
568 outcome.run_failed = true;
569 }
570 }
571 }
572
573 if self.export_requests {
574 std::fs::create_dir_all(&t.output)
575 .with_context(|| format!("Failed to create output directory {:?}", t.output))?;
576 let path = t.output.join("chunked-requests.json");
577 let payload = serde_json::json!({
578 "spec": spec_path.display().to_string(),
579 "base_url": base_url,
580 "base_path": effective_base_path,
581 "operations": export_records,
582 });
583 std::fs::write(&path, serde_json::to_string_pretty(&payload)?)
584 .with_context(|| format!("Failed to write {:?}", path))?;
585 println!(
586 "{p}📝 --export-requests: wrote {} operations to {:?}",
587 export_records.len(),
588 path
589 );
590 }
591
592 Ok(outcome)
593 }
594}
595
596fn plan_target_runs(
600 targets: Vec<TargetConfig>,
601 cli_headers: &HashMap<String, String>,
602 cli_spec: Option<&PathBuf>,
603 output: &Path,
604) -> Vec<TargetRun> {
605 targets
606 .into_iter()
607 .enumerate()
608 .map(|(i, mut t)| {
609 t.normalize_url();
610 let mut headers = cli_headers.clone();
611 if let Some(h) = &t.headers {
612 headers.extend(h.iter().map(|(k, v)| (k.clone(), v.clone())));
613 }
614 if let Some(auth) = &t.auth {
615 headers.insert("Authorization".to_string(), auth.clone());
616 }
617 let name = format!("target_{}", i + 1);
618 TargetRun {
619 prefix: format!("[{name}] "),
620 url: t.url,
621 headers,
622 spec: t.spec.or_else(|| cli_spec.cloned()),
623 output: output.join(name),
624 }
625 })
626 .collect()
627}
628
629fn write_round_record(
633 output: &Path,
634 round: u32,
635 campaign_elapsed: Duration,
636 round_elapsed: Duration,
637 ok: bool,
638 pass: serde_json::Value,
639) -> anyhow::Result<()> {
640 use std::io::Write;
641 let record = serde_json::json!({
642 "round": round,
643 "campaign_elapsed_seconds": campaign_elapsed.as_secs(),
644 "round_elapsed_seconds": round_elapsed.as_secs_f64(),
645 "all_runs_completed": ok,
646 "pass": pass,
647 });
648 let summaries = output.join("round-summaries");
649 std::fs::create_dir_all(&summaries)
650 .with_context(|| format!("failed to create {:?}", summaries))?;
651 std::fs::write(
652 summaries.join(format!("round_{round}.json")),
653 serde_json::to_string_pretty(&record)?,
654 )?;
655 let mut line = serde_json::to_string(&record)?;
656 line.push('\n');
657 std::fs::OpenOptions::new()
658 .create(true)
659 .append(true)
660 .open(output.join("campaign.jsonl"))?
661 .write_all(line.as_bytes())
662 .context("failed to append campaign.jsonl")?;
663 Ok(())
664}
665
666fn export_record(
667 cmd: &ChunkedCommand,
668 label: &str,
669 method: &str,
670 url: &str,
671 op_path: &str,
672 headers: &HashMap<String, String>,
673 r: &ChunkedBenchResult,
674) -> serde_json::Value {
675 serde_json::json!({
676 "operation": label,
677 "method": method,
678 "url": url,
679 "spec_path": op_path,
680 "headers": headers,
681 "chunk_size_bytes": cmd.chunk_size_bytes,
682 "total_size_bytes": cmd.total_size_bytes,
683 "chunk_interval_ms": cmd.chunk_interval_ms,
684 "concurrency": cmd.concurrency,
685 "rps": cmd.rps,
686 "rate_limit_enabled": cmd.rps.is_some(),
687 "new_connection_per_request": cmd.cps,
688 "raw_body": cmd.raw_body,
689 "duration_secs": cmd.duration.as_secs(),
690 "result": {
691 "total_requests": r.total_requests,
692 "successful": r.successful,
693 "failed": r.failed,
694 "bytes_sent": r.bytes_sent,
695 "elapsed_ms": r.elapsed.as_millis(),
696 "req_per_sec": r.req_per_sec,
697 "avg_latency_ms": r.avg_latency_ms,
698 "p50_ms": r.p50_ms,
699 "p95_ms": r.p95_ms,
700 "p99_ms": r.p99_ms,
701 "status_counts": r.status_counts,
702 "error_samples": r.error_samples.iter().map(|s| serde_json::json!({
703 "status": s.status,
704 "server_header": s.server_header,
705 "body_excerpt": s.body_excerpt,
706 })).collect::<Vec<_>>(),
707 },
708 })
709}
710
711fn measured_rps(requests: u64, elapsed: Duration) -> f64 {
713 if elapsed.is_zero() {
714 0.0
715 } else {
716 requests as f64 / elapsed.as_secs_f64()
717 }
718}
719
720fn print_result(prefix: &str, label: &str, r: &ChunkedBenchResult) {
723 let mut out = String::new();
724 let mut line = |s: String| {
725 out.push_str(prefix);
726 out.push_str(&s);
727 out.push('\n');
728 };
729 line(format!("Chunked bench [{}] complete in {:?}", label, r.elapsed));
730 line(format!(" Total requests: {}", r.total_requests));
731 line(format!(" Successful: {}", r.successful));
732 line(format!(" Failed: {}", r.failed));
733 line(format!(" Bytes sent: {}", r.bytes_sent));
734 line(format!(" Throughput: {:.2} req/s", r.req_per_sec));
735 line(format!(
736 " Latency: avg={:.1}ms p50={}ms p95={}ms p99={}ms",
737 r.avg_latency_ms, r.p50_ms, r.p95_ms, r.p99_ms
738 ));
739 if !r.status_counts.is_empty() {
740 line(" Status codes:".to_string());
741 let mut codes: Vec<_> = r.status_counts.iter().collect();
742 codes.sort_by_key(|(k, _)| **k);
743 for (code, n) in codes {
744 line(format!(" {} = {}", code, n));
745 }
746 }
747 if !r.error_samples.is_empty() {
752 line(" Error response samples:".to_string());
753 for s in &r.error_samples {
754 let server = s.server_header.as_deref().unwrap_or("(no Server header)");
755 line(format!(" [{}] Server: {}", s.status, server));
756 if !s.body_excerpt.is_empty() {
757 line(format!(" body: {}", s.body_excerpt.replace('\n', " ")));
758 }
759 }
760 if r.status_counts.keys().any(|c| (500..600).contains(c)) {
761 line(
762 " Hint: 5xx responses with `Server:` revealing a proxy/LB usually \
763 mean the proxy timed out reading from upstream. Each chunked request \
764 takes >= (total_size_bytes / chunk_size_bytes) * chunk_interval_ms; \
765 if that exceeds the proxy's upstream timeout, errors are inevitable."
766 .to_string(),
767 );
768 }
769 }
770 print!("{out}");
771}
772
773#[cfg(test)]
774mod tests {
775 use super::*;
776
777 #[test]
778 fn plan_target_runs_merges_headers_auth_spec_and_output() {
779 let mut target_headers = HashMap::new();
780 target_headers.insert("X-Env".to_string(), "b".to_string());
781 let targets = vec![
782 TargetConfig::from_url("10.0.0.1:8080".to_string()),
783 TargetConfig {
784 url: "https://b.example".to_string(),
785 auth: Some("Bearer t".to_string()),
786 headers: Some(target_headers),
787 spec: Some(PathBuf::from("b.yaml")),
788 },
789 ];
790 let mut cli_headers = HashMap::new();
791 cli_headers.insert("X-Env".to_string(), "cli".to_string());
792 let cli_spec = PathBuf::from("cli.yaml");
793
794 let runs = plan_target_runs(targets, &cli_headers, Some(&cli_spec), Path::new("out"));
795
796 assert_eq!(runs.len(), 2);
797 assert_eq!(runs[0].url, "http://10.0.0.1:8080");
798 assert_eq!(runs[0].headers.get("X-Env").map(String::as_str), Some("cli"));
799 assert_eq!(runs[0].spec.as_deref(), Some(Path::new("cli.yaml")));
800 assert_eq!(runs[0].output, Path::new("out").join("target_1"));
801 assert_eq!(runs[0].prefix, "[target_1] ");
802
803 assert_eq!(runs[1].headers.get("X-Env").map(String::as_str), Some("b"));
804 assert_eq!(runs[1].headers.get("Authorization").map(String::as_str), Some("Bearer t"));
805 assert_eq!(runs[1].spec.as_deref(), Some(Path::new("b.yaml")));
806 assert_eq!(runs[1].output, Path::new("out").join("target_2"));
807 }
808
809 fn cmd(target: Option<&str>, targets_file: Option<PathBuf>) -> ChunkedCommand {
810 ChunkedCommand {
811 target: target.map(str::to_string),
812 targets_file,
813 max_concurrency: 10,
814 spec: None,
815 base_path: None,
816 operation_id: None,
817 method: "POST".to_string(),
818 concurrency: 1,
819 duration: Duration::from_millis(10),
820 chunk_size_bytes: 1024,
821 total_size_bytes: 4096,
822 raw_body: false,
823 chunk_interval_ms: 0,
824 headers: HashMap::new(),
825 insecure: false,
826 validate_requests: false,
827 export_requests: false,
828 output: PathBuf::from("bench-results"),
829 rps: None,
830 cps: false,
831 rounds: None,
832 repeat_until: None,
833 keep_rounds: None,
834 }
835 }
836
837 #[tokio::test]
838 async fn rejects_target_and_targets_file_together() {
839 let c = cmd(Some("http://127.0.0.1:1"), Some(PathBuf::from("t.txt")));
840 assert!(c.execute().await.is_err());
841 }
842
843 #[tokio::test]
844 async fn rejects_neither_target_nor_targets_file() {
845 assert!(cmd(None, None).execute().await.is_err());
846 }
847
848 fn scratch_dir(name: &str) -> PathBuf {
849 let dir = std::env::temp_dir().join(format!(
850 "mf-chunked-{name}-{}-{}",
851 std::process::id(),
852 line!()
853 ));
854 let _ = std::fs::remove_dir_all(&dir);
855 std::fs::create_dir_all(&dir).unwrap();
856 dir
857 }
858
859 #[tokio::test]
860 async fn campaign_rounds_write_records_and_prune_round_dirs() {
861 let mut server = mockito::Server::new_async().await;
862 let _m = server.mock("POST", "/upload").with_status(200).create_async().await;
863 let dir = scratch_dir("campaign");
864 let targets = dir.join("targets.txt");
865 std::fs::write(&targets, format!("{}/upload\n{}/upload\n", server.url(), server.url()))
866 .unwrap();
867 let mut c = cmd(None, Some(targets));
868 c.output = dir.join("out");
869 c.duration = Duration::from_millis(50);
870 c.rounds = Some(3);
871 c.keep_rounds = Some(1);
872
873 assert!(c.execute().await.unwrap());
874
875 let jsonl = std::fs::read_to_string(c.output.join("campaign.jsonl")).unwrap();
876 let rounds: Vec<serde_json::Value> =
877 jsonl.lines().map(|l| serde_json::from_str(l).unwrap()).collect();
878 assert_eq!(rounds.len(), 3);
879 assert_eq!(rounds[2]["round"], 3);
880 assert_eq!(rounds[2]["pass"]["targets"].as_array().unwrap().len(), 2);
881 for n in 1..=3 {
882 assert!(c.output.join(format!("round-summaries/round_{n}.json")).exists());
883 }
884 assert!(!c.output.join("round_1").exists());
885 assert!(!c.output.join("round_2").exists());
886 assert!(c.output.join("round_3/chunked-multi-target-summary.json").exists());
887 let _ = std::fs::remove_dir_all(&dir);
888 }
889
890 #[tokio::test]
891 async fn campaign_repeat_until_stops_on_wall_clock() {
892 let mut server = mockito::Server::new_async().await;
893 let _m = server.mock("POST", "/upload").with_status(200).create_async().await;
894 let dir = scratch_dir("until");
895 let mut c = cmd(Some(&format!("{}/upload", server.url())), None);
896 c.output = dir.clone();
897 c.duration = Duration::from_millis(50);
901 c.repeat_until = Some(Duration::from_millis(500));
902
903 assert!(c.execute().await.unwrap());
904
905 let n = std::fs::read_to_string(dir.join("campaign.jsonl")).unwrap().lines().count();
906 assert!((2..=11).contains(&n), "expected 2-11 rounds of 50ms in 500ms, got {n}");
907 let _ = std::fs::remove_dir_all(&dir);
908 }
909
910 #[tokio::test]
911 async fn raw_body_preserves_json_headers_in_spec_campaign_and_exports() {
912 let mut server = mockito::Server::new_async().await;
913 let m = server
914 .mock("POST", "/upload")
915 .match_header("content-type", "application/json")
916 .match_header("transfer-encoding", "chunked")
917 .match_body(mockito::Matcher::Exact("X".repeat(4096)))
918 .with_status(400)
919 .expect_at_least(2)
920 .create_async()
921 .await;
922 let dir = tempfile::tempdir().unwrap();
923 let spec = dir.path().join("api.json");
924 std::fs::write(
925 &spec,
926 serde_json::json!({
927 "openapi": "3.0.3", "info": {"title": "test", "version": "1"},
928 "paths": {"/upload": {"post": {
929 "operationId": "upload",
930 "requestBody": {"content": {"application/json": {"schema": {
931 "type": "object", "properties": {"name": {"type": "string"}}
932 }}}},
933 "responses": {"200": {"description": "ok"}}
934 }}}
935 })
936 .to_string(),
937 )
938 .unwrap();
939 let targets = dir.path().join("targets.json");
940 std::fs::write(
941 &targets,
942 serde_json::json!([
943 {"url": server.url(), "headers": {"Content-Type": "application/json"}}
944 ])
945 .to_string(),
946 )
947 .unwrap();
948 let mut c = cmd(None, Some(targets));
949 c.spec = Some(spec);
950 c.raw_body = true;
951 c.output = dir.path().join("out");
952 c.rounds = Some(2);
953 c.rps = Some(10);
954 c.cps = true;
955 c.export_requests = true;
956 assert!(c.execute().await.unwrap());
957 m.assert_async().await;
958 let record: serde_json::Value = serde_json::from_slice(
959 &std::fs::read(c.output.join("round-summaries/round_2.json")).unwrap(),
960 )
961 .unwrap();
962 assert!(record["round_elapsed_seconds"].as_f64().unwrap() > 0.0);
963 let pass = &record["pass"];
964 assert_eq!(pass["raw_body"], true);
965 assert_eq!(pass["rps_per_target"], 10);
966 assert_eq!(pass["rate_limit_enabled"], true);
967 assert_eq!(pass["new_connection_per_request"], true);
968 let row = &pass["targets"][0];
969 let requests = row["total_requests"].as_u64().unwrap();
970 let elapsed = row["elapsed_seconds"].as_f64().unwrap();
971 assert!(requests > 0);
972 let measured = row["measured_rps"].as_f64().unwrap();
974 let expected = requests as f64 / elapsed;
975 assert!((measured - expected).abs() <= expected.abs() * 1e-12);
976 assert_eq!(row["status_counts"]["400"], requests);
977 let exported: serde_json::Value = serde_json::from_slice(
978 &std::fs::read(c.output.join("round_2/target_1/chunked-requests.json")).unwrap(),
979 )
980 .unwrap();
981 assert_eq!(exported["operations"][0]["raw_body"], true);
982 assert_eq!(exported["operations"][0]["headers"]["Content-Type"], "application/json");
983 }
984
985 #[tokio::test]
986 async fn automatic_json_and_raw_body_work_without_spec() {
987 let mut server = mockito::Server::new_async().await;
988 let json = server
989 .mock("POST", "/json")
990 .match_header("content-type", "application/vnd.test+json")
991 .match_header("transfer-encoding", "chunked")
992 .match_body(mockito::Matcher::PartialJson(
993 serde_json::json!({"_padding": "X".repeat(4081)}),
994 ))
995 .with_status(200)
996 .expect_at_least(1)
997 .create_async()
998 .await;
999 let raw = server
1000 .mock("POST", "/raw")
1001 .match_header("content-type", "application/vnd.test+json")
1002 .match_header("transfer-encoding", "chunked")
1003 .match_body(mockito::Matcher::Exact("X".repeat(4096)))
1004 .with_status(400)
1005 .expect_at_least(1)
1006 .create_async()
1007 .await;
1008 let mut c = cmd(Some(&format!("{}/json", server.url())), None);
1009 c.headers.insert("Content-Type".into(), "application/vnd.test+json".into());
1010 let (_, summary) = c.execute_pass(Path::new("unused")).await.unwrap();
1011 assert!(summary["total_requests"].as_u64().unwrap() > 0);
1012 assert_eq!(summary["rps_per_target"], serde_json::Value::Null);
1013 assert_eq!(summary["rate_limit_enabled"], false);
1014 assert_eq!(summary["new_connection_per_request"], false);
1015 assert!(summary["measured_rps"].as_f64().unwrap() > 0.0);
1016 json.assert_async().await;
1017 c.target = Some(format!("{}/raw", server.url()));
1018 c.raw_body = true;
1019 let (_, summary) = c.execute_pass(Path::new("unused")).await.unwrap();
1020 assert!(summary["status_counts"]["400"].as_u64().unwrap() > 0);
1021 raw.assert_async().await;
1022 }
1023
1024 #[test]
1025 fn round_records_distinguish_round_time_from_campaign_time() {
1026 let dir = tempfile::tempdir().unwrap();
1027 write_round_record(
1028 dir.path(),
1029 2,
1030 Duration::from_secs(30),
1031 Duration::from_millis(12500),
1032 true,
1033 serde_json::json!({}),
1034 )
1035 .unwrap();
1036 let record: serde_json::Value = serde_json::from_slice(
1037 &std::fs::read(dir.path().join("round-summaries/round_2.json")).unwrap(),
1038 )
1039 .unwrap();
1040 assert_eq!(record["campaign_elapsed_seconds"], 30);
1041 assert_eq!(record["round_elapsed_seconds"], 12.5);
1042 let jsonl: serde_json::Value = serde_json::from_str(
1043 &std::fs::read_to_string(dir.path().join("campaign.jsonl")).unwrap(),
1044 )
1045 .unwrap();
1046 assert_eq!(jsonl, record);
1047 }
1048
1049 #[tokio::test]
1050 async fn rejects_zero_rounds() {
1051 let mut c = cmd(Some("http://127.0.0.1:1"), None);
1052 c.rounds = Some(0);
1053 assert!(c.execute().await.is_err());
1054 }
1055}