open-standards/wai/wai-rs/tests/manifest_reader_equivalence.rs

//! The reference readers build a manifest with `container::read_manifest` and
//! `read_manifest_v2`, which move its members out of one JSON tree (SPEC §2, *Reader
//! limits*: so a parse allocates within a bound), where they used serde's derived reader
//! (`serde_json::from_slice::<Manifest>`). This test holds the two to the same result:
//! for the manifest of every envelope committed under `wai/`, and for each of them with
//! every member, at the manifest's top level and one level down, removed or given a
//! value of each JSON type, both readers read it or both refuse it, and what they read
//! writes back to the same JSON.

use std::path::{Path, PathBuf};

use serde_json::{json, Value};
use wai::container::{check_unique_members, manifest_bytes, read_manifest, read_manifest_v2, Manifest, ManifestV2};

fn envelopes(dir: &Path, out: &mut Vec<PathBuf>) {
    let Ok(rd) = std::fs::read_dir(dir) else { return };
    let mut entries: Vec<_> = rd.flatten().map(|e| e.path()).collect();
    entries.sort();
    for p in entries {
        let name = p.file_name().and_then(|n| n.to_str()).unwrap_or("");
        if p.is_dir() {
            if name != "target" && name != "node_modules" && !name.starts_with('.') {
                envelopes(&p, out);
            }
        } else if name.ends_with(".wai") {
            out.push(p);
        }
    }
}

/// Both readers on `text`: the same verdict, and the same JSON when both read it.
fn same(text: &[u8], v2: bool, what: &str) -> bool {
    let (derived, moved) = if v2 {
        let d = check_unique_members(text).and_then(|_| serde_json::from_slice::<ManifestV2>(text)).map(|m| serde_json::to_string(&m).unwrap());
        let m = read_manifest_v2(text).map(|m| serde_json::to_string(&m).unwrap());
        (d, m)
    } else {
        let d = check_unique_members(text).and_then(|_| serde_json::from_slice::<Manifest>(text)).map(|m| serde_json::to_string(&m).unwrap());
        let m = read_manifest(text).map(|m| serde_json::to_string(&m).unwrap());
        (d, m)
    };
    match (&derived, &moved) {
        (Ok(a), Ok(b)) => assert_eq!(a, b, "{what}: the readers read different manifests"),
        (Err(_), Err(_)) => {}
        _ => panic!("{what}: the derived reader gives {derived:?}, the reference reader {moved:?}"),
    }
    moved.is_ok()
}

/// The values a member is given: one of each JSON type, and a few of each kind the
/// manifest types read (a string, an integer, a float, `null`).
fn replacements() -> Vec<Value> {
    vec![json!(null), json!(true), json!(0), json!(-1), json!(1.5), json!(18446744073709551615u64), json!("s"), json!([]), json!({}), json!([{}]), json!({"x": 1})]
}

fn mutations(text: &[u8], v2: bool, what: &str) -> usize {
    let Ok(Value::Object(top)) = serde_json::from_slice::<Value>(text) else { return 0 };
    let mut n = 0;
    let mut check = |m: &serde_json::Map<String, Value>, why: String| {
        same(serde_json::to_string(m).unwrap().as_bytes(), v2, &format!("{what}: {why}"));
        n += 1;
    };
    for (k, v) in &top {
        let mut m = top.clone();
        m.shift_remove(k);
        check(&m, format!("without {k}"));
        for r in replacements() {
            let mut m = top.clone();
            m.insert(k.clone(), r.clone());
            check(&m, format!("{k} = {r}"));
        }
        // One level down: each member of an object member, and of the first object of
        // an array member (a WAI2's first rendition entry).
        let (inner, at): (Option<&serde_json::Map<String, Value>>, Option<usize>) = match v {
            Value::Object(o) => (Some(o), None),
            Value::Array(a) => match a.first() {
                Some(Value::Object(o)) => (Some(o), Some(0)),
                _ => (None, None),
            },
            _ => (None, None),
        };
        let Some(inner) = inner else { continue };
        let put = |o: serde_json::Map<String, Value>| {
            let mut m = top.clone();
            let nv = match at {
                None => Value::Object(o),
                Some(i) => {
                    let mut a = v.as_array().unwrap().clone();
                    a[i] = Value::Object(o);
                    Value::Array(a)
                }
            };
            m.insert(k.clone(), nv);
            m
        };
        for ik in inner.keys() {
            let mut o = inner.clone();
            o.shift_remove(ik);
            check(&put(o), format!("{k} without {ik}"));
            for r in replacements() {
                let mut o = inner.clone();
                o.insert(ik.clone(), r.clone());
                check(&put(o), format!("{k}.{ik} = {r}"));
            }
        }
    }
    n
}

#[test]
fn the_reference_reader_reads_every_committed_manifest_as_the_derived_reader_does() {
    let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("..");
    let mut files = Vec::new();
    envelopes(&root, &mut files);
    assert!(files.len() >= 400, "found only {} envelopes under {}", files.len(), root.display());
    let (mut read, mut refused, mut mutated) = (0, 0, 0);
    // Mutate the manifests of each distinct set of top-level members once.
    let mut shapes = std::collections::BTreeSet::new();
    for f in &files {
        let b = std::fs::read(f).unwrap();
        let Ok(text) = manifest_bytes(&b) else { continue };
        let v2 = &b[..4] == b"WAI2";
        let what = f.strip_prefix(&root).unwrap_or(f).display().to_string();
        if same(text, v2, &what) {
            read += 1;
        } else {
            refused += 1;
        }
        if let Ok(Value::Object(o)) = serde_json::from_slice::<Value>(text) {
            let keys: Vec<String> = o.keys().cloned().collect();
            if shapes.insert((v2, keys)) {
                mutated += mutations(text, v2, &what);
            }
        }
    }
    eprintln!("{} envelopes: {read} manifests read alike, {refused} refused alike; {mutated} mutations agree", files.len());
    assert!(read >= 300, "only {read} manifests read");
    assert!(mutated >= 1000, "only {mutated} mutations");
}

/// The cases the derived reader treats specially, each read alike: a member whose value
/// is `null` where an `Option` is (`None`, not written back), `target` absent (`null`),
/// a member no revision defines kept in order, a float `min_sim`, and an integer past
/// `u64` as `kid`.
#[test]
fn the_derived_readers_special_cases_are_read_alike() {
    let base = r#"{"wai":"1.0","media":"image","intent":"replicate","model_requirement":{"capability":"wai.image.png","fallback":null,"min_sim":0.5,"z":1,"a":[1]},"conditioning":{"kind":"png","prior":null,"q":{}},"security":null,"refines":null,"derived_from":null,"later":{"b":2},"earlier":3}"#;
    for t in [
        base.to_string(),
        base.replace(r#""min_sim":0.5"#, r#""min_sim":1"#),
        base.replace(r#""security":null"#, r#""security":{"encryption":"none","kid":18446744073709551616}"#),
        base.replace(r#""security":null"#, r#""security":{"kid":-1}"#),
        base.replace(r#""conditioning":{"kind":"png","prior":null,"q":{}}"#, r#""conditioning":{"kind":"png","prior":{"id":"x","sha256":"y","determinism":"z","supersedes":null,"k":[]}}"#),
        base.replace(r#""conditioning":{"kind":"png","prior":null,"q":{}}"#, r#""conditioning":{"kind":"png","prior":{"id":"x","sha256":"y","determinism":"z","supersedes":{"sha256":"w","form":null}}}"#),
        base.replace(r#""conditioning":{"kind":"png","prior":null,"q":{}}"#, r#""conditioning":{"kind":"png","prior":{"id":"x","sha256":"y"}}"#),
        "[]".into(),
        "null".into(),
        "{}".into(),
    ] {
        same(t.as_bytes(), false, &t);
    }
    let v2 = r#"{"wai":"1.1","media":"text","intent":"replicate","renditions":[{"capability":"wai.text.zstd","kind":"zstd","refines":null,"envelope_manifest":{"a":1},"x":[]}],"security":{"encryption":"none"},"z":1}"#;
    for t in [
        v2.to_string(),
        v2.replace(r#""renditions":[{"#, r#""renditions":[1,{"#),
        v2.replace(r#""renditions":["#, r#""renditions":{"a":"#).replace(r#""x":[]}],"#, r#""x":[]}},"#),
        v2.replace(r#""kind":"zstd","#, ""),
    ] {
        same(t.as_bytes(), true, &t);
    }
}