settled-reach/tooling/test_assets.py

#!/usr/bin/env python3
"""`reach assets` connectors against a fake Gradio server (T-1290).

The real services cannot be used in a gate: Stable Audio and Trellis are kept
switched OFF (VRAM on tower-of-joy), and every Gemini call costs money. So the
decisions are exercised without performing them (D-263): a local HTTP server
plays the Gradio apps, and the tests assert on what each connector SENDS and
on how it classifies each kind of failure.

What is pinned:
- audio: submit payload [prompt, duration, steps, cfg], SSE `complete` → file
  download, SSE `error` → fails with the error, relative file URL resolution.
- trellis: the 5-call session sequence with one shared session_hash, the
  9-input image_to_3d payload, the 3-input extract_glb payload, the download.
- image: the generateContent body (aspect ratio as imageConfig, size as a
  prompt hint, inlineData for image-to-image) — built, never sent.
- endpoints: refused connection → "OFF" remedy; HTTP 500 → "rejected";
  401 → credentials; an unreachable non-tower service → network remedy.

Run: .venv/bin/python tooling/test_assets.py
"""

from __future__ import annotations

import contextlib
import io
import json
import socket
import sys
import tempfile
import threading
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from tooling.core.errors import ReachError  # noqa: E402
from tooling.domains.assets import audio, endpoints, image, trellis  # noqa: E402

AUDIO_BYTES = b"RIFF-fake-wav"
GLB_BYTES = b"glTF-fake-glb"


class FakeGradio(BaseHTTPRequestHandler):
    """Just enough of both Gradio apps. Records every request it sees."""

    seen: list[tuple[str, str, object]] = []
    sse_event = "complete"

    def log_message(self, *args):  # silence the default stderr access log
        pass

    def _send(self, code: int, body: bytes, ctype: str = "application/json") -> None:
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _json(self, obj, code: int = 200) -> None:
        self._send(code, json.dumps(obj).encode())

    def do_GET(self):
        FakeGradio.seen.append(("GET", self.path, None))
        if self.path == "/config":
            return self._json({"version": "4.44", "dependencies": [{"api_name": "generate_audio"}, {"api_name": "js_x"}]})
        if self.path == "/info":
            return self._json({"named_endpoints": {"/image_to_3d": {}, "/extract_glb": {}}})
        if self.path.startswith("/gradio_api/call/generate_audio/"):
            if FakeGradio.sse_event == "error":
                body = b"event: error\ndata: \"CUDA out of memory\"\n\n"
            else:
                body = b"event: heartbeat\ndata: null\n\nevent: complete\ndata: [{\"url\": \"/file=out.wav\"}]\n\n"
            return self._send(200, body, "text/event-stream")
        if self.path == "/file=out.wav":
            return self._send(200, AUDIO_BYTES, "audio/wav")
        if self.path == "/file=model.glb":
            return self._send(200, GLB_BYTES, "model/gltf-binary")
        if self.path == "/boom":
            return self._json({"error": "kaboom"}, 500)
        if self.path == "/denied":
            return self._json({"error": "bad key"}, 401)
        return self._json({"error": "not found"}, 404)

    def do_POST(self):
        length = int(self.headers.get("Content-Length", 0))
        raw = self.rfile.read(length)
        try:
            body = json.loads(raw)
        except (json.JSONDecodeError, UnicodeDecodeError):
            body = raw
        FakeGradio.seen.append(("POST", self.path, body))
        if self.path == "/gradio_api/call/generate_audio":
            return self._json({"event_id": "ev123"})
        if self.path == "/upload":
            return self._json(["/tmp/gradio/uploaded.png"])
        if self.path == "/api/preprocess_image_1":
            return self._json({"data": [{"path": "/tmp/pre.png"}]})
        if self.path == "/api/get_seed":
            return self._json({"data": [4242]})
        if self.path == "/api/extract_glb":
            return self._json({"data": [{"url": "/file=model.glb"}, {"url": "/file=model.glb"}]})
        if self.path.startswith("/api/"):
            return self._json({"data": []})
        return self._json({"error": "not found"}, 404)


@contextlib.contextmanager
def fake_server():
    server = ThreadingHTTPServer(("127.0.0.1", 0), FakeGradio)
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    base = f"http://127.0.0.1:{server.server_address[1]}"
    real = endpoints.load_config
    endpoints.load_config = lambda: {"stable_audio_url": base, "trellis_url": base}
    FakeGradio.seen = []
    FakeGradio.sse_event = "complete"
    try:
        yield base
    finally:
        endpoints.load_config = real
        server.shutdown()


def _closed_port() -> int:
    with socket.socket() as s:
        s.bind(("127.0.0.1", 0))
        return s.getsockname()[1]


def quiet(fn, *args, **kwargs):
    with contextlib.redirect_stderr(io.StringIO()):
        return fn(*args, **kwargs)


def check(failures: list[str], cond: bool, message: str) -> None:
    if not cond:
        failures.append(message)


def test_audio(failures: list[str]) -> None:
    with fake_server(), tempfile.TemporaryDirectory() as tmp:
        out = str(Path(tmp) / "x.wav")
        result = quiet(audio.generate, "wind over dunes", duration=5, steps=40, cfg=6.5, output=out, timeout=20)
        posts = [s for s in FakeGradio.seen if s[0] == "POST"]
        check(failures, posts == [("POST", "/gradio_api/call/generate_audio", {"data": ["wind over dunes", 5, 40, 6.5]})],
              f"audio: submit payload changed: {posts}")
        check(failures, Path(out).read_bytes() == AUDIO_BYTES, "audio: downloaded bytes differ")
        check(failures, result["size_bytes"] == len(AUDIO_BYTES) and result["file"] == out,
              f"audio: result keys wrong: {result}")

        FakeGradio.sse_event = "error"
        try:
            quiet(audio.generate, "x", output=out, timeout=20)
            failures.append("audio: an SSE error event did not fail the generation")
        except ReachError as exc:
            check(failures, "CUDA out of memory" in exc.message, f"audio: error detail lost: {exc.message}")

        health = quiet(audio.health)
        check(failures, health["api_endpoints"] == ["generate_audio"], f"audio: js_ endpoints not filtered: {health}")


def test_audio_urls(failures: list[str]) -> None:
    base = "http://h:1"
    cases = {
        "abs": ([{"url": "http://other/f.wav"}], "http://other/f.wav"),
        "rooted": ([{"path": "/file=a.wav"}], "http://h:1/file=a.wav"),
        "bare": (["tmp/a.wav"], "http://h:1/file=tmp/a.wav"),
        "wrapped": ({"data": [{"url": "/x.wav"}]}, "http://h:1/x.wav"),
        "empty": ([], None),
    }
    for name, (payload, want) in cases.items():
        got = audio.extract_file_url(payload, base)
        check(failures, got == want, f"audio url {name}: {got!r} != {want!r}")


def test_trellis(failures: list[str]) -> None:
    with fake_server(), tempfile.TemporaryDirectory() as tmp:
        png = Path(tmp) / "crate.png"
        png.write_bytes(b"\x89PNG-fake")
        out = str(Path(tmp) / "crate.glb")
        result = quiet(trellis.generate, str(png), output=out, simplify=0.9, texture_size=512, seed=7)

        posts = [(p, b) for m, p, b in FakeGradio.seen if m == "POST"]
        paths = [p for p, _ in posts]
        want = ["/api/start_session", "/upload", "/api/preprocess_image_1", "/api/get_seed",
                "/api/image_to_3d", "/api/extract_glb"]
        check(failures, paths == want, f"trellis: call sequence changed: {paths}")
        sessions = {b.get("session_hash") for p, b in posts if p != "/upload" and isinstance(b, dict)}
        check(failures, len(sessions) == 1 and None not in sessions, f"trellis: session not shared: {sessions}")
        by_path = dict(posts)
        check(failures, by_path["/api/get_seed"]["data"] == [True, 7], "trellis: get_seed payload changed")
        check(failures, by_path["/api/image_to_3d"]["data"] ==
              [{"path": "/tmp/pre.png"}, [], False, 4242, 7.5, 12, 3.0, 12, "stochastic"],
              f"trellis: image_to_3d payload changed: {by_path['/api/image_to_3d']['data']}")
        check(failures, by_path["/api/extract_glb"]["data"] == [None, 0.9, 512],
              "trellis: extract_glb payload changed")
        check(failures, Path(out).read_bytes() == GLB_BYTES, "trellis: downloaded bytes differ")
        check(failures, result["seed"] == 4242, "trellis: server seed not reported")

        # batch: one present, one missing input -> summary counts, no sleeping
        plan = trellis.batch_plan(Path(tmp), Path(tmp) / "out", ["crate", "ghost"])
        check(failures, [p[0] for p in plan] == ["crate", "ghost"], "trellis: batch plan ignores --names")
        summary = quiet(trellis.batch, tmp, str(Path(tmp) / "out"), names=["crate", "ghost"],
                        cooldown=0, max_retries=1, retry_delay=0)
        check(failures, (summary["success"], summary["failed"]) == (1, 1),
              f"trellis: batch counts wrong: {summary}")


def test_image_body(failures: list[str]) -> None:
    body = image.build_request_body("a lighthouse", aspect_ratio="16:9", image_size="2K")
    check(failures, body["contents"][0]["parts"] == [{"text": "a lighthouse Resolution: 2K."}],
          f"image: prompt/size hint changed: {body}")
    check(failures, body["generationConfig"] == {"responseModalities": ["TEXT", "IMAGE"],
                                                 "imageConfig": {"aspectRatio": "16:9"}},
          f"image: generationConfig changed: {body['generationConfig']}")
    with tempfile.TemporaryDirectory() as tmp:
        src = Path(tmp) / "ref.jpg"
        src.write_bytes(b"jpegdata")
        parts = image.build_request_body("x", input_image=str(src))["contents"][0]["parts"]
        check(failures, parts[0]["inlineData"]["mimeType"] == "image/jpeg" and len(parts) == 2,
              f"image: image-to-image part wrong: {parts}")
    try:
        image.build_request_body("x", input_image="/nope/missing.png")
        failures.append("image: a missing input image did not fail")
    except ReachError:
        pass


def test_classification(failures: list[str]) -> None:
    port = _closed_port()
    real = endpoints.load_config
    endpoints.load_config = lambda: {"stable_audio_url": f"http://127.0.0.1:{port}"}
    try:
        quiet(audio.health)
        failures.append("endpoints: a refused connection did not fail")
    except ReachError as exc:
        check(failures, "OFF" in exc.message and "do not restart" in (exc.fix or ""),
              f"endpoints: refused connection not reported as OFF: {exc.message} / {exc.fix}")
    finally:
        endpoints.load_config = real

    try:
        endpoints.call(urllib.request.Request(f"http://127.0.0.1:{port}/x"), service="Gemini API",
                       what="t", timeout=2)
        failures.append("endpoints: unreachable cloud API did not fail")
    except ReachError as exc:
        check(failures, "network" in (exc.fix or ""), f"endpoints: cloud outage blamed on VRAM: {exc.fix}")

    with fake_server() as base:
        for path, needle in (("/boom", "check the arguments"), ("/denied", "credentials")):
            try:
                endpoints.call(urllib.request.Request(base + path), service="audio", what="t", timeout=5)
                failures.append(f"endpoints: HTTP error at {path} did not fail")
            except ReachError as exc:
                check(failures, needle in (exc.fix or "") and "rejected" in exc.message,
                      f"endpoints: {path} misclassified: {exc.message} / {exc.fix}")

    try:
        endpoints.get_api_key("SR_TEST_SURELY_UNSET_KEY")
        failures.append("endpoints: a missing API key did not fail")
    except ReachError as exc:
        check(failures, "tracked" in (exc.fix or ""), "endpoints: missing-key remedy lost the never-commit warning")


def main() -> int:
    failures: list[str] = []
    for test in (test_audio, test_audio_urls, test_trellis, test_image_body, test_classification):
        test(failures)
    if failures:
        print("test_assets: FAIL", file=sys.stderr)
        for failure in failures:
            print(f"  - {failure}", file=sys.stderr)
        return 1
    print("test_assets: OK — audio, trellis and image payloads unchanged; OFF, rejected, "
          "credentials and network failures each named correctly; nothing sent to a real service")
    return 0


if __name__ == "__main__":
    sys.exit(main())