open-standards/wai/wai-quantum-qpu/tests/protocol_schema.rs

//! The open provider protocol's messages held to their schema
//! (`schemas/provider.v1.json`): every message the client sends and every
//! answer the reference provider gives, for each operation and each error,
//! with no member the schema does not name; and the protocol's constants are
//! the schema's.

use serde_json::{Value, json};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use wai_quantum::ed25519_dalek::SigningKey;
use wai_quantum::engine::schema::Registry;
use wai_quantum_qpu::Wait;
use wai_quantum_qpu::job::{self, Target};
use wai_quantum_qpu::open_provider::{ERROR_CODES, Local, NATIVE_FORMAT, OPERATIONS, OpenProvider, PROTOCOL};
use wai_quantum_qpu::reference::{Config, Per, Reference};
use wai_quantum_qpu::registry::Kind;

fn schemas() -> Registry {
    let mut r = Registry::empty();
    r.add("provider.v1.json", include_str!("../schemas/provider.v1.json")).unwrap();
    r
}

/// The schema an answer to `op` follows.
fn answer_of(op: &str, answer: &Value) -> &'static str {
    if answer["ok"] == false {
        return "provider.v1.json#/$defs/error";
    }
    match op {
        "describe" => "provider.v1.json#/$defs/describe_answer",
        "estimate" => "provider.v1.json#/$defs/estimate_answer",
        "submit" => "provider.v1.json#/$defs/submit_answer",
        "result" => "provider.v1.json#/$defs/result_answer",
        _ => "provider.v1.json#/$defs/status_answer",
    }
}

#[test]
fn the_schema_names_the_protocol_s_own_constants() {
    let r = schemas();
    let s = r.get("provider.v1.json").unwrap();
    assert_eq!(s["title"], PROTOCOL);
    assert_eq!(wai_quantum::engine::schema::file_of(PROTOCOL).as_deref(), Some("provider.v1.json"));
    assert_eq!(s["$id"], "https://git.transaction.science/Transaction-Science/open-standards/raw/branch/main/wai/wai-quantum-qpu/schemas/provider.v1.json");
    assert_eq!(r.dangling(), Vec::<String>::new());
    let d = &s["$defs"];
    assert_eq!(d["error"]["properties"]["error"]["properties"]["code"]["enum"], json!(ERROR_CODES));
    let ops: Vec<&str> = d["message"]["oneOf"].as_array().unwrap().iter().map(|m| m["$ref"].as_str().unwrap().rsplit('/').next().unwrap()).collect();
    assert_eq!(ops, OPERATIONS);
    for op in OPERATIONS {
        assert_eq!(d[op]["properties"]["op"]["const"], op);
        assert_eq!(d[op]["properties"]["protocol"]["const"], PROTOCOL);
    }
    assert_eq!(d["program"]["properties"]["format"]["const"], NATIVE_FORMAT);
}

#[test]
fn every_message_and_answer_follows_the_protocol_s_schema() {
    let r = schemas();
    let seen: Arc<Mutex<Vec<(Value, Value)>>> = Arc::default();
    use wai_quantum::quantum_toolchain::native::Basis;
    for config in [
        Config::default(),
        Config { price: Some(("0.0003".into(), "USD".into(), Per::Shot)), free_tier: Some("100 shots a day".into()), polls: 2, ..Config::default() },
        Config { basis: Basis::PrxCz, all_to_all: true, ..Config::default() },
        Config { basis: Basis::GpiZz, price: Some(("1".into(), "credits".into(), Per::Job)), ..Config::default() },
    ] {
        let reference = Arc::new(Reference::new(config).unwrap());
        let (log, rf) = (Arc::clone(&seen), Arc::clone(&reference));
        let channel = Local(move |m: &Value| {
            let a = rf.handle(m, None);
            log.lock().unwrap().push((m.clone(), a.clone()));
            a
        });
        let t = Target::new("lab", Kind::Process, Box::new(OpenProvider::new(channel)), None);
        let wait = Wait { interval: Duration::ZERO, polls: 10 };
        // A job end to end through the client: describe, estimate, submit,
        // status until done, result.
        let req = json!({ "qasm": "OPENQASM 3.0; qubit[3] q; bit[3] c; h q[0]; cx q[0], q[2]; t q[1]; c = measure q;", "shots": 64, "seed": 5, "provider": "lab" });
        let body = job::run(&t, &req, true, &SigningKey::from_bytes(&[2; 32]), wait, |_| {}).unwrap();
        assert!(wai_quantum::engine::job::verify(&body).unwrap().0);
        // And cancel, and every error the reference gives, sent as they are.
        let id = reference.handle(&json!({ "protocol": PROTOCOL, "op": "submit", "program": seen.lock().unwrap().iter().rev().find(|(m, _)| m["op"] == "submit").unwrap().0["program"], "shots": 8 }), None)["job"].clone();
        assert!(id.is_string(), "the job this configuration's client submitted, submitted again");
        for m in [
            json!({ "protocol": PROTOCOL, "op": "cancel", "job": id }),
            json!({ "protocol": PROTOCOL, "op": "result", "job": id }),
            json!({ "protocol": PROTOCOL, "op": "status", "job": "none" }),
            json!({ "protocol": PROTOCOL, "op": "teleport" }),
            json!({ "protocol": "wai.quantum.provider/9", "op": "describe" }),
            json!({ "op": "describe" }),
            json!({ "protocol": PROTOCOL, "op": "estimate", "program": { "format": "other" }, "shots": 1 }),
        ] {
            let a = reference.handle(&m, None);
            seen.lock().unwrap().push((m, a));
        }
    }
    let seen = seen.lock().unwrap();
    let mut ops = std::collections::BTreeSet::new();
    let mut problems = Vec::new();
    for (m, a) in seen.iter() {
        let op = m["op"].as_str().unwrap_or("");
        // The client's own messages, and those sent here well formed, follow
        // the message schema; the malformed ones are what the errors answer.
        if OPERATIONS.contains(&op) && m["protocol"] == PROTOCOL && m["program"]["format"] != "other" {
            problems.extend(r.check("provider.v1.json#/$defs/message", m).into_iter().map(|e| format!("message {m}: {e}")));
        }
        problems.extend(r.check(answer_of(op, a), a).into_iter().map(|e| format!("answer to {op}: {e}\n    {a}")));
        ops.insert(if a["ok"] == false { format!("error {}", a["error"]["code"].as_str().unwrap_or("")) } else { op.to_string() });
    }
    assert!(problems.is_empty(), "{}\nA message or an answer holds what the protocol's schema (schemas/provider.v1.json) does not say: name it there with its type and a description. An optional field an answer gains keeps the protocol's version (wai/extensions/quantum-providers.md, section 10); anything else is a new protocol version.", problems.join("\n"));
    let want = ["cancel", "describe", "error invalid", "error not_found", "error refused", "error unsupported", "error unsupported_version", "estimate", "result", "status", "submit"];
    assert_eq!(ops.iter().map(String::as_str).collect::<Vec<_>>(), want);
}