Codex Handbook
core/src/unified_exec/process_tests.rs 210 lines
use super::process::UnifiedExecProcess;use crate::unified_exec::UnifiedExecError;use codex_exec_server::ExecProcess;use codex_exec_server::ExecProcessEventReceiver;use codex_exec_server::ExecProcessFuture;use codex_exec_server::ExecServerError;use codex_exec_server::ProcessId;use codex_exec_server::ProcessSignal;use codex_exec_server::ReadResponse;use codex_exec_server::StartedExecProcess;use codex_exec_server::WriteResponse;use codex_exec_server::WriteStatus;use codex_sandboxing::SandboxType;use pretty_assertions::assert_eq;use std::collections::VecDeque;use std::sync::Arc;use tokio::sync::Mutex;use tokio::sync::watch;use tokio::time::Duration;struct MockExecProcess {    process_id: ProcessId,    write_response: WriteResponse,    read_responses: Mutex<VecDeque<ReadResponse>>,    terminate_error: Option<String>,    wake_tx: watch::Sender<u64>,}impl MockExecProcess {    async fn read(&self) -> Result<ReadResponse, ExecServerError> {        Ok(self            .read_responses            .lock()            .await            .pop_front()            .unwrap_or(ReadResponse {                chunks: Vec::new(),                next_seq: 1,                exited: false,                exit_code: None,                closed: false,                failure: None,            }))    }    async fn terminate(&self) -> Result<(), ExecServerError> {        if let Some(message) = &self.terminate_error {            return Err(ExecServerError::Protocol(message.clone()));        }        Ok(())    }}impl ExecProcess for MockExecProcess {    fn process_id(&self) -> &ProcessId {        &self.process_id    }    fn subscribe_wake(&self) -> watch::Receiver<u64> {        self.wake_tx.subscribe()    }    fn subscribe_events(&self) -> ExecProcessEventReceiver {        ExecProcessEventReceiver::empty()    }    fn read(        &self,        _after_seq: Option<u64>,        _max_bytes: Option<usize>,        _wait_ms: Option<u64>,    ) -> ExecProcessFuture<'_, ReadResponse> {        Box::pin(MockExecProcess::read(self))    }    fn write(&self, _chunk: Vec<u8>) -> ExecProcessFuture<'_, WriteResponse> {        Box::pin(async { Ok(self.write_response.clone()) })    }    fn signal(&self, _signal: ProcessSignal) -> ExecProcessFuture<'_, ()> {        Box::pin(async { Ok(()) })    }    fn terminate(&self) -> ExecProcessFuture<'_, ()> {        Box::pin(MockExecProcess::terminate(self))    }}async fn remote_process(    write_status: WriteStatus,    terminate_error: Option<String>,) -> UnifiedExecProcess {    let (wake_tx, _wake_rx) = watch::channel(0);    let started = StartedExecProcess {        process: Arc::new(MockExecProcess {            process_id: "test-process".to_string().into(),            write_response: WriteResponse {                status: write_status,            },            read_responses: Mutex::new(VecDeque::new()),            terminate_error,            wake_tx,        }),    };    UnifiedExecProcess::from_exec_server_started(started, SandboxType::None)        .await        .expect("remote process should start")}#[tokio::test]async fn remote_write_unknown_process_marks_process_exited() {    let process = remote_process(WriteStatus::UnknownProcess, /*terminate_error*/ None).await;    let err = process        .write(b"hello")        .await        .expect_err("expected write failure");    assert!(matches!(err, UnifiedExecError::WriteToStdin));    assert!(process.has_exited());}#[tokio::test]async fn remote_write_closed_stdin_marks_process_exited() {    let process = remote_process(WriteStatus::StdinClosed, /*terminate_error*/ None).await;    let err = process        .write(b"hello")        .await        .expect_err("expected write failure");    assert!(matches!(err, UnifiedExecError::WriteToStdin));    assert!(process.has_exited());}#[tokio::test]async fn fail_and_terminate_preserves_failure_message() {    let process = remote_process(WriteStatus::Accepted, /*terminate_error*/ None).await;    process.fail_and_terminate("network denied".to_string());    process.fail_and_terminate("second failure".to_string());    assert!(process.has_exited());    assert_eq!(        process.failure_message(),        Some("network denied".to_string())    );}#[tokio::test]async fn remote_terminate_confirmed_updates_state_on_success_only() {    let process = remote_process(        WriteStatus::Accepted,        Some("terminate unavailable".to_string()),    )    .await;    let err = process        .terminate_confirmed()        .await        .expect_err("expected terminate failure");    assert!(matches!(err, UnifiedExecError::ProcessFailed { .. }));    assert!(!process.has_exited());    let process = remote_process(WriteStatus::Accepted, /*terminate_error*/ None).await;    process        .terminate_confirmed()        .await        .expect("terminate should succeed");    assert!(process.has_exited());}#[tokio::test]async fn remote_process_waits_for_early_exit_event() {    let (wake_tx, _wake_rx) = watch::channel(0);    let started = StartedExecProcess {        process: Arc::new(MockExecProcess {            process_id: "test-process".to_string().into(),            write_response: WriteResponse {                status: WriteStatus::Accepted,            },            read_responses: Mutex::new(VecDeque::from([ReadResponse {                chunks: Vec::new(),                next_seq: 2,                exited: true,                exit_code: Some(17),                closed: true,                failure: None,            }])),            terminate_error: None,            wake_tx: wake_tx.clone(),        }),    };    tokio::spawn(async move {        tokio::time::sleep(Duration::from_millis(10)).await;        let _ = wake_tx.send(1);    });    let process = UnifiedExecProcess::from_exec_server_started(started, SandboxType::None)        .await        .expect("remote process should observe early exit");    assert!(process.has_exited());    assert_eq!(process.exit_code(), Some(17));}