mod commands;
use std::ffi::OsStr;
use std::time::{Duration, Instant};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::process::Child;
use vt100;
pub use commands::{Code, AsAnsi};
pub struct PtyTest {
pty: pty_process::Pty,
child: Child,
wait_timeout: Duration,
buf: [u8; 0x1000],
parser: vt100::Parser,
}
#[derive(Debug)]
pub enum Error {
ProcessExited(ScreenDiff, AsciiScreen),
IoError(std::io::Error),
TimeoutForScreenState(ScreenDiff, AsciiScreen),
}
impl std::fmt::Display for Error {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
writeln!(f, "PtyError {{")?;
let maybe_ascii_state = match self {
Error::TimeoutForScreenState(screen_diff, ascii_state) => {
for line in format!("{}", screen_diff).split("\n") {
writeln!(f, " {}", line)?;
}
Some(ascii_state)
}
Error::ProcessExited(screen_diff, ascii_state) => {
for line in format!("{}", screen_diff).split("\n") {
writeln!(f, " {}", line)?;
}
Some(ascii_state)
}
Error::IoError(_) => None,
};
if let Some(ascii_state) = maybe_ascii_state {
writeln!(f, " Screen {{")?;
for line in ascii_state.as_fragments().split("\n") {
if !line.is_empty() {
writeln!(f, " {}", line)?;
}
}
writeln!(f, " }}")?;
}
match self {
Error::TimeoutForScreenState { .. } => {
writeln!(f, " Screen timeout")?;
}
Error::ProcessExited { .. } => {
writeln!(f, " Process exited")?;
}
Error::IoError(err) => {
writeln!(f, " IoError: {:?}", err)?;
}
};
writeln!(f, "}}")?;
Ok(())
}
}
#[macro_export]
macro_rules! ascii_screen {
($($x:tt)*) => {$crate::AsciiScreen::new(file!(), line!(), &[
$($crate::ascii_screen_fragment!{$x}),*
])};
}
pub enum AsciiScreenFragment {
Newline,
String(&'static str),
Underscore(usize),
CursorPosition,
Nothing,
}
impl AsciiScreenFragment {
pub fn by_ident(s: &'static str) -> Self {
let mut underscore_prefix = 0;
for i in s.chars() {
if i == '_' {
underscore_prefix += 1;
continue;
} else {
panic!("unknown character in indent {:?}", i);
}
}
return AsciiScreenFragment::Underscore(underscore_prefix);
}
}
#[derive(Debug)]
pub struct AsciiScreen {
source_info: Option<(&'static str, u32)>,
contents: String,
cursor_rowcol: Option<(u16, u16)>,
}
impl AsciiScreen {
pub fn new(file: &'static str, line: u32, asf_list: &[AsciiScreenFragment]) -> Self {
let mut contents = String::new();
let mut rows = 0;
let mut cols = 0;
let mut cursor_rowcol = None;
for asf in asf_list {
match asf {
AsciiScreenFragment::String(line) => {
cols = 0;
contents.push_str(line);
}
AsciiScreenFragment::Newline => {
contents.push_str("\n");
rows += 1;
}
AsciiScreenFragment::Underscore(added_cols) => {
cols += *added_cols;
}
AsciiScreenFragment::CursorPosition => {
if cursor_rowcol.is_some() {
panic!("more than one cursor defined");
}
cursor_rowcol = Some((rows as u16, cols as u16));
}
AsciiScreenFragment::Nothing => {
continue;
}
}
}
if contents.ends_with("\n") {
contents.pop();
}
AsciiScreen {
source_info: Some((file, line)),
contents,
cursor_rowcol,
}
}
fn as_fragments(&self) -> String {
use std::fmt::Write;
let mut s = String::new();
if let Some((file, line)) = self.source_info {
let _ = writeln!(&mut s, "Source: {}:{}:", file, line);
}
for (_idx, line) in self.contents.split("\n").enumerate() {
let _ = writeln!(&mut s, "{:?}, NL, ", line);
}
s
}
}
#[derive(Debug)]
pub struct Change<E> {
expected: E,
received: E,
}
#[derive(Debug)]
pub struct ScreenDiff {
cursor_pos_diff: Option<Change<Option<(u16, u16)>>>,
content_diff: Option<(String, String)>,
}
impl std::fmt::Display for ScreenDiff {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
writeln!(f, "ScreenDiff {{")?;
match &self.content_diff {
None => {}
Some((expected, found)) => {
let changeset = difference::Changeset::new(&expected, &found, "\n");
use difference::Difference;
for change in changeset.diffs {
match change {
Difference::Rem(x) => {
for line in x.split("\n") {
writeln!(f, " -{}", line)?;
}
}
Difference::Add(x) => {
for line in x.split("\n") {
writeln!(f, " +{}", line)?;
}
}
Difference::Same(x) => {
for line in x.split("\n") {
writeln!(f, " {}", line)?;
}
}
}
}
}
}
if let Some(ref cursor_diff) = self.cursor_pos_diff {
writeln!(f, " Cursor position:")?;
match (cursor_diff.expected, cursor_diff.received) {
(Some((exp_row, exp_col)), Some((rec_row, rec_col))) => {
writeln!(f, " expected: ({}, {})", exp_row, exp_col)?;
writeln!(f, " received: ({}, {})", rec_row, rec_col)?;
}
(Some((exp_row, exp_col)), None) => {
writeln!(f, " expected: ({}, {})", exp_row, exp_col)?;
writeln!(f, " received: (hidden)")?;
}
(None, Some((rec_row, rec_col))) => {
writeln!(f, " expected: (hidden)")?;
writeln!(f, " received: ({}, {})", rec_row, rec_col)?;
}
(None, None) => {
writeln!(f, " both hidden")?;
}
}
}
write!(f, "}}")?;
Ok(())
}
}
#[macro_export]
macro_rules! ascii_screen_fragment {
($x:literal) => { $crate::AsciiScreenFragment::String($x) };
(NL) => { $crate::AsciiScreenFragment::Newline };
($x:ident) => { $crate::AsciiScreenFragment::by_ident(stringify!($x)) };
(^) => { $crate::AsciiScreenFragment::CursorPosition };
(,) => { $crate::AsciiScreenFragment::Nothing };
}
impl PtyTest {
pub async fn new_with_args<S>(program: S, args: Vec<String>, size: &SizeInfo) -> Self
where
S: AsRef<OsStr>,
{
let (pty, pts) = pty_process::open().expect("Failed to create pty");
pty.resize(pty_process::Size::new_with_pixel(
size.lines as u16,
size.cols as u16,
size.width as u16,
size.height as u16,
)).expect("Failed to resize pty");
let child = pty_process::Command::new(program)
.args(args)
.spawn(pts)
.expect("Failed to spawn command");
let parser = vt100::Parser::new(size.lines as u16, size.cols as u16, 0);
Self {
pty,
child,
parser,
buf: [0u8; 0x1000],
wait_timeout: Duration::from_millis(1000),
}
}
pub fn ascii_state(&self) -> AsciiScreen {
let screen = self.parser.screen();
let cursor = if screen.hide_cursor() {
None
} else {
Some(screen.cursor_position())
};
AsciiScreen {
source_info: None,
contents: screen.contents(),
cursor_rowcol: cursor,
}
}
pub fn diff(&self, screen_state: &AsciiScreen) -> Result<(), ScreenDiff> {
let screen = self.parser.screen();
let mut content_diff = None;
let mut cursor_pos_diff = None;
if screen_state.contents != screen.contents() {
content_diff = Some((screen_state.contents.clone(), screen.contents()));
}
let cursor = if screen.hide_cursor() {
None
} else {
Some(screen.cursor_position())
};
if cursor != screen_state.cursor_rowcol {
cursor_pos_diff = Some(Change {
expected: screen_state.cursor_rowcol,
received: cursor,
});
}
if cursor_pos_diff.is_some() || content_diff.is_some() {
return Err(ScreenDiff {
cursor_pos_diff,
content_diff,
});
}
Ok(())
}
pub async fn write_str(&mut self, codes: &str) -> Result<(), Error> {
match self.pty.write_all(codes.as_bytes()).await {
Err(io) => return Err(Error::IoError(io)),
Ok(()) => {}
}
Ok(())
}
pub async fn write<D, T>(&mut self, code: D) -> Result<(), Error>
where
D: AsRef<T>,
T: commands::AsAnsi,
{
let mut s = String::new();
code.as_ref().add_to_string(&mut s);
self.write_str(&s).await
}
pub async fn wait_for(&mut self, screen_state: &AsciiScreen) -> Result<(), Error> {
let end_time = Instant::now() + self.wait_timeout;
let mut diff: Option<ScreenDiff>;
loop {
let current_diff = match self.diff(screen_state) {
Ok(()) => return Ok(()),
Err(d) => d,
};
if Instant::now() >= end_time {
return Err(Error::TimeoutForScreenState(current_diff, self.ascii_state()));
}
diff = Some(current_diff);
if let Ok(Some(_status)) = self.child.try_wait() {
if let Some(d) = diff.take() {
return Err(Error::ProcessExited(d, self.ascii_state()));
}
}
match self.pty.read(&mut self.buf[..]).await {
Ok(0) => {
continue;
}
Ok(nread) => {
self.parser.process(&self.buf[0..nread]);
}
Err(err) => {
return Err(Error::IoError(err));
}
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
}
}
#[derive(Debug)]
pub struct SizeInfo {
pub(crate) lines: usize,
pub(crate) cols: usize,
pub width: usize,
pub height: usize,
}
impl SizeInfo {
pub fn new(cols: usize, lines: usize) -> Self {
Self {
lines,
cols,
width: 10,
height: 5,
}
}
pub fn lines(&self) -> (usize, ()) {
(self.lines, ())
}
pub fn cols(&self) -> (usize, ()) {
(self.cols, ())
}
}