open-standards/wai/wai-quantum-py/tests/test_cli.py

"""The package and the command line give the same answers and the same
receipts: the same job runs to the same result, each verifies the other's
receipts, and both sign with the same user key."""

from __future__ import annotations

import json

import pytest

import wai_quantum as wq
from conftest import BELL, bare

DEMO = wq.Key.demo()

JOBS = [
    ("bell", BELL, {"shots": 1000, "seed": 9}),
    ("rotations", "OPENQASM 3.0; qubit[3] q; ry(0.3) q[0]; cx q[0], q[1]; rz(1.7) q[1]; cx q[1], q[2]; rx(-0.42) q[2];", {"shots": 300, "seed": 4, "observables": ["Z0 Z1", "X2", "Y1"], "amplitudes": ["000", "111"]}),
    ("noisy", "OPENQASM 3.0; qubit[2] q; bit[2] c; h q[0]; cx q[0], q[1]; c = measure q;", {"shots": 500, "seed": 2, "noise": {"depolarizing_2q": 0.05, "readout": 0.02}}),
    ("dynamic", "OPENQASM 3.0; qubit[2] q; bit[2] c; h q[0]; c[0] = measure q[0]; if (c[0]) x q[1]; c[1] = measure q[1];", {"shots": 200, "seed": 6}),
    ("pauli", "OPENQASM 3.0; qubit[6] q; h q; rzz(0.4) q[0], q[1]; rx(0.9) q[3];", {"observables": ["Z0 Z1", "X3"], "method": "pauli"}),
]


def cli_args(spec: dict) -> list:
    out = []
    for k, v in spec.items():
        if k == "observables":
            for o in v:
                out += ["-o", o]
        elif k == "amplitudes":
            for a in v:
                out += ["-a", a]
        elif k == "noise":
            for ch, p in v.items():
                out += ["--noise", f"{ch}={p}"]
        else:
            out += [f"--{k}", str(v)]
    return out


@pytest.mark.parametrize("name,src,spec", JOBS, ids=[j[0] for j in JOBS])
def test_a_job_run_here_is_the_command_lines_job(name, src, spec, run_cli, tmp_path):
    f = tmp_path / "p.qasm"
    f.write_text(src)
    _, out = run_cli("run", str(f), "--demo-key", "--json", *cli_args(spec))
    theirs = json.loads(out)
    ours = wq.run(src, key=DEMO, **spec)
    assert {k: v for k, v in ours.items() if k != "result"} == {k: v for k, v in theirs.items() if k != "result"}
    assert bare(ours["result"]) == bare(theirs["result"])
    # Each verifies the other's receipt.
    assert wq.verify(theirs).valid
    sealed = tmp_path / "ours.json"
    sealed.write_text(json.dumps(ours))
    code, text = run_cli("verify", str(sealed))
    assert code == 0 and "signature    VALID" in text and DEMO.did in text


def test_the_command_lines_receipt_signs_the_canonical_bytes_python_writes(run_cli, tmp_path):
    f = tmp_path / "p.qasm"
    f.write_text(JOBS[1][1])
    _, out = run_cli("run", str(f), "--demo-key", "--json", *cli_args(JOBS[1][2]))
    job = json.loads(out)
    result = dict(job["result"])
    receipt = result.pop("receipt")
    # The digest is BLAKE3 over the request's and the result's canonical
    # text; Python, having read the command line's JSON, writes the same
    # bytes, so it recomputes the same digest.
    assert wq.digest(job["request"], result) == receipt["digest"]
    # And re-sealed here with the same key, the same signature (Ed25519 is
    # deterministic).
    assert wq.seal(job["request"], result, key=DEMO) == receipt


def test_a_forged_result_is_refused_by_both(run_cli, tmp_path):
    job = wq.run(BELL, shots=100, seed=1, key=DEMO)
    forged = json.loads(json.dumps(job))
    forged["result"]["counts"]["00"] += 1
    p = tmp_path / "forged.json"
    p.write_text(json.dumps(forged))
    code, text = run_cli("verify", str(p), check=False)
    assert code == 1 and "INVALID" in text
    assert not wq.verify(forged)


def test_a_key_file_signs_alike_on_both(run_cli, tmp_path):
    k = wq.Key.generate()
    path = k.save(tmp_path / "my.key")
    _, did = run_cli("key", "--key", str(path))
    assert did.strip() == k.did
    f = tmp_path / "p.qasm"
    f.write_text(BELL)
    _, out = run_cli("run", str(f), "--key", str(path), "--json", "--shots", "10")
    assert wq.verify(json.loads(out)).key == k.public_key


def test_both_find_the_same_user_key(run_cli):
    # The command line makes this user's key (in the test's own config
    # directory); the package reads the same file.
    _, did = run_cli("key")
    assert wq.Key.user().did == did.strip()
    _, demo = run_cli("key", "--demo-key")
    assert demo.strip() == DEMO.did


def test_plan_lint_critique_draw_and_methods_are_the_command_lines(run_cli, tmp_path):
    src = JOBS[1][1]
    f = tmp_path / "p.qasm"
    f.write_text(src)
    assert json.loads(run_cli("plan", str(f), "--json", *cli_args(JOBS[1][2]))[1]) == wq.plan(src, **JOBS[1][2])
    assert json.loads(run_cli("lint", str(f), "--json")[1]) == wq.lint(src)
    bad = tmp_path / "bad.qasm"
    bad.write_text("OPENQASM 3.0;\nqubit[1] q;\nfoo q[0];\n")
    code, out = run_cli("lint", str(bad), "--json", check=False)
    assert code == 1 and json.loads(out) == wq.lint(bad)
    assert json.loads(run_cli("critique", str(f), "--json")[1]) == wq.critique(src)
    assert run_cli("draw", str(f))[1] == wq.draw(src)
    assert run_cli("draw", str(f), "--width", "40")[1] == wq.draw(src, width=40)
    assert json.loads(run_cli("methods", "--json")[1]) == wq.methods()


def test_equivalent_and_improve_are_the_command_lines(run_cli, tmp_path):
    a, b = tmp_path / "a.qasm", tmp_path / "b.qasm"
    a.write_text("OPENQASM 3.0; qubit[2] q; h q[0]; h q[0]; t q[1]; t q[1];")
    b.write_text("OPENQASM 3.0; qubit[2] q; s q[1];")
    code, out = run_cli("equivalent", str(a), str(b), "--json", check=False)
    assert code == 0 and json.loads(out) == wq.equivalent(a, b)
    assert json.loads(run_cli("improve", str(a), "--json")[1]) == wq.improve(a)


def test_an_estimate_here_is_the_command_lines(run_cli, tmp_path):
    f = tmp_path / "e.qasm"
    f.write_text("OPENQASM 3.0;\nqubit[3] q;\nh q;\nccx q[0], q[1], q[2];\nrz(0.3) q[0];\n")
    _, out = run_cli("estimate", str(f), "--demo-key", "--json", "--qubit", "us-e3", "--budget", "1e-4", "--max-slowdown", "64")
    theirs = json.loads(out)
    ours = wq.estimate(f, qubit="us-e3", budget=1e-4, max_slowdown=64, calibration="exact_matching", name="e", key=DEMO)
    assert ours == theirs
    assert wq.verify_estimate(theirs)["valid"]


def test_a_run_with_an_intent_and_a_repair_are_the_command_lines(run_cli, tmp_path):
    f, dist = tmp_path / "bell.qasm", tmp_path / "dist.json"
    f.write_text(BELL)
    dist.write_text(json.dumps({"00": 0.5, "11": 0.5}))
    code, out = run_cli("run", str(f), "--shots", "100", "--intent", str(dist), "--json", "--demo-key", check=False)
    cli_job = json.loads(out)
    assert code == 0 and cli_job["result"]["verdict"]["level"] == "proof"
    here = wq.run(BELL, shots=100, intent=[{"distribution": {"00": 0.5, "11": 0.5}}], key=DEMO)
    assert bare(here["result"]) == bare(cli_job["result"])
    assert wq.verify(cli_job).valid
    broken = tmp_path / "broken.qasm"
    broken.write_text("OPENQASM 3.0;\nqubit[2] q\nh q[0];\n")
    assert json.loads(run_cli("repair", str(broken), "--json")[1]) == wq.repair(broken)


def test_error_correction_here_is_the_command_lines(run_cli, tmp_path):
    for args, ours in [
        (["code", "color:5"], lambda: wq.qec_code("color:5", key=DEMO)),
        (["decode", "--code", "surface:5", "--error", "3,12"], lambda: wq.qec_decode("surface:5", error=[3, 12], key=DEMO)),
        (["decode", "--memory", "3,3,0.004", "--shots", "500", "--decoder", "correlated", "--seed", "4"], lambda: wq.qec_decode(memory=(3, 3, 0.004), shots=500, decoder="correlated", seed=4, key=DEMO)),
        (["threshold", "--family", "color", "--distances", "3,5", "--p", "0.03", "--shots", "300"], lambda: wq.qec_threshold("color", distances=[3, 5], p=[0.03], shots=300, key=DEMO)),
    ]:
        _, out = run_cli("qec", *args, "--demo-key", "--json")
        theirs = json.loads(out)
        assert ours() == theirs, args
        assert wq.verify(theirs).valid
    # A certificate the command line wrote checks here.
    _, out = run_cli("qec", "code", "steane", "--demo-key", "--json")
    cert = tmp_path / "cert.json"
    cert.write_text(out)
    assert wq.qec_code("steane", certificate=cert, key=DEMO)["result"]["check"]["verified"]