chore(core) Consolidate apply_patch tests

codex-rs/core/Cargo.toml

@@ -60,6 +60,7 @@ shlex = { workspace = true }
 similar = { workspace = true }
 strum_macros = { workspace = true }
 tempfile = { workspace = true }
+test-case = "3.3.1"
 test-log = { workspace = true }
 thiserror = { workspace = true }
 time = { workspace = true, features = [

codex-rs/core/tests/common/responses.rs

@@ -400,6 +400,21 @@ pub fn ev_apply_patch_function_call(call_id: &str, patch: &str) -> Value {
     })
 }
 
+pub fn ev_apply_patch_shell_call(call_id: &str, patch: &str) -> Value {
+    let args = serde_json::json!({ "command": ["apply_patch", patch] });
+    let arguments = serde_json::to_string(&args).expect("serialize apply_patch arguments");
+
+    ev_function_call(call_id, "shell", &arguments)
+}
+
+pub fn ev_apply_patch_shell_call_via_heredoc(call_id: &str, patch: &str) -> Value {
+    let script = format!("apply_patch <<'EOF'\n{patch}\nEOF\n");
+    let args = serde_json::json!({ "command": ["bash", "-lc", script] });
+    let arguments = serde_json::to_string(&args).expect("serialize apply_patch arguments");
+
+    ev_function_call(call_id, "shell", &arguments)
+}
+
 pub fn sse_failed(id: &str, code: &str, message: &str) -> String {
     sse(vec![serde_json::json!({
         "type": "response.failed",

codex-rs/core/tests/common/test_codex.rs

@@ -24,11 +24,29 @@ use tempfile::TempDir;
 use wiremock::MockServer;
 
 use crate::load_default_config_for_test;
+use crate::responses::ev_apply_patch_custom_tool_call;
+use crate::responses::ev_apply_patch_function_call;
+use crate::responses::ev_apply_patch_shell_call;
+use crate::responses::ev_apply_patch_shell_call_via_heredoc;
+use crate::responses::ev_assistant_message;
+use crate::responses::ev_completed;
+use crate::responses::ev_response_created;
+use crate::responses::mount_sse_sequence;
+use crate::responses::sse;
 use crate::responses::start_mock_server;
 use crate::wait_for_event;
 
 type ConfigMutator = dyn FnOnce(&mut Config) + Send;
 
+/// A collection of different ways the model can output an apply_patch call
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub enum ApplyPatchModelOutput {
+    Freeform,
+    Function,
+    Shell,
+    ShellViaHeredoc,
+}
+
 pub struct TestCodexBuilder {
     config_mutators: Vec<Box<ConfigMutator>>,
 }
@@ -258,6 +276,48 @@ impl TestCodexHarness {
             .expect("output string")
             .to_string()
     }
+
+    pub async fn mount_apply_patch_call(
+        &self,
+        call_id: &str,
+        patch: &str,
+        assistant_msg: &str,
+        output_type: ApplyPatchModelOutput,
+    ) {
+        let ev_fn = match output_type {
+            ApplyPatchModelOutput::Freeform => ev_apply_patch_custom_tool_call,
+            ApplyPatchModelOutput::Function => ev_apply_patch_function_call,
+            ApplyPatchModelOutput::Shell => ev_apply_patch_shell_call,
+            ApplyPatchModelOutput::ShellViaHeredoc => ev_apply_patch_shell_call_via_heredoc,
+        };
+
+        let responses = vec![
+            sse(vec![
+                ev_response_created("resp-1"),
+                ev_fn(call_id, patch),
+                ev_completed("resp-1"),
+            ]),
+            sse(vec![
+                ev_assistant_message("msg-1", assistant_msg),
+                ev_completed("resp-2"),
+            ]),
+        ];
+
+        mount_sse_sequence(self.server(), responses).await;
+    }
+
+    pub async fn get_patch_output(
+        &self,
+        call_id: &str,
+        output_type: ApplyPatchModelOutput,
+    ) -> String {
+        match output_type {
+            ApplyPatchModelOutput::Freeform => self.custom_tool_call_output(call_id).await,
+            ApplyPatchModelOutput::Function
+            | ApplyPatchModelOutput::Shell
+            | ApplyPatchModelOutput::ShellViaHeredoc => self.function_call_stdout(call_id).await,
+        }
+    }
 }
 
 fn custom_tool_call_output<'a>(bodies: &'a [Value], call_id: &str) -> &'a Value {