Add Python REPL test + reenable previously-disabled test

.config/nextest.toml

@@ -8,9 +8,6 @@ status-level = "all"
 failure-output = "immediate-final"
 # Any tests that take longer than 10 minutes should be killed as a failure
 slow-timeout = { period = "60s", terminate-after = 10 }
-# Skip the Python test on CI before we figure out why it failed.
-default-filter = 'not test(test_runner_with_python)'
-
 [test-groups]
 # ensure only one test accessing tun at a time
 tun-access = { max-threads = 1 }

litebox_runner_linux_userland/tests/common/mod.rs

@@ -4,6 +4,9 @@
 use glob::glob;
 use std::path::{Path, PathBuf};
 
+#[cfg(target_os = "linux")]
+pub mod pty;
+
 /// Find all dependencies of a given binary via `ldd`
 #[allow(dead_code, reason = "not used by loader.rs for x86")]
 pub fn find_dependencies(prog: &str) -> Vec<String> {

litebox_runner_linux_userland/tests/common/pty.rs

@@ -0,0 +1,145 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+#![allow(
+    dead_code,
+    reason = "shared test helpers are not used by every test binary"
+)]
+
+pub struct Pty {
+    master: std::fs::File,
+    slave: Option<std::fs::File>,
+}
+
+impl Pty {
+    pub fn open() -> Self {
+        use std::os::fd::FromRawFd;
+
+        let mut master = -1;
+        let mut slave = -1;
+        // SAFETY: master and slave are valid out-pointers, and the optional name/termios/winsize
+        // pointers are null because the test does not need to customize the PTY.
+        let rc = unsafe {
+            libc::openpty(
+                &raw mut master,
+                &raw mut slave,
+                std::ptr::null_mut(),
+                std::ptr::null(),
+                std::ptr::null(),
+            )
+        };
+        assert_eq!(rc, 0, "openpty failed: {}", std::io::Error::last_os_error());
+
+        // SAFETY: master is an owned file descriptor returned by openpty above.
+        let flags = unsafe { libc::fcntl(master, libc::F_GETFL) };
+        assert_ne!(
+            flags,
+            -1,
+            "fcntl(F_GETFL) failed: {}",
+            std::io::Error::last_os_error()
+        );
+        // SAFETY: master is an owned file descriptor returned by openpty above, and flags were read
+        // from the same descriptor.
+        let rc = unsafe { libc::fcntl(master, libc::F_SETFL, flags | libc::O_NONBLOCK) };
+        assert_eq!(
+            rc,
+            0,
+            "fcntl(F_SETFL) failed: {}",
+            std::io::Error::last_os_error()
+        );
+
+        Self {
+            // SAFETY: openpty returned these owned file descriptors and they are not used elsewhere.
+            master: unsafe { std::fs::File::from_raw_fd(master) },
+            // SAFETY: openpty returned these owned file descriptors and they are not used elsewhere.
+            slave: Some(unsafe { std::fs::File::from_raw_fd(slave) }),
+        }
+    }
+
+    pub fn slave_stdio(
+        &self,
+    ) -> (
+        std::process::Stdio,
+        std::process::Stdio,
+        std::process::Stdio,
+    ) {
+        let slave = self.slave.as_ref().expect("PTY slave is already closed");
+        let stdin = slave.try_clone().expect("failed to clone pty slave");
+        let stdout = slave.try_clone().expect("failed to clone pty slave");
+        let stderr = slave.try_clone().expect("failed to clone pty slave");
+        (
+            std::process::Stdio::from(stdin),
+            std::process::Stdio::from(stdout),
+            std::process::Stdio::from(stderr),
+        )
+    }
+
+    pub fn close_slave(&mut self) {
+        drop(self.slave.take());
+    }
+
+    pub fn write_all(&mut self, bytes: &[u8]) {
+        use std::io::Write;
+
+        self.master
+            .write_all(bytes)
+            .expect("failed to write to pty");
+    }
+
+    pub fn wait_for_output(&mut self, output: &mut Vec<u8>, needle: &[u8]) {
+        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
+        loop {
+            self.read_available(output);
+            if output.windows(needle.len()).any(|window| window == needle) {
+                return;
+            }
+            assert!(
+                std::time::Instant::now() < deadline,
+                "timed out waiting for {:?}; output so far:\n{}",
+                String::from_utf8_lossy(needle),
+                String::from_utf8_lossy(output)
+            );
+            std::thread::sleep(std::time::Duration::from_millis(10));
+        }
+    }
+
+    pub fn wait_for_child_exit(
+        &mut self,
+        child: &mut std::process::Child,
+        output: &mut Vec<u8>,
+    ) -> std::process::ExitStatus {
+        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
+        loop {
+            self.read_available(output);
+            if let Some(status) = child.try_wait().expect("failed to wait for child process") {
+                self.read_available(output);
+                return status;
+            }
+            if std::time::Instant::now() >= deadline {
+                let _ = child.kill();
+                let _ = child.wait();
+                panic!(
+                    "timed out waiting for child process to exit; output so far:\n{}",
+                    String::from_utf8_lossy(output)
+                );
+            }
+            std::thread::sleep(std::time::Duration::from_millis(10));
+        }
+    }
+
+    fn read_available(&mut self, output: &mut Vec<u8>) {
+        use std::io::Read;
+
+        let mut buf = [0; 4096];
+        loop {
+            match self.master.read(&mut buf) {
+                Ok(0) => break,
+                Ok(n) => output.extend_from_slice(&buf[..n]),
+                Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => break,
+                Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {}
+                Err(e) if e.raw_os_error() == Some(libc::EIO) => break,
+                Err(e) => panic!("failed to read from pty: {e}"),
+            }
+        }
+    }
+}

litebox_runner_linux_userland/tests/run.rs

@@ -131,7 +131,7 @@ impl Runner {
         self.run_inner(true)
     }
 
-    fn run_inner(&mut self, capture_stdout: bool) -> Vec<u8> {
+    fn prepare_command(&mut self) {
         assert!(!self.has_run);
         self.has_run = true;
         // create tar file using `tar` command with caching
@@ -147,8 +147,12 @@ impl Runner {
             .arg("--initial-files")
             .arg(tar_file)
             .arg(&self.cmd_path)
-            .args(&self.cmd_args)
-            .stderr(std::process::Stdio::inherit());
+            .args(&self.cmd_args);
+    }
+
+    fn run_inner(&mut self, capture_stdout: bool) -> Vec<u8> {
+        self.prepare_command();
+        self.command.stderr(std::process::Stdio::inherit());
         if !capture_stdout {
             self.command.stdout(std::process::Stdio::inherit());
         }
@@ -164,6 +168,21 @@ impl Runner {
         );
         output.stdout
     }
+
+    #[cfg(all(target_arch = "x86_64", target_os = "linux"))]
+    fn spawn_with_stdio(
+        &mut self,
+        stdin: std::process::Stdio,
+        stdout: std::process::Stdio,
+        stderr: std::process::Stdio,
+    ) -> std::process::Child {
+        self.prepare_command();
+        self.command.stdin(stdin).stdout(stdout).stderr(stderr);
+        println!("Running `{:?}`", self.command);
+        self.command
+            .spawn()
+            .expect("Failed to spawn litebox_runner_linux_userland")
+    }
 }
 
 /// Find all C test files in a directory
@@ -306,11 +325,8 @@ fn run_python(args: &[&str]) -> String {
 }
 
 #[cfg(all(target_arch = "x86_64", target_os = "linux"))]
-#[test]
-fn test_runner_with_python() {
-    const HELLO_WORLD_PY: &str = "print(\"Hello, World from litebox!\")";
+fn python_runner(unique_name: &str) -> Runner {
     let python_path = run_which("python3");
-
     let python_guest_dir = python_path.parent().unwrap().to_str().unwrap().to_string();
 
     let python_home = run_python(&["-c", "import sys; print(sys.prefix);"]);
@@ -337,8 +353,8 @@ fn test_runner_with_python() {
     let mut paths_to_stage = std::collections::BTreeSet::new();
     paths_to_stage.extend(python_lib_paths.iter().cloned());
 
-    Runner::new(&python_path, "python_rewriter")
-        .args(["-c", HELLO_WORLD_PY])
+    let mut runner = Runner::new(&python_path, unique_name);
+    runner
         .envs([
             &format!("PYTHONHOME={python_home}"),
             &format!("PYTHONPATH={python_lib_paths_str}"),
@@ -429,10 +445,45 @@ fn test_runner_with_python() {
                     }
                 }
             }
-        })
+        });
+    runner
+}
+
+#[cfg(all(target_arch = "x86_64", target_os = "linux"))]
+#[test]
+fn test_runner_with_python() {
+    const HELLO_WORLD_PY: &str = "print(\"Hello, World from litebox!\")";
+    python_runner("python_rewriter")
+        .args(["-c", HELLO_WORLD_PY])
         .run();
 }
 
+#[cfg(all(target_arch = "x86_64", target_os = "linux"))]
+#[test]
+fn test_runner_with_python_repl_pty() {
+    let mut runner = python_runner("python_repl_pty_rewriter");
+    let mut pty = common::pty::Pty::open();
+    let (stdin, stdout, stderr) = pty.slave_stdio();
+    let mut child = runner.spawn_with_stdio(stdin, stdout, stderr);
+    pty.close_slave();
+
+    let mut output = Vec::new();
+    pty.wait_for_output(&mut output, b">>> ");
+    pty.write_all(b"print(\"hi\")\nexit()\n");
+    let status = pty.wait_for_child_exit(&mut child, &mut output);
+    assert!(
+        status.success(),
+        "Python exited with {status}; output:\n{}",
+        String::from_utf8_lossy(&output)
+    );
+
+    let output = String::from_utf8_lossy(&output).replace("\r\n", "\n");
+    assert!(
+        output.lines().any(|line| line.trim() == "hi"),
+        "Python REPL did not print a standalone hi line; output:\n{output}"
+    );
+}
+
 #[test]
 fn test_tun_with_tcp_socket() {
     let tcp_server_path = PathBuf::from("./tests/net/tcp_server.c");