scripts/verify_hf_stack.py

#!/usr/bin/env python3
"""
End-to-end API checks for the HF/Docker image (same stack as Spaces).

Usage:
  python scripts/verify_hf_stack.py [--base-url http://127.0.0.1:7860]
  HF_VERIFY_BASE_URL=https://your-space.hf.space python scripts/verify_hf_stack.py

Or run the full Docker probe (build, run container, hit API):
  python scripts/verify_hf_stack.py --docker

Deploy to the Hub (requires HF_TOKEN + HF_SPACE_REPO), then verify live:
  HF_TOKEN=... HF_SPACE_REPO=user/slug python scripts/publish_hf_space.py --verify

Prove catalog + kinks + recommendations on a deployed Space (no partner flow):
  HF_VERIFY_BASE_URL=https://owner-slug.hf.space HF_VERIFY_MIN_KINKS=1000 \\
    python scripts/verify_hf_stack.py --smoke
"""
from __future__ import annotations

import argparse
import json
import os
import re
import subprocess
import sys
import time
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen


def _get(url: str, headers: dict[str, str] | None = None) -> dict[str, Any]:
    req = Request(url, headers=headers or {})
    with urlopen(req, timeout=120) as resp:
        return json.loads(resp.read().decode())


def _get_text(url: str, headers: dict[str, str] | None = None) -> str:
    req = Request(url, headers=headers or {})
    with urlopen(req, timeout=120) as resp:
        return resp.read().decode()


def _assert_live_frontend_recs_error_shaping(root: str) -> None:
    """Follow index → app.js → api.js / discover.js and assert error-UI hardening is deployed."""
    if os.environ.get("HF_VERIFY_SKIP_FRONTEND", "").strip().lower() in ("1", "true", "yes", "on"):
        print("verify_hf_stack: skipping live frontend bundle check (HF_VERIFY_SKIP_FRONTEND=1)", file=sys.stderr)
        return
    base = root.rstrip("/")
    html = _get_text(f"{base}/")
    m = re.search(r"""src=["'](/frontend/app\.js\?v=\d+)["']""", html)
    if not m:
        raise RuntimeError("GET /: expected <script src=\"/frontend/app.js?v=…\"> in HTML")
    app_path = m.group(1)
    app_js = _get_text(f"{base}{app_path}")
    m_api = re.search(r"""["']\./api\.js\?v=(\d+)["']""", app_js)
    if not m_api:
        m_api = re.search(r"""["'][^"']*api\.js\?v=(\d+)["']""", app_js)
    if not m_api:
        raise RuntimeError(f"{app_path}: could not find api.js?v=… import")
    api_path = f"/frontend/api.js?v={m_api.group(1)}"
    api_js = _get_text(f"{base}{api_path}")
    for needle in ("_looksLikeHtmlDocument", "_messageForFailedRequest"):
        if needle not in api_js:
            raise RuntimeError(
                f"GET {api_path}: missing {needle!r} — stale UI bundle (redeploy Space / hard-refresh cache)"
            )
    m_disc = re.search(r"""["']\./discover\.js\?v=(\d+)["']""", app_js)
    if not m_disc:
        m_disc = re.search(r"""discover\.js\?v=(\d+)""", app_js)
    if not m_disc:
        raise RuntimeError(f"{app_path}: could not find discover.js?v=… import")
    disc_path = f"/frontend/discover.js?v={m_disc.group(1)}"
    disc_js = _get_text(f"{base}{disc_path}")
    if "Could not load recommendations (server error)" not in disc_js:
        raise RuntimeError(
            f"GET {disc_path}: missing recs error fallback string — stale discover.js (redeploy Space)"
        )
    print(f"verify_hf_stack: frontend recs-error bundle ok ({api_path}, {disc_path})", file=sys.stderr)


def _peek_coep(base: str) -> str:
    url = f"{base.rstrip('/')}/health"
    req = Request(url)
    with urlopen(req, timeout=120) as resp:
        return (resp.headers.get("Cross-Origin-Embedder-Policy") or "").strip()


def _assert_coep_credentialless(base: str) -> None:
    """COEP must allow cross-origin catalog images (see api middleware + tests/test_coep_headers.py).

    On ``*.hf.space``, retry for a while: after ``publish_hf_space`` pushes, the edge may still serve the
    previous image (``require-corp``) until the new container and README ``custom_headers`` propagate.
    """
    root = base.rstrip("/")
    if os.environ.get("HF_VERIFY_SKIP_COEP", "").strip().lower() in ("1", "true", "yes", "on"):
        print("verify_hf_stack: skipping COEP check (HF_VERIFY_SKIP_COEP=1)", file=sys.stderr)
        return
    if "hf.space" not in root:
        coep = _peek_coep(root)
        assert coep == "credentialless", (
            f"Cross-Origin-Embedder-Policy is {coep!r}; expected credentialless for catalog <img> hosts without CORP"
        )
        return

    retries = int(os.environ.get("HF_VERIFY_COEP_RETRIES", "90") or "90")
    interval = float(os.environ.get("HF_VERIFY_COEP_INTERVAL_S", "10") or "10")
    last = ""
    for attempt in range(retries):
        last = _peek_coep(root)
        if last == "credentialless":
            if attempt:
                print(f"verify_hf_stack: COEP credentialless after {attempt} wait(s)", file=sys.stderr)
            return
        if attempt == 0 or (attempt + 1) % 6 == 0:
            print(
                f"verify_hf_stack: COEP is {last!r} (want credentialless); "
                f"waiting for deploy ({attempt + 1}/{retries}, {interval:.0f}s)…",
                file=sys.stderr,
            )
        time.sleep(interval)
    assert last == "credentialless", (
        f"Cross-Origin-Embedder-Policy stayed {last!r} after {retries * interval:.0f}s; "
        "expected credentialless (see README custom_headers + api.py). "
        "Set HF_VERIFY_SKIP_COEP=1 only to bypass."
    )


def _post(url: str, body: dict[str, Any], headers: dict[str, str] | None = None) -> dict[str, Any]:
    data = json.dumps(body).encode()
    h = {"Content-Type": "application/json", **(headers or {})}
    req = Request(url, data=data, headers=h, method="POST")
    with urlopen(req, timeout=120) as resp:
        return json.loads(resp.read().decode())


def _wait_for_health(base: str, timeout_s: float) -> None:
    """Hugging Face Spaces may cold-start; retry /health before failing."""
    url = f"{base.rstrip('/')}/health"
    deadline = time.time() + timeout_s
    last_err: Exception | None = None
    while time.time() < deadline:
        try:
            h = _get(url)
            if h.get("ok") is True:
                return
        except (URLError, HTTPError, TimeoutError, OSError) as e:
            last_err = e
        time.sleep(3.0)
    raise last_err or TimeoutError(f"/health did not return ok within {timeout_s:.0f}s: {url}")


def _require(cond: bool, msg: str) -> None:
    if not cond:
        raise RuntimeError(msg)


def run_smoke_checks(base: str) -> None:
    """Public + authenticated smoke: health stats, ``/kinks``, search, new user, one play, recommendations."""
    root = base.rstrip("/")
    if "hf.space" in root:
        _wait_for_health(root, timeout_s=float(os.environ.get("HF_VERIFY_HEALTH_TIMEOUT_S", "3600")))
    else:
        _wait_for_health(root, timeout_s=120.0)

    health = _get(f"{root}/health")
    _require(health.get("ok") is True, f"health not ok: {health!r}")
    stats = health.get("stats") or {}
    nk = int(stats.get("kinks") or 0)
    min_k = int(os.environ.get("HF_VERIFY_MIN_KINKS", "1") or "1")
    _require(nk >= min_k, f"catalog too small for smoke: kinks={nk}, need>={min_k}")
    print(f"verify_hf_stack smoke: health ok, kinks_in_catalog={nk}", file=sys.stderr)

    _assert_coep_credentialless(root)
    _assert_live_frontend_recs_error_shaping(root)

    probe = _get(f"{root}/health/catalog-sample")
    _require(probe.get("ok") is True, f"catalog-sample not ok: {probe!r}")
    sample_id = probe.get("sample_kink_id")
    _require(bool(sample_id), f"catalog-sample missing kink id: {probe!r}")
    detail = _get(f"{root}/kinks/{sample_id}")
    _require(detail.get("id") == sample_id, f"kink detail mismatch: {detail!r}")
    print(f"verify_hf_stack smoke: catalog-sample id={sample_id!r}", file=sys.stderr)

    u = _post(f"{root}/users", {})
    _require(bool(u.get("id") and u.get("private_token")), f"create user bad payload: {u!r}")
    tok = {"x-private-token": u["private_token"]}
    play = _post(
        f"{root}/users/{u['id']}/plays",
        {"kink_id": sample_id, "interest_state": "love", "directions": ["together"]},
        headers=tok,
    )
    _require(bool(play.get("plays", {}).get(sample_id)), f"play not saved: {play!r}")
    rec = _get(f"{root}/users/{u['id']}/recommendations?limit=8", headers=tok)
    nrec = len(rec.get("items", []))
    _require(nrec >= 1, f"recommendations empty: {rec!r}")
    print(f"verify_hf_stack smoke: recommendations count={nrec} for new user", file=sys.stderr)
    print("verify_hf_stack: SMOKE OK", file=sys.stderr)


def run_checks(base: str) -> None:
    root = base.rstrip("/")
    # Direct app host (proxied Gradio/FastAPI). The hub page huggingface.co/spaces/... is not the API origin.
    if "hf.space" in root:
        # First boot may download a multi-GB catalog from the Hub before /health responds.
        _wait_for_health(root, timeout_s=float(os.environ.get("HF_VERIFY_HEALTH_TIMEOUT_S", "3600")))
        expect_k = int(os.environ.get("HF_VERIFY_EXPECT_KINKS", "0") or "0")
        min_k = int(os.environ.get("HF_VERIFY_MIN_KINKS", "0") or "0")
        if expect_k > 0 or min_k > 0:
            h = _get(f"{root}/health")
            nk = int((h.get("stats") or {}).get("kinks") or 0)
            if expect_k > 0:
                assert nk == expect_k, f"catalog count mismatch: kinks={nk}, expected={expect_k}"
            assert nk >= min_k, f"catalog too small: kinks={nk}, need>={min_k}"
    else:
        # Local/docker: catalog bootstrap may run in a background thread after bind.
        _wait_for_health(root, timeout_s=120.0)

    _assert_coep_credentialless(root)
    _assert_live_frontend_recs_error_shaping(root)

    a = _post(f"{root}/users", {})
    b = _post(f"{root}/users", {})
    assert a["id"] and a["private_token"]
    assert b["id"] and b["private_token"]

    tok_a = {"x-private-token": a["private_token"]}
    tok_b = {"x-private-token": b["private_token"]}

    _post(
        f"{root}/users/{a['id']}/partners",
        {"partner_id": b["id"]},
        headers=tok_a,
    )
    u_b = _get(f"{root}/users/{b['id']}", headers=tok_b)
    assert a["id"] in (u_b.get("incoming_partner_requests") or []), u_b

    u_b2 = _post(
        f"{root}/users/{b['id']}/partner-requests/accept",
        {"from_user_id": a["id"]},
        headers=tok_b,
    )
    assert a["id"] in u_b2["partners"], u_b2

    grp = _post(
        f"{root}/users/{a['id']}/partner-groups",
        {"name": "Verify", "member_ids": [b["id"]]},
        headers=tok_a,
    )
    assert grp.get("id")

    klist = _get(f"{root}/kinks?limit=1")
    kitems = klist.get("items") or []
    assert kitems, ("no kinks in catalog", klist)
    sample_id = kitems[0].get("id")
    assert sample_id, kitems[0]

    play = _post(
        f"{root}/users/{a['id']}/plays",
        {"kink_id": sample_id, "interest_state": "love", "directions": ["together"]},
        headers=tok_a,
    )
    assert play["plays"].get(sample_id), play

    rec = _get(
        f"{root}/users/{a['id']}/recommendations?limit=5",
        headers=tok_a,
    )
    assert len(rec.get("items", [])) >= 1, rec

    sr = _get(f"{root}/search?q=massage&limit=5")
    assert len(sr.get("items", [])) >= 1, sr

    print("verify_hf_stack: OK", file=sys.stderr)


def docker_probe(image: str) -> int:
    port = "7860"
    cmd = [
        "docker",
        "run",
        "--rm",
        "-d",
        "-p",
        f"127.0.0.1:0:{port}",
        "-e",
        "KINK_STORE_PATH=/tmp/hf_verify_store.db",
        # Image defaults require Hub dataset; use bundled 5-kink seed for a fast local probe.
        "-e",
        "KINK_HF_REQUIRE_FULL_CATALOG=0",
        "-e",
        "KINK_HF_DATASET_REPO=",
        image,
    ]
    cid = subprocess.check_output(cmd, text=True).strip()
    try:
        port_line = subprocess.check_output(["docker", "port", cid, f"{port}/tcp"], text=True).strip()
        # e.g. 0.0.0.0:32768
        host_port = port_line.rsplit(":", 1)[-1]
        base = f"http://127.0.0.1:{host_port}"
        deadline = time.time() + 180
        last_err: Exception | None = None
        while time.time() < deadline:
            try:
                _get(f"{base}/health")
                last_err = None
                break
            except (URLError, HTTPError, TimeoutError, OSError) as e:
                last_err = e
                time.sleep(1)
        if last_err is not None:
            raise last_err
        run_checks(base)
        return 0
    finally:
        subprocess.run(["docker", "stop", cid], capture_output=True)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--base-url", default=os.environ.get("HF_VERIFY_BASE_URL", ""))
    ap.add_argument(
        "--require-live",
        action="store_true",
        help="Require HF_LIVE_URL (or --base-url) for a deployed Space; exit 2 if missing",
    )
    ap.add_argument("--docker", action="store_true", help="Build image and run docker probe")
    ap.add_argument(
        "--smoke",
        action="store_true",
        help="Shorter live checks: health + catalog size + /kinks + /search + new user + recommendations (no partners)",
    )
    ap.add_argument("--image", default="kink-cli-hf:latest")
    args = ap.parse_args()
    if args.require_live:
        live = (args.base_url or os.environ.get("HF_LIVE_URL", "")).strip()
        if not live:
            print(
                "Set HF_LIVE_URL or pass --base-url to the deployed Space (*.hf.space).",
                file=sys.stderr,
            )
            return 2
        args.base_url = live
    if args.docker:
        subprocess.run(
            ["docker", "build", "-t", args.image, "."],
            cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
            check=True,
        )
        return docker_probe(args.image)
    base = args.base_url or "http://127.0.0.1:7860"
    if args.smoke:
        run_smoke_checks(base)
    else:
        run_checks(base)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())