Skip to main content

mockforge_bench/
executor.rs

1//! k6 execution and output handling
2
3use crate::error::{BenchError, Result};
4use indicatif::{ProgressBar, ProgressStyle};
5use std::path::Path;
6use std::process::Stdio;
7use tokio::io::{AsyncBufReadExt, BufReader};
8use tokio::process::Command as TokioCommand;
9
10/// k6 executor
11pub struct K6Executor {
12    k6_path: String,
13}
14
15impl K6Executor {
16    /// Create a new k6 executor
17    pub fn new() -> Result<Self> {
18        let k6_path = which::which("k6")
19            .map_err(|_| BenchError::K6NotFound)?
20            .to_string_lossy()
21            .to_string();
22
23        Ok(Self { k6_path })
24    }
25
26    /// Check if k6 is installed
27    pub fn is_k6_installed() -> bool {
28        which::which("k6").is_ok()
29    }
30
31    /// Get k6 version
32    pub async fn get_version(&self) -> Result<String> {
33        let output = TokioCommand::new(&self.k6_path)
34            .arg("version")
35            .output()
36            .await
37            .map_err(|e| BenchError::K6ExecutionFailed(e.to_string()))?;
38
39        if !output.status.success() {
40            return Err(BenchError::K6ExecutionFailed("Failed to get k6 version".to_string()));
41        }
42
43        Ok(String::from_utf8_lossy(&output.stdout).trim().to_string())
44    }
45
46    /// Execute a k6 script
47    pub async fn execute(
48        &self,
49        script_path: &Path,
50        output_dir: Option<&Path>,
51        verbose: bool,
52    ) -> Result<K6Results> {
53        println!("Starting load test...\n");
54
55        let mut cmd = TokioCommand::new(&self.k6_path);
56        cmd.arg("run");
57
58        // Add output options
59        if let Some(dir) = output_dir {
60            let summary_path = dir.join("summary.json");
61            cmd.arg("--summary-export").arg(summary_path);
62        }
63
64        // Add verbosity
65        if verbose {
66            cmd.arg("--verbose");
67        }
68
69        cmd.arg(script_path);
70        cmd.stdout(Stdio::piped());
71        cmd.stderr(Stdio::piped());
72
73        let mut child = cmd.spawn().map_err(|e| BenchError::K6ExecutionFailed(e.to_string()))?;
74
75        let stdout = child
76            .stdout
77            .take()
78            .ok_or_else(|| BenchError::K6ExecutionFailed("Failed to capture stdout".to_string()))?;
79
80        let stderr = child
81            .stderr
82            .take()
83            .ok_or_else(|| BenchError::K6ExecutionFailed("Failed to capture stderr".to_string()))?;
84
85        // Stream output
86        let stdout_reader = BufReader::new(stdout);
87        let stderr_reader = BufReader::new(stderr);
88
89        let mut stdout_lines = stdout_reader.lines();
90        let mut stderr_lines = stderr_reader.lines();
91
92        // Create progress indicator
93        let spinner = ProgressBar::new_spinner();
94        spinner.set_style(
95            ProgressStyle::default_spinner().template("{spinner:.green} {msg}").unwrap(),
96        );
97        spinner.set_message("Running load test...");
98
99        // Read output lines
100        tokio::spawn(async move {
101            while let Ok(Some(line)) = stdout_lines.next_line().await {
102                spinner.set_message(line.clone());
103                if !line.is_empty() && !line.contains("running") && !line.contains("default") {
104                    println!("{}", line);
105                }
106            }
107            spinner.finish_and_clear();
108        });
109
110        tokio::spawn(async move {
111            while let Ok(Some(line)) = stderr_lines.next_line().await {
112                if !line.is_empty() {
113                    eprintln!("{}", line);
114                }
115            }
116        });
117
118        // Wait for completion
119        let status =
120            child.wait().await.map_err(|e| BenchError::K6ExecutionFailed(e.to_string()))?;
121
122        if !status.success() {
123            return Err(BenchError::K6ExecutionFailed(format!(
124                "k6 exited with status: {}",
125                status
126            )));
127        }
128
129        // Parse results if output directory was specified
130        let results = if let Some(dir) = output_dir {
131            Self::parse_results(dir)?
132        } else {
133            K6Results::default()
134        };
135
136        Ok(results)
137    }
138
139    /// Parse k6 results from JSON output
140    fn parse_results(output_dir: &Path) -> Result<K6Results> {
141        let summary_path = output_dir.join("summary.json");
142
143        if !summary_path.exists() {
144            return Ok(K6Results::default());
145        }
146
147        let content = std::fs::read_to_string(summary_path)
148            .map_err(|e| BenchError::ResultsParseError(e.to_string()))?;
149
150        let json: serde_json::Value = serde_json::from_str(&content)
151            .map_err(|e| BenchError::ResultsParseError(e.to_string()))?;
152
153        let duration_values = &json["metrics"]["http_req_duration"]["values"];
154
155        Ok(K6Results {
156            total_requests: json["metrics"]["http_reqs"]["values"]["count"].as_u64().unwrap_or(0),
157            failed_requests: json["metrics"]["http_req_failed"]["values"]["passes"]
158                .as_u64()
159                .unwrap_or(0),
160            avg_duration_ms: duration_values["avg"].as_f64().unwrap_or(0.0),
161            p95_duration_ms: duration_values["p(95)"].as_f64().unwrap_or(0.0),
162            p99_duration_ms: duration_values["p(99)"].as_f64().unwrap_or(0.0),
163            rps: json["metrics"]["http_reqs"]["values"]["rate"].as_f64().unwrap_or(0.0),
164            vus_max: json["metrics"]["vus_max"]["values"]["value"].as_u64().unwrap_or(0) as u32,
165            min_duration_ms: duration_values["min"].as_f64().unwrap_or(0.0),
166            max_duration_ms: duration_values["max"].as_f64().unwrap_or(0.0),
167            med_duration_ms: duration_values["med"].as_f64().unwrap_or(0.0),
168            p90_duration_ms: duration_values["p(90)"].as_f64().unwrap_or(0.0),
169        })
170    }
171}
172
173impl Default for K6Executor {
174    fn default() -> Self {
175        Self::new().expect("k6 not found")
176    }
177}
178
179/// k6 test results
180#[derive(Debug, Clone, Default)]
181pub struct K6Results {
182    pub total_requests: u64,
183    pub failed_requests: u64,
184    pub avg_duration_ms: f64,
185    pub p95_duration_ms: f64,
186    pub p99_duration_ms: f64,
187    pub rps: f64,
188    pub vus_max: u32,
189    pub min_duration_ms: f64,
190    pub max_duration_ms: f64,
191    pub med_duration_ms: f64,
192    pub p90_duration_ms: f64,
193}
194
195impl K6Results {
196    /// Get error rate as a percentage
197    pub fn error_rate(&self) -> f64 {
198        if self.total_requests == 0 {
199            return 0.0;
200        }
201        (self.failed_requests as f64 / self.total_requests as f64) * 100.0
202    }
203
204    /// Get success rate as a percentage
205    pub fn success_rate(&self) -> f64 {
206        100.0 - self.error_rate()
207    }
208}
209
210#[cfg(test)]
211mod tests {
212    use super::*;
213
214    #[test]
215    fn test_k6_results_error_rate() {
216        let results = K6Results {
217            total_requests: 100,
218            failed_requests: 5,
219            avg_duration_ms: 100.0,
220            p95_duration_ms: 200.0,
221            p99_duration_ms: 300.0,
222            ..Default::default()
223        };
224
225        assert_eq!(results.error_rate(), 5.0);
226        assert_eq!(results.success_rate(), 95.0);
227    }
228
229    #[test]
230    fn test_k6_results_zero_requests() {
231        let results = K6Results::default();
232        assert_eq!(results.error_rate(), 0.0);
233    }
234}