use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use std::sync::mpsc;
use std::thread;
use std::time::{Duration, Instant};
use portable_pty::{CommandBuilder, NativePtySystem, PtySize, PtySystem};
use crate::assertion::AssertionFailure;
use crate::frame::TerminalFrame;
use crate::scenario;
use crate::step::Step;
const DEFAULT_COLS: u16 = 80;
const DEFAULT_ROWS: u16 = 24;
const DEFAULT_READ_TIMEOUT: Duration = Duration::from_millis(500);
pub struct PtySession {
binary_path: PathBuf,
cols: u16,
rows: u16,
output_buffer: Vec<u8>,
writer: Box<dyn Write + Send>,
output_receiver: mpsc::Receiver<Vec<u8>>,
child: Box<dyn portable_pty::Child + Send + Sync>,
}
impl PtySession {
pub fn spawn(binary_path: &Path) -> Result<Self, PtySessionError> {
Self::spawn_with_size(binary_path, DEFAULT_COLS, DEFAULT_ROWS, &[], &[], None)
}
pub fn spawn_with_size(
binary_path: &Path,
cols: u16,
rows: u16,
args: &[&str],
env_vars: &[(&str, &str)],
workdir: Option<&Path>,
) -> Result<Self, PtySessionError> {
let pty_system = NativePtySystem::default();
let pair = pty_system
.openpty(PtySize {
rows,
cols,
pixel_width: 0,
pixel_height: 0,
})
.map_err(|err| PtySessionError::PtyCreation(err.to_string()))?;
let mut command = CommandBuilder::new(binary_path);
if let Some(directory) = workdir {
command.cwd(directory);
}
for arg in args {
command.arg(arg);
}
for (key, value) in env_vars {
command.env(key, value);
}
let child = pair
.slave
.spawn_command(command)
.map_err(|err| PtySessionError::SpawnFailed(err.to_string()))?;
let reader = pair
.master
.try_clone_reader()
.map_err(|err| PtySessionError::PtyCreation(err.to_string()))?;
let writer = pair
.master
.take_writer()
.map_err(|err| PtySessionError::PtyCreation(err.to_string()))?;
let (output_sender, output_receiver) = mpsc::channel();
thread::spawn(move || read_pty_output(reader, &output_sender));
Ok(Self {
binary_path: binary_path.to_path_buf(),
cols,
rows,
output_buffer: Vec::new(),
writer,
output_receiver,
child,
})
}
pub fn write_bytes(&mut self, data: &[u8]) -> Result<(), PtySessionError> {
self.writer
.write_all(data)
.map_err(|err| PtySessionError::WriteFailed(err.to_string()))?;
self.writer
.flush()
.map_err(|err| PtySessionError::WriteFailed(err.to_string()))?;
Ok(())
}
pub fn write_text(&mut self, text: &str) -> Result<(), PtySessionError> {
self.write_bytes(text.as_bytes())
}
pub fn press_key(&mut self, key: &str) -> Result<(), PtySessionError> {
let bytes = key_to_bytes(key);
self.write_bytes(&bytes)
}
pub fn drain_output(&mut self, timeout: Duration) {
let deadline = std::time::Instant::now() + timeout;
while let Ok(chunk) = self.output_receiver.recv_timeout(
deadline
.checked_duration_since(std::time::Instant::now())
.unwrap_or(Duration::ZERO),
) {
self.output_buffer.extend_from_slice(&chunk);
}
}
pub fn capture_frame(&mut self) -> TerminalFrame {
self.drain_output(DEFAULT_READ_TIMEOUT);
TerminalFrame::new(self.cols, self.rows, &self.output_buffer)
}
pub fn wait_for_text(
&mut self,
needle: &str,
timeout: Duration,
) -> Result<TerminalFrame, PtySessionError> {
let deadline = std::time::Instant::now() + timeout;
loop {
self.drain_output(Duration::from_millis(100));
let frame = TerminalFrame::new(self.cols, self.rows, &self.output_buffer);
if !frame.find_text(needle).is_empty() {
return Ok(frame);
}
if std::time::Instant::now() >= deadline {
return Err(PtySessionError::Timeout(format!(
"Text '{needle}' did not appear within {}ms",
timeout.as_millis()
)));
}
}
}
pub fn wait_for_stable_frame(
&mut self,
stable_duration: Duration,
timeout: Duration,
) -> Result<TerminalFrame, PtySessionError> {
let deadline = std::time::Instant::now() + timeout;
let mut previous_text = String::new();
let mut stable_since = std::time::Instant::now();
let mut seen_change = false;
loop {
self.drain_output(Duration::from_millis(100));
let frame = TerminalFrame::new(self.cols, self.rows, &self.output_buffer);
let current_text = frame.all_text();
if current_text != previous_text {
previous_text = current_text;
stable_since = std::time::Instant::now();
seen_change = true;
} else if seen_change
&& std::time::Instant::now().duration_since(stable_since) >= stable_duration
{
return Ok(frame);
}
if std::time::Instant::now() >= deadline {
return Err(PtySessionError::Timeout(
"Frame did not stabilize within timeout".to_string(),
));
}
}
}
pub fn run_eventually(
&mut self,
timeout: Duration,
poll: Duration,
predicate: &(dyn Fn(&TerminalFrame) -> crate::assertion::MatchResult + Send + Sync),
) -> Result<TerminalFrame, PtySessionError> {
scenario::eventually_loop(
timeout,
poll,
|| {
self.drain_output(Duration::ZERO);
TerminalFrame::new(self.cols, self.rows, &self.output_buffer)
},
predicate,
thread::sleep,
Instant::now,
)
.map_err(PtySessionError::Assertion)
}
pub fn execute_steps(&mut self, steps: &[Step]) -> Result<TerminalFrame, PtySessionError> {
let mut current_frame = None;
for step in steps {
match step {
Step::WriteText(text) => {
self.write_text(text)?;
current_frame = None;
}
Step::PressKey(key) => {
self.press_key(key)?;
current_frame = None;
}
Step::Sleep(duration) => {
thread::sleep(*duration);
current_frame = None;
}
Step::WaitForText { needle, timeout_ms } => {
let timeout = Duration::from_millis(u64::from(*timeout_ms));
current_frame = Some(self.wait_for_text(needle, timeout)?);
}
Step::WaitForStableFrame {
stable_ms,
timeout_ms,
} => {
let stable = Duration::from_millis(u64::from(*stable_ms));
let timeout = Duration::from_millis(u64::from(*timeout_ms));
current_frame = Some(self.wait_for_stable_frame(stable, timeout)?);
}
Step::Eventually {
timeout,
poll,
predicate,
} => {
current_frame =
Some(self.run_eventually(*timeout, *poll, predicate.as_ref())?);
}
Step::Capture | Step::CaptureLabeled { .. } => {
current_frame = Some(self.capture_frame());
}
Step::ViewingPause(_) => {
}
}
}
let final_frame = match current_frame {
Some(frame) => frame,
None => self.capture_frame(),
};
Ok(final_frame)
}
pub fn cols(&self) -> u16 {
self.cols
}
pub fn rows(&self) -> u16 {
self.rows
}
pub fn binary_path(&self) -> &Path {
&self.binary_path
}
pub fn wait_for_exit(&mut self, timeout: Duration) -> Result<bool, PtySessionError> {
let deadline = std::time::Instant::now() + timeout;
loop {
match self.child.try_wait() {
Ok(Some(status)) => return Ok(status.success()),
Ok(None) => {
if std::time::Instant::now() >= deadline {
return Err(PtySessionError::Timeout(
"Process did not exit within timeout".into(),
));
}
thread::sleep(Duration::from_millis(50));
}
Err(err) => {
return Err(PtySessionError::Timeout(format!(
"Error waiting for process exit: {err}"
)));
}
}
}
}
}
impl Drop for PtySession {
fn drop(&mut self) {
let _ = self.child.kill();
let _ = self.child.wait();
}
}
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum PtySessionError {
#[error("PTY creation failed: {0}")]
PtyCreation(String),
#[error("Failed to spawn binary: {0}")]
SpawnFailed(String),
#[error("Write failed: {0}")]
WriteFailed(String),
#[error("Timeout: {0}")]
Timeout(String),
#[error("Assertion failed: {0}")]
Assertion(Box<AssertionFailure>),
}
fn read_pty_output(mut reader: Box<dyn Read + Send>, sender: &mpsc::Sender<Vec<u8>>) {
let mut buffer = [0u8; 4096];
loop {
match reader.read(&mut buffer) {
Ok(0) | Err(_) => break,
Ok(bytes_read) => {
if sender.send(buffer[..bytes_read].to_vec()).is_err() {
break;
}
}
}
}
}
fn key_to_bytes(key: &str) -> Vec<u8> {
match key.to_lowercase().as_str() {
"enter" | "return" => vec![b'\r'],
"tab" => vec![b'\t'],
"backtab" | "shift+tab" => vec![0x1b, b'[', b'Z'],
"escape" | "esc" => vec![0x1b],
"backspace" => vec![0x7f],
"up" => vec![0x1b, b'[', b'A'],
"down" => vec![0x1b, b'[', b'B'],
"right" => vec![0x1b, b'[', b'C'],
"left" => vec![0x1b, b'[', b'D'],
"home" => vec![0x1b, b'[', b'H'],
"end" => vec![0x1b, b'[', b'F'],
"delete" => vec![0x1b, b'[', b'3', b'~'],
"pageup" => vec![0x1b, b'[', b'5', b'~'],
"pagedown" => vec![0x1b, b'[', b'6', b'~'],
"space" => vec![b' '],
other => {
if let Some(character) = other.strip_prefix("ctrl+")
&& character.len() == 1
&& let Some(byte) = character.bytes().next()
&& byte.to_ascii_lowercase().is_ascii_lowercase()
{
let ctrl_byte = byte.to_ascii_lowercase() - b'a' + 1;
return vec![ctrl_byte];
}
other.as_bytes().to_vec()
}
}
}
#[must_use]
pub struct PtySessionBuilder {
binary_path: PathBuf,
cols: u16,
rows: u16,
args: Vec<String>,
env_vars: Vec<(String, String)>,
workdir: Option<PathBuf>,
}
impl PtySessionBuilder {
pub fn new(binary_path: impl Into<PathBuf>) -> Self {
Self {
binary_path: binary_path.into(),
cols: DEFAULT_COLS,
rows: DEFAULT_ROWS,
args: Vec::new(),
env_vars: Vec::new(),
workdir: None,
}
}
pub fn size(mut self, cols: u16, rows: u16) -> Self {
self.cols = cols;
self.rows = rows;
self
}
pub fn args<I, S>(mut self, args: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.args.extend(args.into_iter().map(Into::into));
self
}
pub fn env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.env_vars.push((key.into(), value.into()));
self
}
pub fn workdir(mut self, directory: impl Into<PathBuf>) -> Self {
self.workdir = Some(directory.into());
self
}
pub fn spawn(self) -> Result<PtySession, PtySessionError> {
let arg_refs: Vec<&str> = self.args.iter().map(String::as_str).collect();
let env_refs: Vec<(&str, &str)> = self
.env_vars
.iter()
.map(|(key, value)| (key.as_str(), value.as_str()))
.collect();
PtySession::spawn_with_size(
&self.binary_path,
self.cols,
self.rows,
&arg_refs,
&env_refs,
self.workdir.as_deref(),
)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn pty_session_builder_forwards_args() {
let builder = PtySessionBuilder::new("/bin/echo").args(["--help", "--version"]);
assert_eq!(
builder.args,
vec!["--help".to_string(), "--version".to_string()]
);
}
#[test]
fn pty_session_builder_args_appends_across_calls() {
let builder = PtySessionBuilder::new("/bin/echo")
.args(["one"])
.args(vec![String::from("two"), String::from("three")]);
assert_eq!(builder.args, vec!["one", "two", "three"]);
}
#[test]
fn key_to_bytes_returns_ctrl_a() {
let bytes = key_to_bytes("ctrl+a");
assert_eq!(bytes, vec![0x01]);
}
#[test]
fn key_to_bytes_returns_ctrl_z() {
let bytes = key_to_bytes("ctrl+z");
assert_eq!(bytes, vec![0x1a]);
}
#[test]
fn key_to_bytes_ctrl_multi_char_falls_through() {
let bytes = key_to_bytes("ctrl+ab");
assert_eq!(bytes, "ctrl+ab".as_bytes());
}
#[test]
fn key_to_bytes_ctrl_non_alpha_falls_through() {
let bytes = key_to_bytes("ctrl+[");
assert_eq!(bytes, "ctrl+[".as_bytes());
}
#[test]
fn key_to_bytes_known_keys() {
assert_eq!(key_to_bytes("enter"), vec![b'\r']);
assert_eq!(key_to_bytes("tab"), vec![b'\t']);
assert_eq!(key_to_bytes("escape"), vec![0x1b]);
assert_eq!(key_to_bytes("backspace"), vec![0x7f]);
assert_eq!(key_to_bytes("space"), vec![b' ']);
}
#[test]
fn key_to_bytes_backtab_sends_csi_z() {
assert_eq!(key_to_bytes("backtab"), vec![0x1b, b'[', b'Z']);
assert_eq!(key_to_bytes("shift+tab"), vec![0x1b, b'[', b'Z']);
}
#[test]
fn key_to_bytes_unknown_key_returns_raw_bytes() {
let bytes = key_to_bytes("x");
assert_eq!(bytes, vec![b'x']);
}
#[test]
fn execute_steps_captures_input_after_an_earlier_wait() {
let mut session = PtySessionBuilder::new("/bin/sh")
.args([
"-c",
"printf 'ready\\n'; read value; printf 'done:%s\\n' \"$value\"; sleep 60",
])
.spawn()
.expect("failed to spawn interactive shell script");
let steps = [
Step::wait_for_text("ready", 3_000),
Step::write_text("hello\n"),
];
let frame = session
.execute_steps(&steps)
.expect("step execution should succeed");
assert!(
frame.all_text().contains("done:hello"),
"final frame must include output produced after the wait"
);
}
#[test]
fn wait_for_stable_frame_times_out_when_no_output() {
let temp_dir = tempfile::tempdir().expect("failed to create temp dir");
let script_path = temp_dir.path().join("silent.sh");
std::fs::write(&script_path, "#!/bin/sh\nsleep 60\n").expect("failed to write script");
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(&script_path, std::fs::Permissions::from_mode(0o755))
.expect("failed to set permissions");
}
let mut session = PtySession::spawn(&script_path).expect("failed to spawn silent script");
let result =
session.wait_for_stable_frame(Duration::from_millis(200), Duration::from_millis(800));
assert!(
matches!(result, Err(PtySessionError::Timeout(_))),
"should timeout when no frame change is observed"
);
}
#[test]
fn wait_for_stable_frame_returns_after_content_stabilizes() {
let temp_dir = tempfile::tempdir().expect("failed to create temp dir");
let script_path = temp_dir.path().join("greet.sh");
std::fs::write(&script_path, "#!/bin/sh\necho hello\nsleep 60\n")
.expect("failed to write script");
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(&script_path, std::fs::Permissions::from_mode(0o755))
.expect("failed to set permissions");
}
let mut session = PtySession::spawn(&script_path).expect("failed to spawn greet script");
let frame = session
.wait_for_stable_frame(Duration::from_millis(300), Duration::from_secs(5))
.expect("frame should stabilize");
let text = frame.all_text();
assert!(
text.contains("hello"),
"stable frame should contain echoed output, got: '{text}'"
);
}
}