From 11b48643da786ce9996eae59daa050ba92a6ee9d Mon Sep 17 00:00:00 2001
From: kajacx <kajacx@gmail.com>
Date: Sat, 22 Jul 2023 15:25:10 +0200
Subject: [PATCH] Add first redirect io test

---
 .../src/wasi/preview2/wasi_ctx_builder.rs     | 16 ++++
 tests/wasi_components/hello_wasi/protocol.wit |  2 +-
 tests/wasi_components/io_redirect/guest.rs    | 22 +++++
 tests/wasi_components/io_redirect/host.rs     | 95 +++++++++++++++++++
 .../wasi_components/io_redirect/protocol.wit  |  7 ++
 5 files changed, 141 insertions(+), 1 deletion(-)
 create mode 100644 tests/wasi_components/io_redirect/guest.rs
 create mode 100644 tests/wasi_components/io_redirect/host.rs
 create mode 100644 tests/wasi_components/io_redirect/protocol.wit

diff --git a/crates/wasm-bridge-js/src/wasi/preview2/wasi_ctx_builder.rs b/crates/wasm-bridge-js/src/wasi/preview2/wasi_ctx_builder.rs
index 6dd8fb9f..0038a3df 100644
--- a/crates/wasm-bridge-js/src/wasi/preview2/wasi_ctx_builder.rs
+++ b/crates/wasm-bridge-js/src/wasi/preview2/wasi_ctx_builder.rs
@@ -12,4 +12,20 @@ impl WasiCtxBuilder {
     pub fn build(self, _table: &mut Table) -> Result<WasiCtx> {
         Ok(WasiCtx::new())
     }
+
+    pub fn inherit_stdin(self) -> Self {
+        self // TODO: this could (in theory) be implemented with prompt, but no one probably wants that
+    }
+
+    pub fn inherit_stdout(self) -> Self {
+        self
+    }
+
+    pub fn inherit_stderr(self) -> Self {
+        self
+    }
+
+    pub fn inherit_stdio(self) -> Self {
+        self.inherit_stdin().inherit_stdout().inherit_stderr()
+    }
 }
diff --git a/tests/wasi_components/hello_wasi/protocol.wit b/tests/wasi_components/hello_wasi/protocol.wit
index 8d428b08..806be529 100644
--- a/tests/wasi_components/hello_wasi/protocol.wit
+++ b/tests/wasi_components/hello_wasi/protocol.wit
@@ -1,4 +1,4 @@
-package component-test:wit-protocol
+package component-test:wasi-protocol
 
 world test-world {
   import add-one: func(num: s32) -> s32
diff --git a/tests/wasi_components/io_redirect/guest.rs b/tests/wasi_components/io_redirect/guest.rs
new file mode 100644
index 00000000..51fd5f6d
--- /dev/null
+++ b/tests/wasi_components/io_redirect/guest.rs
@@ -0,0 +1,22 @@
+wit_bindgen::generate!({
+    path: "../protocol.wit",
+    world: "io-redirect",
+});
+
+struct GuestImpl;
+
+impl IoRedirect for GuestImpl {
+    fn readln_from_stdin() -> Option<String> {
+        None // TODO:
+    }
+
+    fn writeln_to_stdout(line: String) {
+        println!("{line}");
+    }
+
+    fn writeln_to_stderr(line: String) {
+        eprintln!("{line}");
+    }
+}
+
+export_io_redirect!(GuestImpl);
diff --git a/tests/wasi_components/io_redirect/host.rs b/tests/wasi_components/io_redirect/host.rs
new file mode 100644
index 00000000..78c5751e
--- /dev/null
+++ b/tests/wasi_components/io_redirect/host.rs
@@ -0,0 +1,95 @@
+use wasm_bridge::{
+    component::{Component, Linker},
+    Config, Engine, Result, Store,
+};
+
+use wasm_bridge::wasi::preview2::*;
+
+wasm_bridge::component::bindgen!({
+    path: "../protocol.wit",
+    world: "io-redirect",
+    async: true,
+    with: {
+       "wasi:cli-base/stdin": wasi::cli_base::stdin,
+       "wasi:cli-base/stdout": wasi::cli_base::stdout,
+       "wasi:cli-base/stderr": wasi::cli_base::stderr,
+    }
+});
+
+struct State {
+    table: Table,
+    wasi: WasiCtx,
+}
+
+impl WasiView for State {
+    fn table(&self) -> &Table {
+        &self.table
+    }
+    fn table_mut(&mut self) -> &mut Table {
+        &mut self.table
+    }
+    fn ctx(&self) -> &WasiCtx {
+        &self.wasi
+    }
+    fn ctx_mut(&mut self) -> &mut WasiCtx {
+        &mut self.wasi
+    }
+}
+
+pub async fn run_test(component_bytes: &[u8]) -> Result<()> {
+    no_config(component_bytes).await?;
+    inherit(component_bytes).await?;
+
+    Ok(())
+}
+
+async fn no_config(component_bytes: &[u8]) -> Result<()> {
+    let mut config = Config::new();
+    config.wasm_component_model(true);
+    config.async_support(true);
+
+    let mut table = Table::new();
+    let wasi = WasiCtxBuilder::new().build(&mut table)?;
+
+    let engine = Engine::new(&config)?;
+    let mut store = Store::new(&engine, State { table, wasi });
+
+    let component = Component::new(&store.engine(), &component_bytes)?;
+
+    let mut linker = Linker::new(store.engine());
+    wasi::command::add_to_linker(&mut linker)?;
+
+    let (instance, _) = IoRedirect::instantiate_async(&mut store, &component, &linker).await?;
+
+    let result = instance.call_readln_from_stdin(&mut store).await?;
+    assert_eq!(result, None);
+
+    instance.call_writeln_to_stdout(&mut store, "NO_PRINT").await?;
+    instance.call_writeln_to_stderr(&mut store, "NO_PRINT").await?;
+
+    Ok(())
+}
+
+async fn inherit(component_bytes: &[u8]) -> Result<()> {
+    let mut config = Config::new();
+    config.wasm_component_model(true);
+    config.async_support(true);
+
+    let mut table = Table::new();
+    let wasi = WasiCtxBuilder::new().inherit_stdio().build(&mut table)?;
+
+    let engine = Engine::new(&config)?;
+    let mut store = Store::new(&engine, State { table, wasi });
+
+    let component = Component::new(&store.engine(), &component_bytes)?;
+
+    let mut linker = Linker::new(store.engine());
+    wasi::command::add_to_linker(&mut linker)?;
+
+    let (instance, _) = IoRedirect::instantiate_async(&mut store, &component, &linker).await?;
+
+    instance.call_writeln_to_stdout(&mut store, "PRINT_OUT_1").await?;
+    instance.call_writeln_to_stderr(&mut store, "PRINT_ERR_1").await?;
+
+    Ok(())
+}
diff --git a/tests/wasi_components/io_redirect/protocol.wit b/tests/wasi_components/io_redirect/protocol.wit
new file mode 100644
index 00000000..e5c4983b
--- /dev/null
+++ b/tests/wasi_components/io_redirect/protocol.wit
@@ -0,0 +1,7 @@
+package component-test:wasi-protocol
+
+world io-redirect {
+  export readln-from-stdin: func() -> option<string>
+  export writeln-to-stdout: func(line: string)
+  export writeln-to-stderr: func(line: string)
+}