fix(test): capture worker stdout and stderr in test output (#14410)

1 files changed, 59 insertions, 5 deletions

diff --git a/runtime/ops/io.rs b/runtime/ops/io.rs
index d54c66efd..f18624eb2 100644
--- a/runtime/ops/io.rs
+++ b/runtime/ops/io.rs
@@ -77,13 +77,67 @@ pub fn init() -> Extension {
     .build()
 }
 
-pub fn init_stdio() -> Extension {
+pub enum StdioPipe {
+  Inherit,
+  File(StdFile),
+}
+
+impl Default for StdioPipe {
+  fn default() -> Self {
+    Self::Inherit
+  }
+}
+
+impl Clone for StdioPipe {
+  fn clone(&self) -> Self {
+    match self {
+      StdioPipe::Inherit => StdioPipe::Inherit,
+      StdioPipe::File(pipe) => StdioPipe::File(pipe.try_clone().unwrap()),
+    }
+  }
+}
+
+/// Specify how stdin, stdout, and stderr are piped.
+/// By default, inherits from the process.
+#[derive(Clone, Default)]
+pub struct Stdio {
+  pub stdin: StdioPipe,
+  pub stdout: StdioPipe,
+  pub stderr: StdioPipe,
+}
+
+pub fn init_stdio(stdio: Stdio) -> Extension {
+  // todo(dsheret): don't do this? Taking out the writers was necessary to prevent invalid handle panics
+  let stdio = Rc::new(RefCell::new(Some(stdio)));
+
   Extension::builder()
-    .state(|state| {
+    .state(move |state| {
+      let stdio = stdio
+        .borrow_mut()
+        .take()
+        .expect("Extension only supports being used once.");
       let t = &mut state.resource_table;
-      t.add(StdFileResource::stdio(&STDIN_HANDLE, "stdin"));
-      t.add(StdFileResource::stdio(&STDOUT_HANDLE, "stdout"));
-      t.add(StdFileResource::stdio(&STDERR_HANDLE, "stderr"));
+      t.add(StdFileResource::stdio(
+        match &stdio.stdin {
+          StdioPipe::Inherit => &STDIN_HANDLE,
+          StdioPipe::File(pipe) => pipe,
+        },
+        "stdin",
+      ));
+      t.add(StdFileResource::stdio(
+        match &stdio.stdout {
+          StdioPipe::Inherit => &STDOUT_HANDLE,
+          StdioPipe::File(pipe) => pipe,
+        },
+        "stdout",
+      ));
+      t.add(StdFileResource::stdio(
+        match &stdio.stderr {
+          StdioPipe::Inherit => &STDERR_HANDLE,
+          StdioPipe::File(pipe) => pipe,
+        },
+        "stderr",
+      ));
       Ok(())
     })
     .build()