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

//! The manifest types are `#[non_exhaustive]`, so code outside the crate builds
//! them with constructors and builders. These tests run from outside the crate:
//! the builders write exactly the JSON the earlier struct literals wrote, a
//! committed envelope round-trips byte for byte, and members the Rust types model
//! only as `extra` (or a `security` object) survive a parse and a re-serialization.

use wai::container::{
    Conditioning, Manifest, ManifestV2, ModelRequirement, Prior, Rendition, RenditionMeta, Security,
    Supersedes, Wai, WaiMulti,
};

#[test]
fn builders_write_the_json_a_literal_wrote() {
    let m = Manifest::new(
        "image",
        "replicate",
        ModelRequirement::new("wai.neural.bmshj2018").with_fallback("wai.image.png"),
        Conditioning::new("bmshj2018").with_prior(Prior::new("q3", "ab".repeat(32), "int-rans")),
    )
    .with_target(serde_json::json!({"w": 256, "h": 256}));
    assert_eq!(
        serde_json::to_string(&m).unwrap(),
        format!(
            r#"{{"wai":"1.0","media":"image","intent":"replicate","model_requirement":{{"capability":"wai.neural.bmshj2018","fallback":"wai.image.png"}},"conditioning":{{"kind":"bmshj2018","prior":{{"id":"q3","sha256":"{}","determinism":"int-rans"}}}},"target":{{"w":256,"h":256}}}}"#,
            "ab".repeat(32)
        )
    );
    let plain = Manifest::new("text", "replicate", ModelRequirement::new("wai.text.zstd"), Conditioning::new("zstd"));
    assert_eq!(
        serde_json::to_string(&plain).unwrap(),
        r#"{"wai":"1.0","media":"text","intent":"replicate","model_requirement":{"capability":"wai.text.zstd"},"conditioning":{"kind":"zstd"},"target":null}"#
    );
    let semantic = ModelRequirement::new("wai.semantic.clip").with_min_sim(0.5).with_sim_metric("dreamsim");
    assert_eq!(
        serde_json::to_string(&semantic).unwrap(),
        r#"{"capability":"wai.semantic.clip","min_sim":0.5,"sim_metric":"dreamsim"}"#
    );
    let v2 = ManifestV2::new("audio", "replicate", vec![
        RenditionMeta::new("wai.neural.encodec32", "encodec_tokens").with_prior(Prior::new("e", "cd".repeat(32), "int")),
        RenditionMeta::new("wai.meta.hdr10plus", "descriptor").companion_to("wai.video.av1"),
    ])
    .with_target(serde_json::json!({"sr": 48000}));
    assert_eq!(
        serde_json::to_string(&v2).unwrap(),
        format!(
            r#"{{"wai":"1.1","media":"audio","intent":"replicate","renditions":[{{"capability":"wai.neural.encodec32","kind":"encodec_tokens","prior":{{"id":"e","sha256":"{}","determinism":"int"}}}},{{"capability":"wai.meta.hdr10plus","kind":"descriptor","role":"companion","companion_of":"wai.video.av1"}}],"target":{{"sr":48000}}}}"#,
            "cd".repeat(32)
        )
    );
    let sup = Prior::new("b", "bb".repeat(32), "int").with_supersedes(Supersedes::new("aa".repeat(32)).with_form("wai.prior.set"));
    assert!(serde_json::to_string(&sup).unwrap().ends_with(&format!(
        r#""supersedes":{{"sha256":"{}","form":"wai.prior.set"}}}}"#, "aa".repeat(32))));
    let with_sec = plain.with_wai("1.1").with_security(Security::new("mls-sframe", 7));
    assert!(serde_json::to_string(&with_sec).unwrap().contains(r#""security":{"encryption":"mls-sframe","kid":7}"#));
}

#[test]
fn the_committed_int_mlicpp_envelope_round_trips() {
    let d = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/vectors/int_mlicpp");
    for name in ["stream.wai", "stream_set.wai"] {
        let bytes = std::fs::read(d.join(name)).unwrap();
        assert_eq!(Wai::from_bytes(&bytes).unwrap().to_bytes().unwrap(), bytes, "{name}");
    }
}

#[test]
fn security_and_unknown_members_survive() {
    let v1 = r#"{"wai":"1.0","media":"text","intent":"replicate","model_requirement":{"capability":"wai.text.zstd","later_req":1},"conditioning":{"kind":"zstd","later_cond":[1,2]},"security":{"encryption":"mls-sframe","kid":3,"later_sec":"x"},"target":null,"later_top":{"a":true}}"#;
    let mut env = b"WAI1".to_vec();
    env.extend_from_slice(&(v1.len() as u32).to_le_bytes());
    env.extend_from_slice(v1.as_bytes());
    env.extend_from_slice(&2u32.to_le_bytes());
    env.extend_from_slice(b"ct");
    let w = Wai::from_bytes(&env).unwrap();
    assert_eq!(w.manifest.security.as_ref().and_then(|s| s.encryption.as_deref()), Some("mls-sframe"));
    assert!(!w.manifest.security.as_ref().unwrap().is_plaintext());
    assert_eq!(w.to_bytes().unwrap(), env, "nothing is dropped and the key order is kept");

    let v2 = r#"{"wai":"1.1","media":"video","intent":"replicate","renditions":[{"capability":"wai.video.av1","kind":"av1","id":"main","later_entry":{"n":1}}],"target":null,"later_top":2}"#;
    let manifest: ManifestV2 = serde_json::from_str(v2).unwrap();
    let m = WaiMulti::from_parts(manifest, vec![b"AV1".to_vec()]).unwrap();
    let back = WaiMulti::from_bytes(&m.to_bytes().unwrap()).unwrap();
    assert_eq!(serde_json::to_string(&back.manifest).unwrap(), v2);
    assert_eq!(back.manifest.renditions[0].extra["later_entry"]["n"], 1);
    // The payloads must match the entries.
    let one: ManifestV2 = serde_json::from_str(v2).unwrap();
    assert!(WaiMulti::from_parts(one, vec![]).is_err());
    // Rendition builders carry the identity members.
    let r = Rendition::media("wai.image.png", "png", b"PNG".to_vec()).with_id("p").with_payload_digest();
    assert_eq!(r.meta.id.as_deref(), Some("p"));
    assert_eq!(r.meta.sha256.as_deref(), Some(wai::pin::sha256_hex(b"PNG").as_str()));
}