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

"""The service client: against a stub that records what it is sent (how the
client names itself, its JSON, how it follows a job), and against the job
service built from this checkout, run here. No test reaches the public
service."""

from __future__ import annotations

import json
import threading
import time
from http.server import BaseHTTPRequestHandler, HTTPServer

import pytest

import wai_quantum as wq
from conftest import BELL, bare

DEMO = wq.Key.demo()

#: A stub's answer that is none: the connection closes once the request is read.
DROP = (None, None)


# -- a stub, for what the client sends ------------------------------------------


class Stub:
    """A server that answers from a script and keeps every request."""

    def __init__(self, answers):
        self.answers = list(answers)
        self.seen = []
        stub = self

        class H(BaseHTTPRequestHandler):
            def log_message(self, *a):
                pass

            def _any(self):
                n = int(self.headers.get("content-length") or 0)
                body = self.rfile.read(n) if n else b""
                stub.seen.append({"method": self.command, "path": self.path, "headers": dict(self.headers), "body": json.loads(body) if body else None})
                status, answer = stub.answers.pop(0) if stub.answers else (500, {"error": "the stub ran out of answers"})
                if status is None:
                    return
                raw = answer if isinstance(answer, bytes) else json.dumps(answer).encode()
                self.send_response(status)
                self.send_header("content-type", "application/json")
                self.send_header("content-length", str(len(raw)))
                self.end_headers()
                self.wfile.write(raw)

            do_GET = do_POST = do_DELETE = _any

        self.server = HTTPServer(("127.0.0.1", 0), H)
        self.url = f"http://127.0.0.1:{self.server.server_port}/quantum-jobs"
        threading.Thread(target=self.server.serve_forever, daemon=True).start()

    def close(self):
        self.server.shutdown()


@pytest.fixture
def stub():
    made = []

    def make(answers):
        s = Stub(answers)
        made.append(s)
        return s

    yield make
    for s in made:
        s.close()


def test_the_client_names_itself_and_sends_json(stub):
    done = {"id": "j1", "status": "done", "method": "stabilizer", "request": {"qasm": BELL}, "result": {"counts": {"00": 1}}}
    s = stub([(202, {"id": "j1", "status": "queued"}), (200, {"id": "j1", "status": "running"}), (200, done)])
    job = wq.Service(s.url, timeout=10).run(BELL, shots=1, wait=30, poll=0.01)
    assert isinstance(job, wq.Job) and job.counts == {"00": 1}
    post, *gets = s.seen
    assert post["method"] == "POST" and post["path"] == "/quantum-jobs/v1/run" and post["body"] == {"qasm": BELL, "shots": 1}
    assert [g["path"] for g in gets] == ["/quantum-jobs/v1/jobs/j1"] * 2
    for r in s.seen:
        assert r["headers"]["User-Agent"] == f"wai-quantum-python/{wq.__version__}" == wq.USER_AGENT
        assert r["headers"]["Accept"] == "application/json"
    assert post["headers"]["Content-Type"] == "application/json"


def test_the_base_may_name_v1_or_not_and_the_environment_may_name_it(stub, monkeypatch):
    s = stub([(200, {"ok": True}), (200, {"ok": True}), (200, {"ok": True})])
    assert wq.Service(s.url + "/v1").health() and wq.Service(s.url + "/").health()
    monkeypatch.setenv("WAI_QUANTUM_SERVICE", s.url)
    assert wq.Service().base == s.url and wq.Service().health()
    assert [r["path"] for r in s.seen] == ["/quantum-jobs/v1/health"] * 3
    monkeypatch.delenv("WAI_QUANTUM_SERVICE")
    assert wq.Service().base == wq.DEFAULT_SERVICE


def test_a_refusal_carries_the_services_status_and_words(stub):
    s = stub([(400, {"error": "line 1: unknown gate foo"}), (409, {"error": "the job has already finished"}), (502, b"<html>bad gateway</html>")])
    svc = wq.Service(s.url)
    with pytest.raises(wq.ServiceError) as e:
        svc.submit("qubit[1] q; foo q[0];")
    assert e.value.status == 400 and e.value.message == "line 1: unknown gate foo" and e.value.body == {"error": "line 1: unknown gate foo"}
    with pytest.raises(wq.ServiceError) as e:
        svc.cancel("j/1")
    assert e.value.status == 409 and s.seen[1]["path"] == "/quantum-jobs/v1/jobs/j%2F1" and s.seen[1]["method"] == "DELETE"
    with pytest.raises(wq.ServiceError, match="not JSON"):
        svc.describe()


def test_a_failed_or_slow_job_raises(stub):
    s = stub([(200, {"id": "j", "status": "failed", "error": "the worker ran out of memory"}), (200, {"id": "k", "status": "running"}), (200, {"id": "k", "status": "running"})])
    svc = wq.Service(s.url)
    with pytest.raises(wq.ServiceError, match="ran out of memory"):
        svc.result("j")
    with pytest.raises(wq.ServiceError, match="still running"):
        svc.result("k", wait=0.05, poll=0.1)


def test_a_job_whose_answer_is_lost_is_sent_again_under_its_key(stub):
    s = stub([DROP, (200, {"idempotency": {"on": ["POST /v1/jobs", "POST /v1/run"]}}), (202, {"id": "j1", "status": "queued"})])
    assert wq.Service(s.url, timeout=10).submit(BELL, shots=1)["id"] == "j1"
    assert [(r["method"], r["path"]) for r in s.seen] == [("POST", "/quantum-jobs/v1/jobs"), ("GET", "/quantum-jobs/v1"), ("POST", "/quantum-jobs/v1/jobs")]
    first, again = s.seen[0]["headers"]["Idempotency-Key"], s.seen[2]["headers"]["Idempotency-Key"]
    assert first == again and len(first) == 32
    # A key given is the key sent.
    s = stub([(202, {"id": "j2", "status": "queued"})])
    wq.Service(s.url).submit(BELL, shots=1, idempotency_key="mine-1")
    assert s.seen[0]["headers"]["Idempotency-Key"] == "mine-1" and s.seen[0]["body"] == {"qasm": BELL, "shots": 1}


def test_a_job_is_not_sent_again_to_a_service_that_keeps_no_keys(stub):
    s = stub([DROP, (200, {"service": "wai-quantum-service"})])
    with pytest.raises(wq.ServiceError) as e:
        wq.Service(s.url, timeout=10).submit(BELL, shots=1)
    assert e.value.status is None and [r["method"] for r in s.seen] == ["POST", "GET"]


def test_an_interrupted_job_ends_the_wait(stub):
    s = stub([(200, {"id": "j", "status": "interrupted", "error": "interrupted: the service stopped while this job ran"})])
    with pytest.raises(wq.ServiceError, match="interrupted"):
        wq.Service(s.url).result("j", wait=30, poll=0.01)
    assert len(s.seen) == 1


def test_an_unreachable_service_raises_without_a_status():
    with pytest.raises(wq.ServiceError) as e:
        wq.Service("http://127.0.0.1:9", timeout=2).describe()
    assert e.value.status is None


def test_verify_is_offline(stub):
    s = stub([])
    job = wq.run(BELL, shots=10, key=DEMO)
    assert wq.Service(s.url).verify(job).valid and s.seen == []


# -- the job service, run here ----------------------------------------------------


def test_the_service_describes_itself(service_url):
    svc = wq.Service(service_url)
    assert svc.health()
    d = svc.describe()
    assert d["service"] == "wai-quantum-service" and len(d["receipt_key"]) == 64


def test_a_job_run_on_the_service_is_the_job_run_here(service_url):
    svc = wq.Service(service_url)
    spec = {"shots": 700, "seed": 11, "observables": ["Z0 Z1"]}
    theirs = svc.run(BELL, **spec)
    ours = wq.run(BELL, key=DEMO, **spec)
    assert theirs.status == "done" and theirs["request"] == ours["request"]
    assert (theirs["method"], theirs["why"], theirs["estimate"]) == (ours["method"], ours["why"], ours["estimate"])
    assert bare(theirs.result) == bare(ours.result)
    # The service adds what the result came from, its host named in it.
    p = theirs.result["provenance"]
    assert p["schema"] == "wai.quantum.provenance/1" and p["software"]["host"].startswith("wai-quantum-service")
    # Its receipt, sealed with the service's key, verifies here.
    v = svc.verify(theirs)
    assert v.valid and v.key == svc.describe()["receipt_key"]
    forged = json.loads(json.dumps(theirs))
    forged["result"]["counts"]["00"] = forged["result"]["counts"].get("00", 0) + 1
    assert not wq.verify(forged)


def test_submit_status_and_result(service_url):
    svc = wq.Service(service_url)
    rec = svc.submit("OPENQASM 3.0; qubit[3] q; h q[0]; cx q[0], q[1]; ry(0.2) q[2];", observables=["Z2"])
    assert rec["status"] in ("queued", "running", "done") and rec["id"]
    job = svc.result(rec["id"], wait=120, poll=0.1)
    assert job.status == "done" and job.verify().valid
    assert svc.status(rec["id"])["status"] == "done"
    with pytest.raises(wq.ServiceError) as e:
        svc.status("no-such-job")
    assert e.value.status == 404


def test_cancel_a_queued_job_and_a_running_one(service_url):
    svc = wq.Service(service_url)
    # One worker: a long job holds it, so the next one waits in the queue.
    long = "OPENQASM 3.0; qubit[22] q;\n" + "".join(f"rx(0.{k % 9 + 1}) q[{k % 22}];\ncx q[{k % 22}], q[{(k + 7) % 22}];\n" for k in range(4000))
    a = svc.submit(long, observables=["Z0"], method="statevector")
    b = svc.submit(BELL, shots=10)
    assert svc.cancel(b["id"])["status"] == "cancelled"
    assert svc.status(b["id"])["status"] == "cancelled"
    with pytest.raises(wq.ServiceError, match="cancelled"):
        svc.result(b["id"], wait=5)
    # The long one, queued or running, is cancelled too, and stays so.
    assert svc.cancel(a["id"])["status"] == "cancelled"
    time.sleep(0.5)
    assert svc.status(a["id"])["status"] == "cancelled"
    with pytest.raises(wq.ServiceError) as e:
        svc.cancel(a["id"])
    assert e.value.status == 409


def test_a_job_sent_again_under_its_key_is_the_same_job(service_url):
    svc = wq.Service(service_url)
    a = svc.submit(BELL, shots=10, idempotency_key="py-key-1")
    b = svc.submit(BELL, shots=10, idempotency_key="py-key-1")
    assert a["id"] == b["id"]
    assert svc.run(BELL, shots=10, idempotency_key="py-key-1", wait=120, poll=0.1).status == "done"
    assert svc.submit(BELL, shots=10)["id"] != a["id"], "a fresh key is a new job"
    with pytest.raises(wq.ServiceError) as e:
        svc.submit(BELL, shots=11, idempotency_key="py-key-1")
    assert e.value.status == 422 and e.value.message.endswith("[J034]")
    with pytest.raises(wq.ServiceError) as e:
        svc.submit(BELL, shots=10, idempotency_key="has space")
    assert e.value.status == 400 and e.value.message.endswith("[J034]")


def test_the_service_refuses_what_the_engine_refuses_with_its_words(service_url):
    svc = wq.Service(service_url)
    with pytest.raises(wq.ServiceError) as e:
        svc.run("OPENQASM 3.0;\nqubit[1] q;\nfoo q[0];\n")
    assert e.value.status == 400
    with pytest.raises(wq.EngineError) as here:
        wq.run("OPENQASM 3.0;\nqubit[1] q;\nfoo q[0];\n", key=DEMO)
    assert e.value.message == here.value.message


def test_a_wasmq_envelope_and_a_qudit_job_run_on_the_service(service_url):
    s = wq.Service(service_url, timeout=60)
    job = s.wasmq("v2-t-toffoli", wait=240, poll=0.2)
    assert (job["status"], job["method"]) == ("done", "wasmq")
    assert wq.verify(job).valid
    w = wq.wasmq_verify(job, "v2-t-toffoli")
    assert w["valid"] and w["reproduced"] is True, w
    run = job["result"]["energy"]["run"]
    assert (run.get("joules_micro") or 0) >= 1 or (set(run) == {"unmetered"} and run["unmetered"]), run
    e = next(x for x in wq.wasmq_examples() if x["name"] == "surface3-decode64")
    by_bytes = s.wasmq(e["envelope"], wait=240, poll=0.2)
    assert by_bytes["result"]["counts"] == wq.wasmq(e["envelope"], key=DEMO)["result"]["counts"]
    q = s.qudit({"d": 6, "qudits": 2, "gates": [{"gate": "fourier", "qudits": [0]}, {"gate": "csum", "qudits": [0, 1]}], "shots": 120}, wait=240, poll=0.2)
    assert q["result"]["schema"] == "wai.quantum.qudit/1" and wq.qudit_check(q)["verified"]
    with pytest.raises(wq.ServiceError) as err:
        s.qudit({"d": 5, "qudits": 1, "gates": []})
    assert err.value.status == 422 and "J091" in str(err.value)