use crate::common::{ExponentialBackoff, TestRepo, wt_bin};
use portable_pty::CommandBuilder;
use std::io::Read;
use std::time::{Duration, Instant};
const OUTPUT_POLL_INTERVAL_MS: u64 = 10;
const STABLE_READ_THRESHOLD: u32 = 4;
const MIN_DRAIN_WAIT_MS: u64 = 1000;
#[derive(Debug, Clone)]
pub enum CaptureStrategy {
TimeInterval(Duration),
ByteInterval(usize),
}
#[derive(Debug, Clone)]
pub struct ProgressiveCaptureOptions {
pub strategy: CaptureStrategy,
pub timeout: Duration,
pub terminal_size: (u16, u16),
pub min_change_threshold: usize,
}
impl Default for ProgressiveCaptureOptions {
fn default() -> Self {
Self {
strategy: CaptureStrategy::TimeInterval(Duration::from_millis(50)),
timeout: Duration::from_secs(10),
terminal_size: (48, 150),
min_change_threshold: 100, }
}
}
impl ProgressiveCaptureOptions {
pub fn with_byte_interval(byte_interval: usize) -> Self {
Self {
strategy: CaptureStrategy::ByteInterval(byte_interval),
timeout: Duration::from_secs(10),
terminal_size: (48, 150),
min_change_threshold: 100,
}
}
}
#[derive(Debug, Clone)]
pub struct OutputSnapshot {
pub timestamp: Duration,
visible_text: String,
}
impl OutputSnapshot {
pub fn visible_text(&self) -> &str {
&self.visible_text
}
}
#[derive(Debug)]
pub struct ProgressiveOutput {
pub stages: Vec<OutputSnapshot>,
pub exit_code: i32,
pub total_duration: Duration,
}
impl ProgressiveOutput {
pub fn initial(&self) -> &OutputSnapshot {
self.stages.first().unwrap()
}
pub fn final_snapshot(&self) -> &OutputSnapshot {
self.stages.last().unwrap()
}
pub fn final_output(&self) -> &str {
self.final_snapshot().visible_text()
}
pub fn snapshot_at(&self, target: Duration) -> &OutputSnapshot {
self.stages
.iter()
.min_by_key(|s| s.timestamp.abs_diff(target))
.unwrap()
}
pub fn samples(&self, count: usize) -> Vec<&OutputSnapshot> {
if self.stages.is_empty() || count == 0 {
return vec![];
}
if count >= self.stages.len() {
return self.stages.iter().collect();
}
let step = (self.stages.len() - 1) as f64 / (count - 1) as f64;
(0..count)
.map(|i| {
let index = ((i as f64 * step).round() as usize).min(self.stages.len() - 1);
&self.stages[index]
})
.collect()
}
pub fn dots_per_stage(&self) -> Vec<usize> {
self.stages
.iter()
.map(|s| s.visible_text.matches('·').count())
.collect()
}
pub fn verify_progressive_filling(&self) -> Result<(), String> {
let dots = self.dots_per_stage();
if dots.is_empty() {
return Err("No snapshots captured".to_string());
}
if dots.len() < 2 {
return Err("Need at least 2 snapshots to verify progressive filling".to_string());
}
let (peak_index, peak_dots) = dots
.iter()
.enumerate()
.max_by_key(|(_, count)| *count)
.unwrap();
let peak_dots = *peak_dots;
if peak_dots == 0 {
return Err(format!(
"Progressive filling verification failed: no placeholder dots observed in any snapshot. \
This suggests data filled too quickly to observe progressive rendering, or the command \
doesn't use progressive rendering. Dots progression: {:?}",
dots
));
}
if peak_index == dots.len() - 1 {
return Err(format!(
"Progressive filling verification failed: peak dots ({}) at final snapshot. \
Expected dots to decrease after peak. Dots progression: {:?}",
peak_dots, dots
));
}
let final_dots = dots[dots.len() - 1];
if final_dots >= peak_dots {
return Err(format!(
"Progressive filling verification failed: final dots ({}) >= peak dots ({}). \
Expected decrease after peak. Dots progression: {:?}",
final_dots, peak_dots, dots
));
}
Ok(())
}
}
pub fn capture_progressive_output(
repo: &TestRepo,
subcommand: &str,
args: &[&str],
options: ProgressiveCaptureOptions,
) -> ProgressiveOutput {
let start_time = Instant::now();
let pair = super::open_pty_with_size(options.terminal_size.0, options.terminal_size.1);
let mut cmd = CommandBuilder::new(wt_bin());
cmd.arg(subcommand);
for arg in args {
cmd.arg(arg);
}
cmd.cwd(repo.root_path());
configure_pty_environment(&mut cmd, repo);
let mut child = pair.slave.spawn_command(cmd).unwrap_or_else(|_| {
panic!(
"Failed to spawn 'wt {}' in PTY at {}",
subcommand,
repo.root_path().display()
)
});
drop(pair.slave);
let mut parser = vt100::Parser::new(
options.terminal_size.0,
options.terminal_size.1,
0, );
let mut reader = pair.master.try_clone_reader().unwrap();
let mut snapshots = Vec::new();
let mut last_snapshot_time = Instant::now();
let mut last_snapshot_bytes = 0;
let mut last_snapshot_text = String::new();
let mut total_bytes = 0;
loop {
let mut temp_buf = [0u8; 4096];
let drain_and_capture_final = |parser: &mut vt100::Parser,
reader: &mut Box<dyn std::io::Read + Send>,
snapshots: &mut Vec<OutputSnapshot>,
last_snapshot_text: &str,
total_bytes: &mut usize,
start_time: Instant| {
let backoff = ExponentialBackoff::default();
let mut attempt = 0u32;
let mut consecutive_no_data = 0u32;
let drain_start = Instant::now();
loop {
if drain_start.elapsed() > backoff.timeout {
break;
}
let mut buf = [0u8; 4096];
match reader.read(&mut buf) {
Ok(0) => {
consecutive_no_data += 1;
let min_wait_elapsed =
drain_start.elapsed() >= Duration::from_millis(MIN_DRAIN_WAIT_MS);
if consecutive_no_data >= STABLE_READ_THRESHOLD && min_wait_elapsed {
break;
}
backoff.sleep(attempt);
attempt += 1;
}
Ok(n) => {
consecutive_no_data = 0;
attempt = 0;
*total_bytes += n;
parser.process(&buf[..n]);
}
Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
backoff.sleep(attempt);
attempt += 1;
}
Err(_) => break,
}
}
let screen = parser.screen();
let final_text = screen.contents();
if final_text != last_snapshot_text {
snapshots.push(OutputSnapshot {
timestamp: start_time.elapsed(),
visible_text: final_text,
});
}
};
match reader.read(&mut temp_buf) {
Ok(0) => {
drain_and_capture_final(
&mut parser,
&mut reader,
&mut snapshots,
&last_snapshot_text,
&mut total_bytes,
start_time,
);
break;
}
Ok(n) => {
total_bytes += n;
parser.process(&temp_buf[..n]);
let should_snapshot = match options.strategy {
CaptureStrategy::TimeInterval(interval) => {
last_snapshot_time.elapsed() >= interval
}
CaptureStrategy::ByteInterval(byte_interval) => {
total_bytes >= last_snapshot_bytes + byte_interval
}
};
if should_snapshot {
let screen = parser.screen();
let current_text = screen.contents();
let meets_threshold = match options.strategy {
CaptureStrategy::TimeInterval(_) => {
last_snapshot_time.elapsed().as_millis()
>= options.min_change_threshold as u128
}
CaptureStrategy::ByteInterval(_) => {
total_bytes >= last_snapshot_bytes + options.min_change_threshold
}
};
if current_text != last_snapshot_text && meets_threshold {
snapshots.push(OutputSnapshot {
timestamp: start_time.elapsed(),
visible_text: current_text.clone(),
});
last_snapshot_text = current_text;
last_snapshot_time = Instant::now();
last_snapshot_bytes = total_bytes;
}
}
}
Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
if let Ok(Some(_)) = child.try_wait() {
drain_and_capture_final(
&mut parser,
&mut reader,
&mut snapshots,
&last_snapshot_text,
&mut total_bytes,
start_time,
);
break;
}
if start_time.elapsed() > options.timeout {
panic!(
"Timeout waiting for command to complete after {:?}. \
Captured {} snapshots, last at {:?}, {} total bytes",
options.timeout,
snapshots.len(),
snapshots.last().map(|s| s.timestamp),
total_bytes
);
}
std::thread::sleep(Duration::from_millis(OUTPUT_POLL_INTERVAL_MS));
}
Err(e) => panic!("Failed to read PTY output: {}", e),
}
}
let exit_status = child
.wait()
.unwrap_or_else(|_| panic!("Failed to wait for 'wt {}' to exit", subcommand));
let exit_code = exit_status.exit_code() as i32;
let total_duration = start_time.elapsed();
ProgressiveOutput {
stages: snapshots,
exit_code,
total_duration,
}
}
fn configure_pty_environment(cmd: &mut CommandBuilder, repo: &TestRepo) {
cmd.env_clear();
cmd.env(
"HOME",
home::home_dir().unwrap().to_string_lossy().to_string(),
);
cmd.env(
"PATH",
std::env::var("PATH").unwrap_or_else(|_| "/usr/bin:/bin".to_string()),
);
for (key, value) in repo.test_env_vars() {
cmd.env(key, value);
}
cmd.env("WORKTRUNK_PLACEHOLDER_REVEAL_MS", "0");
for key in [
"LLVM_PROFILE_FILE",
"CARGO_LLVM_COV",
"CARGO_LLVM_COV_TARGET_DIR",
] {
if let Ok(val) = std::env::var(key) {
cmd.env(key, val);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_progressive_capture_options_default() {
let opts = ProgressiveCaptureOptions::default();
assert!(matches!(opts.strategy, CaptureStrategy::TimeInterval(_)));
assert_eq!(opts.timeout, Duration::from_secs(10));
assert_eq!(opts.terminal_size, (48, 150));
assert_eq!(opts.min_change_threshold, 100);
}
#[test]
fn test_progressive_capture_options_byte_interval() {
let opts = ProgressiveCaptureOptions::with_byte_interval(500);
assert!(matches!(opts.strategy, CaptureStrategy::ByteInterval(500)));
assert_eq!(opts.timeout, Duration::from_secs(10));
assert_eq!(opts.terminal_size, (48, 150));
assert_eq!(opts.min_change_threshold, 100);
}
#[test]
fn test_progressive_output_samples() {
let stages = vec![
OutputSnapshot {
timestamp: Duration::from_millis(0),
visible_text: "stage 0".to_string(),
},
OutputSnapshot {
timestamp: Duration::from_millis(50),
visible_text: "stage 1".to_string(),
},
OutputSnapshot {
timestamp: Duration::from_millis(100),
visible_text: "stage 2".to_string(),
},
OutputSnapshot {
timestamp: Duration::from_millis(150),
visible_text: "stage 3".to_string(),
},
];
let output = ProgressiveOutput {
stages,
exit_code: 0,
total_duration: Duration::from_millis(150),
};
let samples = output.samples(2);
assert_eq!(samples.len(), 2);
assert_eq!(samples[0].visible_text, "stage 0");
assert_eq!(samples[1].visible_text, "stage 3");
let samples = output.samples(3);
assert_eq!(samples.len(), 3);
assert_eq!(samples[0].visible_text, "stage 0");
assert_eq!(samples[1].visible_text, "stage 2");
assert_eq!(samples[2].visible_text, "stage 3");
}
#[test]
fn test_dots_decrease_verification() {
let stages = vec![
OutputSnapshot {
timestamp: Duration::from_millis(0),
visible_text: "· · · · ·".to_string(),
},
OutputSnapshot {
timestamp: Duration::from_millis(50),
visible_text: "data · · ·".to_string(),
},
OutputSnapshot {
timestamp: Duration::from_millis(100),
visible_text: "data data ·".to_string(),
},
OutputSnapshot {
timestamp: Duration::from_millis(150),
visible_text: "data data data".to_string(),
},
];
let output = ProgressiveOutput {
stages,
exit_code: 0,
total_duration: Duration::from_millis(150),
};
assert!(output.verify_progressive_filling().is_ok());
let dots = output.dots_per_stage();
assert_eq!(dots, vec![5, 3, 1, 0]);
}
}