use std::collections::VecDeque;
use std::io::{Read, Write};
use std::time::Duration;
#[derive(Debug, Clone)]
pub struct RecordedInteraction {
pub direction: InteractionDirection,
pub data: Vec<u8>,
pub timestamp: Duration,
}
impl RecordedInteraction {
#[must_use]
pub fn input(data: impl Into<Vec<u8>>, timestamp: Duration) -> Self {
Self {
direction: InteractionDirection::Input,
data: data.into(),
timestamp,
}
}
#[must_use]
pub fn output(data: impl Into<Vec<u8>>, timestamp: Duration) -> Self {
Self {
direction: InteractionDirection::Output,
data: data.into(),
timestamp,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum InteractionDirection {
Input,
Output,
}
#[derive(Debug)]
pub struct TestSession {
responses: VecDeque<Vec<u8>>,
interactions: Vec<RecordedInteraction>,
start_time: std::time::Instant,
closed: bool,
}
impl TestSession {
#[must_use]
pub fn new() -> Self {
Self {
responses: VecDeque::new(),
interactions: Vec::new(),
start_time: std::time::Instant::now(),
closed: false,
}
}
#[must_use]
pub fn builder() -> TestSessionBuilder {
TestSessionBuilder::new()
}
pub fn queue_response(&mut self, data: impl Into<Vec<u8>>) {
self.responses.push_back(data.into());
}
pub fn queue_response_str(&mut self, s: &str) {
self.queue_response(s.as_bytes().to_vec());
}
pub fn queue_line(&mut self, line: &str) {
let mut data = line.as_bytes().to_vec();
data.extend_from_slice(b"\r\n");
self.queue_response(data);
}
pub fn send(&mut self, data: &[u8]) {
let elapsed = self.start_time.elapsed();
self.interactions
.push(RecordedInteraction::input(data.to_vec(), elapsed));
}
pub fn receive(&mut self) -> Option<Vec<u8>> {
let data = self.responses.pop_front()?;
let elapsed = self.start_time.elapsed();
self.interactions
.push(RecordedInteraction::output(data.clone(), elapsed));
Some(data)
}
#[must_use]
pub fn interactions(&self) -> &[RecordedInteraction] {
&self.interactions
}
#[must_use]
pub fn inputs(&self) -> Vec<&RecordedInteraction> {
self.interactions
.iter()
.filter(|i| i.direction == InteractionDirection::Input)
.collect()
}
#[must_use]
pub fn outputs(&self) -> Vec<&RecordedInteraction> {
self.interactions
.iter()
.filter(|i| i.direction == InteractionDirection::Output)
.collect()
}
#[must_use]
pub fn has_pending(&self) -> bool {
!self.responses.is_empty()
}
pub const fn close(&mut self) {
self.closed = true;
}
#[must_use]
pub const fn is_closed(&self) -> bool {
self.closed
}
#[must_use]
pub fn all_input_str(&self) -> String {
let bytes: Vec<u8> = self.inputs().iter().flat_map(|i| i.data.clone()).collect();
String::from_utf8_lossy(&bytes).into_owned()
}
#[must_use]
pub fn all_output_str(&self) -> String {
let bytes: Vec<u8> = self.outputs().iter().flat_map(|i| i.data.clone()).collect();
String::from_utf8_lossy(&bytes).into_owned()
}
pub fn assert_sent(&self, needle: &str) {
let input = self.all_input_str();
assert!(
input.contains(needle),
"Expected to send {needle:?}, but sent:\n{input}"
);
}
pub fn assert_not_sent(&self, needle: &str) {
let input = self.all_input_str();
assert!(
!input.contains(needle),
"Expected NOT to send {needle:?}, but found it in:\n{input}"
);
}
}
impl Default for TestSession {
fn default() -> Self {
Self::new()
}
}
impl Read for TestSession {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
if self.closed {
return Ok(0);
}
match self.receive() {
Some(data) => {
let len = buf.len().min(data.len());
buf[..len].copy_from_slice(&data[..len]);
Ok(len)
}
None => Err(std::io::Error::new(
std::io::ErrorKind::WouldBlock,
"No data available",
)),
}
}
}
impl Write for TestSession {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
if self.closed {
return Err(std::io::Error::new(
std::io::ErrorKind::BrokenPipe,
"Session closed",
));
}
self.send(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
#[derive(Debug, Default)]
pub struct TestSessionBuilder {
responses: Vec<Vec<u8>>,
}
impl TestSessionBuilder {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn response(mut self, data: impl Into<Vec<u8>>) -> Self {
self.responses.push(data.into());
self
}
#[must_use]
pub fn response_str(self, s: &str) -> Self {
self.response(s.as_bytes().to_vec())
}
#[must_use]
pub fn line(self, line: &str) -> Self {
let mut data = line.as_bytes().to_vec();
data.extend_from_slice(b"\r\n");
self.response(data)
}
#[must_use]
pub fn prompt(self, prompt: &str) -> Self {
self.response_str(prompt)
}
#[must_use]
pub fn login_sequence(self, username: &str, password: &str) -> Self {
self.prompt("Login: ")
.line(username)
.prompt("Password: ")
.line(password)
.line("Welcome!")
.prompt("$ ")
}
#[must_use]
pub fn build(self) -> TestSession {
let mut session = TestSession::new();
for response in self.responses {
session.queue_response(response);
}
session
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_session_basic() {
let mut session = TestSession::new();
session.queue_response_str("Hello");
session.queue_response_str("World");
assert_eq!(session.receive(), Some(b"Hello".to_vec()));
session.send(b"test");
assert_eq!(session.receive(), Some(b"World".to_vec()));
assert_eq!(session.inputs().len(), 1);
assert_eq!(session.outputs().len(), 2);
}
#[test]
fn test_session_builder() {
let mut session = TestSession::builder()
.prompt("$ ")
.line("output line")
.build();
assert_eq!(session.receive(), Some(b"$ ".to_vec()));
assert_eq!(session.receive(), Some(b"output line\r\n".to_vec()));
}
#[test]
fn test_session_assertions() {
let mut session = TestSession::new();
session.send(b"hello world");
session.send(b"test");
session.assert_sent("hello");
session.assert_sent("world");
session.assert_not_sent("goodbye");
}
}