#!/usr/bin/env python3 """ AI Room Designer — Stage 3 tests (LOCAL RIG ONLY, $0 spend, no network). A. Gemini request/response path — fake transport captures the exact request (URL, header name, payload shape) and returns a canned Gemini-style response; the backend must decode the image out of it. Proves wiring without a key. B. QC retry — scripted backend returns a REFRAMED image first (the known Stage 2 Budget Refresh v1 failure pattern: arch collapsed to 0.102 when the camera moved), gate FAILs it, the loop regenerates WITH the framing lock and passes. Uses the real test5 kitchen photo. C. Retry exhaustion — backend always reframes: after 1 + max_retries attempts the best attempt is surfaced and the job is FLAGGED. D. Dry-run end-to-end through the real HTTP server (mock provider): /api/config, /api/generate, app 200, demos.json 6 entries. Run: python3 tests/test_stage3.py """ import base64 import io import json import shutil import subprocess import sys import time import urllib.request from pathlib import Path ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(ROOT / "pipeline")) import config # noqa: E402 import generate as gen # noqa: E402 from backends import GeminiBackend, MockBackend # noqa: E402 import pipeline as pl # noqa: E402 from PIL import Image # noqa: E402 TMP = ROOT / "tests" / "_stage3_tmp" TMP.mkdir(parents=True, exist_ok=True) fails = [] def check(name, ok, detail=""): print(("PASS " if ok else "FAIL ") + name + (f" — {detail}" if detail else "")) if not ok: fails.append(name) def fresh_job(src_dir, mode, name): d = TMP / name if d.exists(): shutil.rmtree(d) d.mkdir(parents=True) shutil.copy(src_dir / "before.jpg", d / "before.jpg") meta = {"mode": mode, "mode_label": mode} (d / "job.json").write_text(json.dumps(meta)) return d # ---------------------------------------------------------------- A: Gemini -- captured = {} def fake_transport(url, headers, payload): captured.update(url=url, headers=headers, payload=payload) buf = io.BytesIO() Image.new("RGB", (64, 64), (200, 180, 150)).save(buf, format="PNG") return {"candidates": [{"content": {"parts": [ {"text": "done"}, {"inline_data": {"mime_type": "image/png", "data": base64.b64encode(buf.getvalue()).decode()}}]}}], "usageMetadata": {"totalTokenCount": 1234}} be = GeminiBackend(key="TEST-KEY-NOT-REAL", transport=fake_transport) out = be.edit(str(ROOT / "tests/test2-living/before.jpg"), "PROMPT", str(TMP / "a.png")) check("A1 backend returns an image file", Path(out).exists()) check("A2 URL targets configured model", captured["url"].endswith(f"/models/{config.DEFAULT_MODEL}:generateContent"), captured["url"]) check("A3 key travels in x-goog-api-key header only", captured["headers"].get("x-goog-api-key") == "TEST-KEY-NOT-REAL" and "TEST-KEY" not in json.dumps(captured["payload"])) parts = captured["payload"]["contents"][0]["parts"] check("A4 payload = inline image + prompt text", "inline_data" in parts[0] and parts[1].get("text") == "PROMPT") check("A5 payload asks for IMAGE modality + configured size", captured["payload"]["generationConfig"]["responseModalities"] == ["TEXT", "IMAGE"] and captured["payload"]["generationConfig"]["imageConfig"]["imageSize"] == "1K") check("A6 usage metadata captured", be.last_usage == {"totalTokenCount": 1234}) # ------------------------------------------------------------- B: QC retry -- job = fresh_job(ROOT / "tests/test5-kitchen-budget", "budget-refresh", "retry") res = gen.generate(job, backend=MockBackend(script="reframe,grade")) att = [a for a in res["attempts"] if "scores" in a] check("B1 first (reframed) attempt FAILs the gate", att[0]["verdict"].startswith("FAIL"), f"arch={att[0]['scores']['arch_edge_corr']}") check("B2 retry fired with the framing lock", len(att) == 2 and att[1]["framing_lock"] is True) check("B3 retry PASSes and job completes unflagged", res["status"] == "done" and res["flagged"] is False and att[1]["verdict"].startswith("PASS"), f"arch={att[1]['scores']['arch_edge_corr']}") check("B4 job finished with qc/plan/composite", all((job / f).exists() for f in ("qc.json", "plan.json", "before_after.png"))) # ------------------------------------------- C: exhaustion -> best + flag -- job = fresh_job(ROOT / "tests/test5-kitchen-budget", "budget-refresh", "exhaust") res = gen.generate(job, backend=MockBackend(script="reframe")) att = [a for a in res["attempts"] if "scores" in a] check("C1 stops at 1 + max_retries attempts", len(att) == 1 + config.max_retries(), f"{len(att)} attempts") check("C2 best attempt surfaced + job FLAGGED", res["status"] == "done-flagged" and res["flagged"] is True and (job / "after_matched.png").exists()) qc = json.loads((job / "qc.json").read_text()) check("C3 flag recorded in qc.json gate_final", qc["gate_final"].startswith("FLAGGED")) # ------------------------------------------ D: HTTP end-to-end (dry-run) -- env = {"ROOM_DESIGNER_PROVIDER": "mock", "ROOM_DESIGNER_MOCK_SCRIPT": "reframe,grade", "PATH": "/usr/bin:/bin:/usr/local/bin"} proc = subprocess.Popen([sys.executable, str(ROOT / "app/server.py")], cwd=ROOT, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) try: import socket s = socket.socket() s.settimeout(1) for _ in range(50): try: s.connect(("127.0.0.1", 8765)); s.close(); break except OSError: time.sleep(0.2) cfg = json.loads(urllib.request.urlopen("http://127.0.0.1:8765/api/config").read()) check("D1 /api/config honest about dry-run", cfg["provider"] == "mock" and cfg["dry_run"] is True and cfg["live"] is False) code = urllib.request.urlopen("http://127.0.0.1:8765/app/").status demos = json.loads(urllib.request.urlopen("http://127.0.0.1:8765/app/demos.json").read()) check("D2 app 200 + 6 demo results load", code == 200 and len(demos) == 6) data_url = "data:image/jpeg;base64," + base64.b64encode( (ROOT / "tests/test5-kitchen-budget/before.jpg").read_bytes()).decode() req = urllib.request.Request("http://127.0.0.1:8765/api/generate", data=json.dumps({"image_data_url": data_url, "mode": "budget-refresh", "room": "kitchen"}).encode(), headers={"Content-Type": "application/json"}) r = json.loads(urllib.request.urlopen(req, timeout=60).read()) check("D3 upload generates end-to-end via HTTP", r["status"] == "done" and r["dry_run"] is True and r["retries_used"] == 1, f"status={r['status']} retries={r.get('retries_used')}") check("D4 result payload carries slider images + CAD plan", r["before"].endswith("before.jpg") and r["after"].endswith("after_matched.png") and r["plan"]["total_cad_high"] > 0 and r["qc"]["gate"].startswith("PASS")) finally: proc.terminate() print("\n" + ("ALL STAGE 3 TESTS PASSED" if not fails else f"FAILURES: {fails}")) sys.exit(1 if fails else 0)