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, 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"})
                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_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_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