open-standards/wai/wai-quantum-service/tests/metrics.rs

//! `GET /metrics` of the service, run as it is deployed (`wai-quantum-service serve`) on loopback:
//! it reads strictly as the text format, every family present from the start; a job run moves the
//! jobs held by status and the requests by route, method and status; `HEAD` is answered as `GET`,
//! its length stated and without the text; another method is `405`; a request through a front is
//! a path the service does not have, `404`; and no job, client or key of those requests is a
//! label value in it.

use serde_json::{Value, json};
use std::io::{Read, Write};
use std::net::{TcpListener, TcpStream};
use std::process::{Command, Stdio};
use std::time::{Duration, Instant};
use wai_metrics::read::{Family, identifier_like, label_values, parse, value};

mod common;

use common::Running;

fn free_port() -> u16 {
    TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap().port()
}

/// One request, and its head and body as sent.
fn raw(port: u16, method: &str, path: &str, extra: &str, body: &str) -> (String, String) {
    let mut s = TcpStream::connect(("127.0.0.1", port)).unwrap();
    s.set_read_timeout(Some(Duration::from_secs(120))).unwrap();
    write!(s, "{method} {path} HTTP/1.1\r\nHost: 127.0.0.1\r\nContent-Type: application/json\r\nContent-Length: {}\r\n{extra}Connection: close\r\n\r\n{body}", body.len()).unwrap();
    let mut out = String::new();
    s.read_to_string(&mut out).unwrap();
    let (head, body) = out.split_once("\r\n\r\n").unwrap_or((out.as_str(), ""));
    (head.to_string(), body.to_string())
}

fn status(head: &str) -> u16 {
    head.split_whitespace().nth(1).and_then(|c| c.parse().ok()).unwrap_or(0)
}

fn header<'a>(head: &'a str, name: &str) -> Option<&'a str> {
    head.lines().find_map(|l| l.split_once(':').filter(|(n, _)| n.eq_ignore_ascii_case(name)).map(|(_, v)| v.trim()))
}

/// `GET /metrics`, read strictly, no label value naming a party.
fn scrape(port: u16) -> Vec<Family> {
    let (head, body) = raw(port, "GET", "/metrics", "", "");
    assert_eq!(status(&head), 200, "{head}");
    assert_eq!(header(&head, "content-type"), Some(wai_metrics::CONTENT_TYPE), "{head}");
    let f = parse(&body).unwrap_or_else(|e| panic!("the exposition does not read: {e:?}\n{body}"));
    for v in label_values(&f) {
        assert!(!identifier_like(v), "a label value that names a party: {v}\n{body}");
    }
    f
}

/// The service, and the directory of its store, removed once it has ended.
struct Service(Running, std::path::PathBuf);

impl Drop for Service {
    fn drop(&mut self) {
        self.0.stop();
        let _ = std::fs::remove_dir_all(&self.1);
    }
}

fn start() -> (Service, u16) {
    let dir = std::env::temp_dir().join(format!("wai-qsvc-metrics-{}", std::process::id()));
    let _ = std::fs::remove_dir_all(&dir);
    std::fs::create_dir_all(&dir).unwrap();
    let port = free_port();
    let child = Command::new(env!("CARGO_BIN_EXE_wai-quantum-service"))
        .arg("serve")
        .env("WAI_QSVC_ADDR", format!("127.0.0.1:{port}"))
        .env("WAI_QSVC_STORE", dir.join("store"))
        .env("WAI_QSVC_KEY", dir.join("signing.key"))
        .env("WAI_QSVC_JOB_GB", "2")
        .env("WAI_QSVC_BUDGET_GB", "4")
        .env("WAI_QSVC_WORKERS", "1")
        .env("WAI_QSVC_ENERGY", "0")
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .spawn()
        .expect("the service starts");
    let service = Service(Running::new(child), dir);
    let t = Instant::now();
    while TcpStream::connect(("127.0.0.1", port)).is_err() || status(&raw(port, "GET", "/v1/health", "", "").0) != 200 {
        assert!(t.elapsed() < Duration::from_secs(60), "the service did not answer");
        std::thread::sleep(Duration::from_millis(50));
    }
    (service, port)
}

#[test]
fn the_service_answers_metrics_that_move_with_its_jobs() {
    let (_service, port) = start();

    // Before any job: every family, the jobs at 0 in every status.
    let f = scrape(port);
    for s in ["queued", "running", "done", "failed", "cancelled", "interrupted"] {
        assert_eq!(value(&f, "wai_quantum_jobs", &[("status", s)]), Some(0.0), "{s}");
    }
    assert_eq!(value(&f, "wai_quantum_memory_held_bytes", &[]), Some(0.0));
    assert_eq!(value(&f, "wai_quantum_energy_unmetered", &[]), Some(1.0), "the service was started with energy metering off");
    assert_eq!(value(&f, "wai_quantum_restored_jobs", &[("outcome", "finished")]), Some(0.0), "an empty store");
    for pool in ["reading", "checks"] {
        for event in ["lent", "moved", "kept", "taken_back"] {
            assert_eq!(value(&f, "wai_quantum_lent_turns_total", &[("pool", pool), ("event", event)]), Some(0.0), "{pool} {event}: no turn lent yet");
        }
    }
    assert!(f.iter().any(|x| x.name == "wai_quantum_requests_total"));
    assert_eq!(value(&f, "wai_quantum_requests_total", &[("route", "/v1/health"), ("method", "GET"), ("code", "200")]).map(|v| v >= 1.0), Some(true), "the health checks while it started");

    // A job run: done, held, and its requests counted by route, not by its id.
    let bell = json!({ "qasm": "OPENQASM 3.0; qubit[2] q; bit[2] c; h q[0]; cx q[0], q[1]; c = measure q;", "shots": 100 });
    let (head, body) = raw(port, "POST", "/v1/run", "", &bell.to_string());
    let run = status(&head);
    let made: Value = serde_json::from_str(&body).unwrap_or(Value::Null);
    // On a loaded host the run may still be running when its wait ends: 202, then polled.
    assert!(run == 200 || run == 202, "{head}\n{body}");
    let id = made["id"].as_str().unwrap_or_else(|| panic!("the job's id: {body}")).to_string();
    let t = Instant::now();
    let mut polls = 0;
    loop {
        let (h, b) = raw(port, "GET", &format!("/v1/jobs/{id}?wait=20"), "", "");
        assert_eq!(status(&h), 200, "{h}\n{b}");
        polls += 1;
        let j: Value = serde_json::from_str(&b).unwrap_or(Value::Null);
        if j["status"] != "queued" && j["status"] != "running" {
            assert_eq!(j["status"], "done", "{b}");
            break;
        }
        assert!(t.elapsed() < Duration::from_secs(180), "the job did not finish: {b}");
    }
    let f = scrape(port);
    assert_eq!(value(&f, "wai_quantum_jobs", &[("status", "done")]), Some(1.0));
    let code = run.to_string();
    assert_eq!(value(&f, "wai_quantum_requests_total", &[("route", "/v1/run"), ("method", "POST"), ("code", code.as_str())]), Some(1.0));
    assert_eq!(value(&f, "wai_quantum_requests_total", &[("route", "/v1/jobs/{id}"), ("method", "GET"), ("code", "200")]), Some(f64::from(polls)), "each poll, by the route's pattern");
    assert_eq!(value(&f, "wai_quantum_requests_total", &[("route", "/metrics"), ("method", "GET"), ("code", "200")]), Some(1.0), "the scrape before this one");
    let text = raw(port, "GET", "/metrics", "", "").1;
    assert!(!text.contains(&id), "the job's id is in the exposition");

    // HEAD as GET, its length stated and without the text; another method 405.
    let (head, body) = raw(port, "HEAD", "/metrics", "", "");
    assert_eq!(status(&head), 200, "{head}");
    assert_eq!(header(&head, "content-type"), Some(wai_metrics::CONTENT_TYPE), "{head}");
    assert!(header(&head, "content-length").and_then(|n| n.parse::<usize>().ok()).is_some_and(|n| n > 0), "{head}");
    assert!(body.is_empty(), "HEAD was answered with the text: {body:?}");
    let (head, _) = raw(port, "POST", "/metrics", "", "{}");
    assert_eq!((status(&head), header(&head, "allow")), (405, Some("GET, HEAD")), "{head}");

    // Through a front, by any of the fields a front sets: a path the service does not have.
    for front in wai_metrics::FRONT_FIELDS {
        let (head, body) = raw(port, "GET", "/metrics", &format!("{front}: 192.0.2.7\r\n"), "");
        assert_eq!(status(&head), 404, "{front}: {head}");
        assert!(!body.contains("# HELP"), "{front}: a request through a front was answered the counters");
    }
    let f = scrape(port);
    assert_eq!(value(&f, "wai_quantum_requests_total", &[("route", "/metrics"), ("method", "GET"), ("code", "404")]), Some(wai_metrics::FRONT_FIELDS.len() as f64));
    assert_eq!(value(&f, "wai_quantum_requests_total", &[("route", "/metrics"), ("method", "HEAD"), ("code", "200")]), Some(1.0));
}