226 lines
9.7 KiB
Python
226 lines
9.7 KiB
Python
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from PIL import Image
|
|
|
|
import screen_designs
|
|
|
|
|
|
class ScreenDesignsTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp_dir = tempfile.TemporaryDirectory()
|
|
self.root = Path(self.temp_dir.name)
|
|
|
|
def tearDown(self):
|
|
self.temp_dir.cleanup()
|
|
|
|
def test_catalogue_is_well_formed(self):
|
|
names = [screen["name"] for screen in screen_designs.SCREENS]
|
|
self.assertEqual(len(names), len(set(names)),
|
|
"screen names must be unique")
|
|
refs = [screen["ref"] for screen in screen_designs.SCREENS
|
|
if screen["ref"]]
|
|
self.assertEqual(len(refs), len(set(refs)),
|
|
"screenshot refs must be unique")
|
|
for screen in screen_designs.SCREENS:
|
|
self.assertTrue(screen["prompt"].strip())
|
|
self.assertIn(screen["kind"], ("scene", "ui"))
|
|
self.assertTrue(screen["regions"])
|
|
for region in screen["regions"]:
|
|
x0, y0, x1, y1, label = region
|
|
self.assertLess(x0, x1)
|
|
self.assertLess(y0, y1)
|
|
self.assertTrue(0 <= x0 <= 1 and 0 <= x1 <= 1)
|
|
self.assertTrue(0 <= y0 <= 1 and 0 <= y1 <= 1)
|
|
self.assertTrue(label.strip())
|
|
|
|
def test_draft_wireframe_matches_reference_dimensions(self):
|
|
output = self.root / "draft.png"
|
|
screen = screen_designs.SCREENS[0]
|
|
screen_designs._draw_wireframe(screen, output)
|
|
with Image.open(output) as image:
|
|
self.assertEqual(image.size,
|
|
(screen_designs.WIDTH, screen_designs.HEIGHT))
|
|
|
|
def test_catalogue_covers_the_shipped_screens(self):
|
|
names = [screen["name"] for screen in screen_designs.SCREENS]
|
|
required = [
|
|
"venue_overview", "targeting_npc", "access_inspection",
|
|
"radio_report", "late_night_avenue", "night_flashlight",
|
|
"radial_menu", "shift_review", "start_menu", "pause_overlay",
|
|
"settings_panel", "emergency_radio", "wand_screening",
|
|
"joke_shop_queue", "studio_arrival", "crowd_panic",
|
|
"responder_arrival", "body_camera_timeline", "escalation_hud",
|
|
]
|
|
for name in required:
|
|
self.assertIn(name, names)
|
|
self.assertGreaterEqual(len(screen_designs.SCREENS), len(required))
|
|
for screen in screen_designs.SCREENS:
|
|
self.assertTrue(screen["title"].strip())
|
|
|
|
def test_list_command_prints_the_catalogue(self):
|
|
import contextlib
|
|
import io
|
|
buffer = io.StringIO()
|
|
with contextlib.redirect_stdout(buffer):
|
|
code = screen_designs.main(["list"])
|
|
self.assertEqual(code, 0)
|
|
output = buffer.getvalue()
|
|
self.assertIn("escalation_hud", output)
|
|
self.assertIn("wand_screening", output)
|
|
self.assertIn("Total:", output)
|
|
|
|
def test_judge_dry_run_builds_requests_without_network(self):
|
|
import contextlib
|
|
import io
|
|
if not screen_designs.SCREENSHOTS_DIR.is_dir():
|
|
self.skipTest("no screenshots checked out")
|
|
buffer = io.StringIO()
|
|
with contextlib.redirect_stdout(buffer):
|
|
code = screen_designs.main(["judge", "--dry-run"])
|
|
# A dry run must not call the API and must not crash on missing verdicts.
|
|
self.assertEqual(code, 0)
|
|
output = buffer.getvalue()
|
|
self.assertIn("dry run", output)
|
|
self.assertIn("Summary:", output)
|
|
|
|
def test_compose_prompt_appends_kind_specific_style(self):
|
|
scene = {"kind": "scene", "prompt": "a street"}
|
|
ui = {"kind": "ui", "prompt": "a panel"}
|
|
self.assertIn(screen_designs.SCENE_STYLE,
|
|
screen_designs._compose_prompt(scene))
|
|
self.assertIn(screen_designs.UI_STYLE,
|
|
screen_designs._compose_prompt(ui))
|
|
self.assertIn("16:9", screen_designs._compose_prompt(scene))
|
|
|
|
def test_mae_is_zero_for_identical_images(self):
|
|
path = self.root / "identical.png"
|
|
Image.new("RGB", (1280, 720), (40, 80, 120)).save(path)
|
|
left = screen_designs._load_rgb(path)
|
|
right = screen_designs._load_rgb(path)
|
|
self.assertAlmostEqual(screen_designs._mae(left, right), 0.0, places=1)
|
|
|
|
def test_histogram_distance_is_zero_for_identical_images(self):
|
|
path = self.root / "identical.png"
|
|
Image.new("RGB", (1280, 720), (40, 80, 120)).save(path)
|
|
left = screen_designs._load_rgb(path)
|
|
right = screen_designs._load_rgb(path)
|
|
self.assertAlmostEqual(
|
|
screen_designs._histogram_distance(left, right), 0.0, places=3)
|
|
|
|
def test_blank_frame_is_flagged_as_bad(self):
|
|
reference = self.root / "ref.png"
|
|
capture = self.root / "cap.png"
|
|
Image.new("RGB", (1280, 720), (200, 200, 200)).save(reference)
|
|
Image.new("RGB", (1280, 720), (0, 0, 0)).save(capture)
|
|
report = screen_designs.compare_screenshot(reference, capture)
|
|
self.assertEqual(report["quality"], "bad")
|
|
self.assertTrue(any("near-blank" in e for e in report["errors"]))
|
|
|
|
def test_wrong_aspect_is_flagged_as_bad(self):
|
|
reference = self.root / "ref.png"
|
|
capture = self.root / "cap.png"
|
|
Image.new("RGB", (1280, 720), (100, 120, 140)).save(reference)
|
|
Image.new("RGB", (640, 640), (100, 120, 140)).save(capture)
|
|
report = screen_designs.compare_screenshot(reference, capture)
|
|
self.assertEqual(report["quality"], "bad")
|
|
self.assertTrue(any("aspect" in e for e in report["errors"]))
|
|
|
|
|
|
class VisionJudgeTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp_dir = tempfile.TemporaryDirectory()
|
|
self.root = Path(self.temp_dir.name)
|
|
self.capture = self.root / "capture.png"
|
|
self.design = self.root / "design.png"
|
|
Image.new("RGB", (64, 36), (20, 30, 40)).save(self.capture)
|
|
Image.new("RGB", (64, 36), (200, 190, 180)).save(self.design)
|
|
self.screen = screen_designs.SCREENS[0]
|
|
|
|
def tearDown(self):
|
|
self.temp_dir.cleanup()
|
|
|
|
def test_judge_request_embeds_both_images(self):
|
|
body = screen_designs.build_judge_request(
|
|
"model-x", self.screen, self.capture, self.design)
|
|
self.assertEqual(body["model"], "model-x")
|
|
self.assertEqual(body["temperature"], 0)
|
|
self.assertEqual(body["response_format"], {"type": "json_object"})
|
|
parts = body["messages"][0]["content"]
|
|
self.assertEqual(parts[0]["type"], "text")
|
|
images = [p for p in parts if p["type"] == "image_url"]
|
|
self.assertEqual(
|
|
len(images), 2, "design and capture are both attached")
|
|
for image in images:
|
|
self.assertTrue(image["image_url"]["url"].startswith(
|
|
"data:image/png;base64,"))
|
|
self.assertIn(self.screen["title"], parts[0]["text"])
|
|
self.assertIn("mockup", parts[0]["text"].lower())
|
|
|
|
def test_parse_plain_json_verdict(self):
|
|
payload = {"choices": [{"message": {"content":
|
|
'{"match": true, "confidence": 0.9, "differences": ["style"]}'}}]}
|
|
verdict = screen_designs.parse_judge_response(payload)
|
|
self.assertTrue(verdict["match"])
|
|
self.assertAlmostEqual(verdict["confidence"], 0.9)
|
|
self.assertEqual(verdict["differences"], ["style"])
|
|
|
|
def test_parse_fenced_json_verdict(self):
|
|
payload = {"choices": [{"message": {"content":
|
|
'Sure:\n```json\n{"match": false, "confidence": 0.7}\n```'}}]}
|
|
verdict = screen_designs.parse_judge_response(payload)
|
|
self.assertFalse(verdict["match"])
|
|
self.assertAlmostEqual(verdict["confidence"], 0.7)
|
|
|
|
def test_parse_content_parts(self):
|
|
payload = {"choices": [{"message": {"content": [
|
|
{"type": "text", "text": '{"match": true, "confidence": 0.5}'}]}}]}
|
|
verdict = screen_designs.parse_judge_response(payload)
|
|
self.assertTrue(verdict["match"])
|
|
|
|
def test_parse_non_json_is_unknown(self):
|
|
payload = {"choices": [{"message": {"content": "no idea"}}]}
|
|
verdict = screen_designs.parse_judge_response(payload)
|
|
self.assertIsNone(verdict["match"])
|
|
self.assertTrue(verdict["differences"])
|
|
|
|
def test_parse_missing_choices_is_unknown(self):
|
|
verdict = screen_designs.parse_judge_response({})
|
|
self.assertIsNone(verdict["match"])
|
|
|
|
def test_classify_thresholds(self):
|
|
self.assertEqual(screen_designs.classify(
|
|
{"match": True, "confidence": 0.1}), "match")
|
|
self.assertEqual(screen_designs.classify(
|
|
{"match": False, "confidence": 0.9}), "bad")
|
|
self.assertEqual(screen_designs.classify(
|
|
{"match": False, "confidence": 0.3}), "review")
|
|
self.assertEqual(screen_designs.classify(
|
|
{"match": None, "confidence": 0.0}), "unknown")
|
|
self.assertEqual(screen_designs.classify(
|
|
{"match": False, "confidence": 0.55}, threshold=0.5), "bad")
|
|
|
|
def test_env_file_does_not_override_existing(self):
|
|
env_path = self.root / ".env"
|
|
env_path.write_text('OPENROUTER_API_KEY=from-file\nOTHER=1\n',
|
|
encoding="utf-8")
|
|
import os
|
|
old = os.environ.get("OPENROUTER_API_KEY")
|
|
os.environ["OPENROUTER_API_KEY"] = "from-env"
|
|
os.environ.pop("OTHER", None)
|
|
try:
|
|
screen_designs._load_env_file(env_path)
|
|
self.assertEqual(os.environ["OPENROUTER_API_KEY"], "from-env")
|
|
self.assertEqual(os.environ.get("OTHER"), "1")
|
|
finally:
|
|
if old is None:
|
|
os.environ.pop("OPENROUTER_API_KEY", None)
|
|
else:
|
|
os.environ["OPENROUTER_API_KEY"] = old
|
|
os.environ.pop("OTHER", None)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|