use std::io::{self, Read, Write};
use std::process::Command;
use std::sync::mpsc::{self, TryRecvError};
use std::time::{Duration, Instant};
use portable_pty::{CommandBuilder, MasterPty, PtySize, native_pty_system};
#[allow(dead_code)]
pub struct PtySession {
master: Box<dyn MasterPty + Send>,
child: Box<dyn portable_pty::Child + Send>,
writer: Box<dyn Write + Send>,
reader_rx: mpsc::Receiver<Vec<u8>>,
}
#[allow(dead_code)]
impl PtySession {}
impl PtySession {
pub fn spawn(args: &[&str]) -> io::Result<Self> {
let pty_system = native_pty_system();
let pty_pair = pty_system
.openpty(PtySize {
rows: 24,
cols: 80,
pixel_width: 0,
pixel_height: 0,
})
.map_err(pty_err_to_io)?;
let mut cmd = CommandBuilder::new("oxi");
cmd.args(args);
cmd.env("OXI_NO_USER_CONFIG", "1");
cmd.cwd(".");
let child = pty_pair.slave.spawn_command(cmd).map_err(pty_err_to_io)?;
let reader = pty_pair.master.try_clone_reader().map_err(pty_err_to_io)?;
let writer = pty_pair.master.take_writer().map_err(pty_err_to_io)?;
drop(pty_pair.slave);
let (tx, rx) = mpsc::channel::<Vec<u8>>();
std::thread::spawn(move || {
let mut reader = reader;
let mut buf = [0u8; 4096];
loop {
match reader.read(&mut buf) {
Ok(0) | Err(_) => break,
Ok(n) => {
if tx.send(buf[..n].to_vec()).is_err() {
break;
}
}
}
}
});
Ok(Self {
master: pty_pair.master,
child,
writer,
reader_rx: rx,
})
}
pub fn read_until(&mut self, pattern: &str, timeout: Duration) -> io::Result<String> {
let deadline = Instant::now() + timeout;
let mut buf = String::new();
loop {
match self.reader_rx.try_recv() {
Ok(bytes) => {
buf.push_str(&String::from_utf8_lossy(&bytes));
if buf.contains(pattern) {
return Ok(buf);
}
}
Err(TryRecvError::Empty) => {
if Instant::now() >= deadline {
return Ok(buf);
}
std::thread::sleep(Duration::from_millis(20));
}
Err(TryRecvError::Disconnected) => {
return Ok(buf);
}
}
}
}
#[allow(dead_code)]
pub fn send_line(&mut self, text: &str) -> io::Result<()> {
self.writer.write_all(text.as_bytes())?;
self.writer.write_all(b"\r")?; self.writer.flush()
}
#[allow(dead_code)]
pub fn send_raw(&mut self, bytes: &[u8]) -> io::Result<()> {
self.writer.write_all(bytes)?;
self.writer.flush()
}
#[allow(dead_code)]
pub fn resize(&self, cols: u16, rows: u16) -> io::Result<()> {
self.master
.resize(PtySize {
rows,
cols,
pixel_width: 0,
pixel_height: 0,
})
.map_err(pty_err_to_io)
}
pub fn try_wait(&mut self) -> io::Result<Option<u32>> {
self.child
.try_wait()
.map(|opt| opt.map(|status| status.exit_code()))
}
pub fn kill(&mut self) -> io::Result<()> {
self.child.kill()
}
}
impl Drop for PtySession {
fn drop(&mut self) {
let _ = self.child.kill();
let _ = self.child.wait();
}
}
fn pty_err_to_io(e: anyhow::Error) -> io::Error {
io::Error::other(e)
}
#[track_caller]
pub fn assert_output_contains(haystack: &str, needle: &str) {
assert!(
haystack.contains(needle),
"expected PTY output to contain {:?}\n\nactual output:\n---\n{}\n---",
needle,
haystack
);
}
pub fn oxi_binary_available() -> bool {
Command::new("oxi")
.arg("--version")
.output()
.map(|o| o.status.success())
.unwrap_or(false)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_assert_output_contains_pass() {
assert_output_contains("hello world", "hello");
}
#[test]
#[should_panic(expected = "expected PTY output to contain")]
fn test_assert_output_contains_fail() {
assert_output_contains("hello world", "missing");
}
}