qwen-image/tour-comfy/test/test_scenery_case.py

#!/usr/bin/env python3
import argparse
import contextlib
import os
import re
import shutil
import sys
import tempfile
import traceback
from pathlib import Path
from unittest.mock import patch

IMAGE_EXTS = [".png", ".jpg", ".jpeg", ".webp"]


def slug(value: str) -> str:
    return re.sub(r"[^a-zA-Z0-9_.-]+", "_", value).strip("_") or "case"


def find_image(folder: Path, stem: str) -> Path | None:
    for file in sorted(folder.iterdir()):
        if file.is_file() and file.stem.lower() == stem.lower() and file.suffix.lower() in IMAGE_EXTS:
            return file
    return None


def find_cases(root: Path) -> list[dict]:
    cases = []

    for folder, _, _ in os.walk(root):
        folder = Path(folder)

        prompt_file = folder / "prompt.md"
        image1 = find_image(folder, "image1")
        image2 = find_image(folder, "image2")

        if prompt_file.exists() and image1 and image2:
            cases.append({
                "folder": folder,
                "image1": image1,
                "image2": image2,
                "prompt": prompt_file,
            })

    return cases


def expected_prompt_refinement(prompt: str) -> str:
    # Match de intent uit je bestaande test:
    # "image 1" / "Image 2" moet "Picture 1" / "Picture 2" worden.
    prompt = re.sub(r"\bimage\s+1\b", "Picture 1", prompt, flags=re.IGNORECASE)
    prompt = re.sub(r"\bimage\s+2\b", "Picture 2", prompt, flags=re.IGNORECASE)
    return prompt


def run_case(edit_api, HTTPException, case: dict, output_dir: Path, check_refinement: bool) -> tuple[bool, str]:
    folder = case["folder"]
    case_name = slug(folder.relative_to(folder.parents[0]).as_posix())

    model_filename = f"{case_name}__image1{case['image1'].suffix.lower()}"
    scene_filename = f"{case_name}__image2{case['image2'].suffix.lower()}"

    model_path = output_dir / model_filename
    scene_path = output_dir / scene_filename

    shutil.copy2(case["image1"], model_path)
    shutil.copy2(case["image2"], scene_path)

    prompt = case["prompt"].read_text(encoding="utf-8").strip()

    saved_prompts = []

    def mock_save_db_prompt(ptype, prompt_text, meta):
        saved_prompts.append(prompt_text)

    # Test 1: missing model image -> 404
    req_missing = edit_api.SceneryRequest(
        model_filename=f"missing_{case_name}.png",
        scene_image=scene_filename,
        prompt=prompt,
    )

    try:
        edit_api.generate_scenery(req_missing)
        return False, "Missing-model check faalde: er kwam geen HTTPException."
    except HTTPException as exc:
        if exc.status_code != 404:
            return False, f"Missing-model check gaf status {exc.status_code}, verwacht 404."
        if "Model image not found" not in str(exc.detail):
            return False, f"Missing-model check detail onverwacht: {exc.detail}"

    # Test 2: model en scene zijn hetzelfde bestand -> 400
    req_identical = edit_api.SceneryRequest(
        model_filename=model_filename,
        scene_image=model_filename,
        prompt=prompt,
    )

    try:
        edit_api.generate_scenery(req_identical)
        return False, "Identical-image check faalde: er kwam geen HTTPException."
    except HTTPException as exc:
        if exc.status_code != 400:
            return False, f"Identical-image check gaf status {exc.status_code}, verwacht 400."
        if "cannot be the same image" not in str(exc.detail):
            return False, f"Identical-image check detail onverwacht: {exc.detail}"

    # Test 3: geldige case, prompt refinement controleren, geen echte thread starten
    req_valid = edit_api.SceneryRequest(
        model_filename=model_filename,
        scene_image=scene_filename,
        prompt=prompt,
    )

    with patch("edit_api._load_wireframe_dir", return_value=str(output_dir), create=True), \
            patch("edit_api.database.save_db_prompt", side_effect=mock_save_db_prompt), \
            patch("threading.Thread.start"):

        edit_api.generate_scenery(req_valid)

    if len(saved_prompts) != 1:
        return False, f"save_db_prompt werd {len(saved_prompts)} keer aangeroepen, verwacht 1."

    if check_refinement:
        expected = expected_prompt_refinement(prompt)
        actual = saved_prompts[0]

        if actual != expected:
            return False, (
                "Prompt refinement klopt niet.\n"
                f"Expected: {expected!r}\n"
                f"Actual:   {actual!r}"
            )

    return True, "OK"


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--cases-dir",
        default="scenery_cases",
        help="Root-folder met subfolders die image1.*, image2.* en prompt.md bevatten.",
    )
    parser.add_argument(
        "--output-dir",
        default=None,
        help="Optioneel vaste output-dir. Default gebruikt een tijdelijke map.",
    )
    parser.add_argument(
        "--no-refinement-check",
        action="store_true",
        help="Alleen controleren dat generate_scenery werkt, niet de exacte prompt-output.",
    )

    args = parser.parse_args()

    project_root = Path.cwd()
    sys.path.insert(0, str(project_root))

    import edit_api
    from fastapi import HTTPException

    cases_root = Path(args.cases_dir).resolve()

    if not cases_root.exists():
        print(f"ERROR: cases-dir bestaat niet: {cases_root}")
        return 2

    cases = find_cases(cases_root)

    if not cases:
        print(f"Geen testcases gevonden onder: {cases_root}")
        print("Verwachte structuur per testcase:")
        print("  some_case/")
        print("    image1.png")
        print("    image2.png")
        print("    prompt.md")
        return 2

    if args.output_dir:
        output_dir = Path(args.output_dir).resolve()
        output_dir.mkdir(parents=True, exist_ok=True)
        cleanup = contextlib.nullcontext(output_dir)
    else:
        cleanup = tempfile.TemporaryDirectory(prefix="scenery_validation_")

    passed = 0
    failed = 0

    with cleanup as tmp:
        output_dir = Path(tmp) if not isinstance(tmp, Path) else tmp

        print(f"Cases root : {cases_root}")
        print(f"Output dir : {output_dir}")
        print(f"Cases      : {len(cases)}")
        print()

        for case in cases:
            label = case["folder"].relative_to(cases_root)

            try:
                ok, message = run_case(
                    edit_api=edit_api,
                    HTTPException=HTTPException,
                    case=case,
                    output_dir=output_dir,
                    check_refinement=not args.no_refinement_check,
                )

                if ok:
                    passed += 1
                    print(f"[PASS] {label}")
                else:
                    failed += 1
                    print(f"[FAIL] {label}")
                    print(message)
                    print()

            except Exception:
                failed += 1
                print(f"[ERROR] {label}")
                traceback.print_exc()
                print()

    print()
    print(f"Resultaat: {passed} passed, {failed} failed")

    return 0 if failed == 0 else 1


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