asterinas/test/nixos/common/framework/src/session.rs

286 lines
8.9 KiB
Rust

// SPDX-License-Identifier: MPL-2.0
use rexpect::session::PtySession;
use super::Error;
/// Describes a runtime session configuration.
///
/// A session descriptor defines how to interact with a particular execution context,
/// such as a shell environment, container, or remote session. It specifies the prompt
/// pattern (by `expect_prompt`) and the commands needed to enter (by `cmd_to_enter`)
/// and exit (by `cmd_to_exit`) the context.
///
/// # Example
///
/// The session descriptor uses a fluent builder API:
///
/// ```rust
/// use nixos_test_framework::SessionDesc;
///
/// let desc = SessionDesc::new()
/// .expect_prompt("/ #")
/// .cmd_to_enter("podman run -it alpine")
/// .cmd_to_exit("exit");
/// ```
pub struct SessionDesc {
prompt: &'static str,
enter_command: &'static str,
exit_cmd: &'static str,
}
impl SessionDesc {
/// Creates a new session descriptor with empty fields.
///
/// Use the builder methods to configure the session before use.
pub fn new() -> Self {
Self {
prompt: "",
enter_command: "",
exit_cmd: "",
}
}
/// Sets the expected prompt pattern for this session.
pub fn expect_prompt(mut self, prompt: &'static str) -> Self {
self.prompt = prompt;
self
}
/// Sets the command used to enter this session.
pub fn cmd_to_enter(mut self, enter_command: &'static str) -> Self {
self.enter_command = enter_command;
self
}
/// Sets the command used to exit this session.
pub fn cmd_to_exit(mut self, exit_cmd: &'static str) -> Self {
self.exit_cmd = exit_cmd;
self
}
}
impl Default for SessionDesc {
fn default() -> Self {
Self::new()
}
}
/// An interactive session for running commands in a test environment.
///
/// `Session` provides a high-level interface for interacting with the test environment.
/// It manages execution contexts and handles nested environments automatically.
///
/// It uses a [`SessionDesc`] to track the current prompt and the correct command
/// to exit the current context.
pub struct Session {
desc: SessionDesc,
pty_session: PtySession,
}
impl Session {
/// Creates a new session with the given PTY session and descriptor.
pub(super) fn new(desc: SessionDesc, pty_session: PtySession) -> Self {
Self { desc, pty_session }
}
fn output_error(error: &Error) {
match error {
Error::EOF {
expected,
got,
exit_code,
} => {
println!("=== EOF Error Details ===");
println!("Expected: {}", expected);
println!(
"Got: {}",
String::from_utf8_lossy(&strip_ansi_escapes::strip(got))
);
println!("Exit code: {:?}", exit_code);
println!("========================");
}
Error::Timeout {
expected,
got,
timeout,
} => {
println!("=== Timeout Error Details ===");
println!("Expected: {}", expected);
println!(
"Got: {}",
String::from_utf8_lossy(&strip_ansi_escapes::strip(got))
);
println!("Timeout: {:?}", timeout);
println!("============================");
}
_ => {}
}
}
/// Executes a command in the current session.
///
/// This method runs the specified command and waits for it to complete.
/// The command is considered complete when the session prompt reappears.
///
/// Returns an error if the command times out or the session terminates unexpectedly.
///
/// # Example
///
/// ```rust
/// use nixos_test_framework::*;
///
/// fn example(nixos_shell: &mut Session) -> Result<(), Error> {
/// // Execute simple commands
/// nixos_shell.run_cmd("ls -la")?;
/// nixos_shell.run_cmd("cd /tmp")?;
/// nixos_shell.run_cmd("mkdir test_dir")?;
///
/// Ok(())
/// }
/// ```
pub fn run_cmd(&mut self, command: &str) -> Result<(), Error> {
println!("--> Running: {}", command);
self.pty_session.send_line(command)?;
// Read and consume the echoed command line
self.pty_session.exp_string(command).unwrap();
if let Err(e) = self.pty_session.exp_string(self.desc.prompt) {
Self::output_error(&e);
return Err(e);
}
Ok(())
}
/// Executes a command and verifies its output contains expected text.
///
/// This method runs the command and checks that the specified string appears
/// in the output. This is useful for validating command results.
///
/// Returns an error if:
/// - The expected string is not found in the output
/// - The command times out
/// - The session terminates unexpectedly
///
/// # Example
///
/// ```rust
/// use nixos_test_framework::*;
///
/// fn example(nixos_shell: &mut Session) -> Result<(), Error> {
/// // Verify output contains expected string
/// nixos_shell.run_cmd_and_expect("echo 'Hello, World!'", "Hello")?;
///
/// // Check if a file exists
/// nixos_shell.run_cmd_and_expect("ls /etc/hostname", "/etc/hostname")?;
///
/// // Verify system information
/// nixos_shell.run_cmd_and_expect("cat /etc/os-release", "NixOS")?;
///
/// Ok(())
/// }
/// ```
pub fn run_cmd_and_expect(&mut self, command: &str, expected: &str) -> Result<(), Error> {
println!("--> Running: {} (expecting: {})", command, expected);
self.pty_session.send_line(command)?;
// Read and consume the echoed command line
self.pty_session.exp_string(command).unwrap();
match self.pty_session.exp_string(self.desc.prompt) {
Ok(unread) => {
let cleaned_unread =
String::from_utf8_lossy(&strip_ansi_escapes::strip(&unread)).to_string();
if !cleaned_unread.contains(expected) {
println!("=== Unexpected Output ===");
println!("Expected: {}", expected);
println!("Output before prompt:\n{}", cleaned_unread);
println!("=========================");
return Err(Error::EOF {
expected: expected.to_string(),
got: cleaned_unread,
exit_code: None,
});
}
}
Err(e) => {
Self::output_error(&e);
return Err(e);
}
}
Ok(())
}
pub(super) fn run<F>(&mut self, test_ops: F) -> Result<(), Error>
where
F: FnOnce(&mut Session) -> Result<(), Error>,
{
(test_ops)(self)
}
/// Enters a nested session, runs operations, and automatically exits.
///
/// This method is used to work with nested environments like containers, SSH sessions,
/// or any interactive shell.
///
/// Returns an error if entering, running operations, or exiting fails.
///
/// # Example
///
/// ```rust
/// use nixos_test_framework::*;
///
/// fn container_test(nixos_shell: &mut Session) -> Result<(), Error> {
/// // Define the container session
/// let container_session_desc = SessionDesc::new()
/// .expect_prompt("/ #")
/// .cmd_to_enter("podman run -it docker.io/library/alpine")
/// .cmd_to_exit("exit");
///
/// // Enter container, run tests, and automatically exit
/// nixos_shell.enter_session_and_run(container_session_desc, |alpine_shell| {
/// alpine_shell.run_cmd_and_expect(
/// "cat /etc/os-release",
/// "Alpine"
/// )?;
/// Ok(())
/// })?;
///
/// // Back in the host - container has been exited
/// nixos_shell.run_cmd("echo 'Back on host'")?;
/// Ok(())
/// }
/// ```
pub fn enter_session_and_run<F>(&mut self, desc: SessionDesc, test_ops: F) -> Result<(), Error>
where
F: FnOnce(&mut Session) -> Result<(), Error>,
{
let old_desc = std::mem::replace(&mut self.desc, desc);
if let Err(e) = self.run_cmd(self.desc.enter_command) {
self.desc = old_desc;
return Err(e);
}
let res = self.run(test_ops);
let exit_cmd = self.desc.exit_cmd;
self.desc = old_desc;
let exit_res = self.run_cmd(exit_cmd);
res?;
exit_res?;
Ok(())
}
pub(super) fn shutdown(&mut self) -> Result<(), Error> {
self.pty_session.send_line(self.desc.exit_cmd)?;
self.pty_session.process.wait()?;
Ok(())
}
}