sausage/sdk/tests/wasi.rs

//! End-to-end protocol checks against a WASI plugin under Wasmtime.
use sausage_plugin::serde_json::{Value, json};
use std::{
    io::{BufRead, BufReader, Write},
    path::{Path, PathBuf},
    process::{Command, Stdio},
};
struct Process {
    child: std::process::Child,
    input: Option<std::process::ChildStdin>,
    output: BufReader<std::process::ChildStdout>,
}
impl Process {
    fn send(&mut self, v: Value) {
        writeln!(self.input.as_mut().unwrap(), "{v}").unwrap();
    }
    fn read(&mut self) -> Value {
        loop {
            let mut s = String::new();
            assert_ne!(self.output.read_line(&mut s).unwrap(), 0, "unexpected EOF");
            if let Ok(v) = sausage_plugin::serde_json::from_str::<Value>(&s)
                && v["jsonrpc"] == "2.0"
            {
                return v;
            }
        }
    }
    fn request(&mut self, method: &str, params: Value) {
        self.send(json!({"jsonrpc":"2.0","id":"probe-1","method":method,"params":params}));
    }
    fn answer(&mut self, r: &Value, result: Value) {
        self.send(json!({"jsonrpc":"2.0","id":r["id"],"result":result}));
    }
}
fn start(wasm: &Path, dir: &Path) -> Process {
    let mut child = Command::new("wasmtime")
        .args(["run", "--dir"])
        .arg(format!("{}::/", dir.display()))
        .arg(wasm)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()
        .unwrap();
    let input = child.stdin.take();
    let output = BufReader::new(child.stdout.take().unwrap());
    Process {
        child,
        input,
        output,
    }
}
#[test]
fn protocol() {
    let root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    let target = root.join("../target/sdk/probe");
    assert!(
        Command::new("cargo")
            .env("CARGO_TARGET_DIR", &target)
            .args(["build", "--offline", "--manifest-path"])
            .arg(root.join("tests/plugins/probe/Cargo.toml"))
            .args(["--target", "wasm32-wasip2"])
            .status()
            .unwrap()
            .success()
    );
    let wasm = target.join("wasm32-wasip2/debug/sdk-probe.wasm");
    let dir = std::env::temp_dir().join(format!("sausage-sdk-{}", std::process::id()));
    std::fs::create_dir_all(&dir).unwrap();
    let mut p = start(&wasm, &dir);
    p.request("rpc.discover", json!({}));
    let d = p.read();
    assert_eq!(d["result"]["x-sausage-grace"], 12);
    assert_eq!(
        d["result"]["methods"][0]["description"],
        "Add two numbers.\n\nThe result is their sum."
    );
    assert!(p.child.wait().unwrap().success());
    let mut p = start(&wasm, &dir);
    p.request("add", json!({"a":2,"b":4}));
    assert_eq!(p.read()["result"], 6);
    assert!(p.child.wait().unwrap().success());
    let mut p = start(&wasm, &dir);
    p.request("fail", json!({}));
    assert_eq!(p.read()["error"]["code"], -32010);
    assert!(p.child.wait().unwrap().success());
    let mut p = start(&wasm, &dir);
    p.request("exercise", json!({}));
    let a = p.read();
    let b = p.read();
    assert_eq!(a["method"], "other.first");
    assert_eq!(b["method"], "core.call");
    assert_eq!(b["params"]["detach"], true);
    p.send(json!({"jsonrpc":"2.0","id":9876,"result":{}}));
    p.send(json!({"jsonrpc":"2.0","method":"core.notify","params":{"call":a["id"],"name":"tick","params":{}}}));
    p.answer(&b, json!({}));
    p.answer(&a, json!({}));
    assert_eq!(p.read()["method"], "core.notify");
    let grant = p.read();
    assert_eq!(grant["method"], "core.lock");
    p.answer(&grant, json!({"held":true}));
    assert_eq!(p.read()["method"], "core.unlock");
    let attempt = p.read();
    assert_eq!(attempt["params"]["wait"], false);
    p.answer(&attempt, json!({"held":false}));
    assert_eq!(p.read()["result"], "done");
    assert!(p.child.wait().unwrap().success());
    let mut err = String::new();
    std::io::Read::read_to_string(p.child.stderr.as_mut().unwrap(), &mut err).unwrap();
    assert!(err.contains("unknown ID"));
    assert!(err.contains("notification: tick"));
    let mut p = start(&wasm, &dir);
    p.request("rpc.retired", json!({"previous":"hash"}));
    assert_eq!(p.read()["result"], "hash");
    assert!(p.child.wait().unwrap().success());
    let mut p = start(&wasm, &dir);
    p.request("late", json!({}));
    let grant = p.read();
    let barrier = p.read();
    p.answer(&grant, json!({"held":true}));
    assert_eq!(p.read()["method"], "core.unlock");
    p.answer(&barrier, json!({}));
    let kept = p.read();
    p.answer(&kept, json!({"held":true}));
    assert_eq!(p.read()["result"], "late done");
    assert!(p.child.wait().unwrap().success());
    let mut p = start(&wasm, &dir);
    p.request("wait", json!({}));
    assert_eq!(p.read()["method"], "other.wait");
    drop(p.input.take());
    assert_eq!(p.read()["params"]["name"], "stopped");
    assert!(p.child.wait().unwrap().success());
    assert_eq!(std::fs::read(dir.join("stopped")).unwrap(), b"yes");
    let mut p = start(&wasm, &dir);
    p.request("start", json!({}));
    let first = p.read();
    let second = p.read();
    let outgoing = if first["method"] == "other.background" {
        first
    } else {
        second
    };
    assert_eq!(outgoing["method"], "other.background");
    p.answer(&outgoing, json!({}));
    assert!(p.child.wait().unwrap().success());
    std::fs::remove_dir_all(dir).unwrap();
}