test_nodes.py

"""Unit tests for the bound functions — no network, runs in a second.

    python apps/05_workflow1111/test_nodes.py
"""

import base64
import io
import os
import sys
import tempfile
import traceback

from PIL import Image

# Windows consoles default to cp1252 and would die on the box-drawing output.
for _stream in (sys.stdout, sys.stderr):
    try:
        _stream.reconfigure(encoding="utf-8", errors="replace")
    except (AttributeError, ValueError):
        pass

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import nodes as N  # noqa: E402

PASS = FAIL = 0
FAILURES = []


def check(name, fn):
    global PASS, FAIL
    try:
        fn()
    except Exception as e:
        FAIL += 1
        FAILURES.append((name, traceback.format_exc()))
        print(f"  FAIL  {name}: {type(e).__name__}: {e}")
    else:
        PASS += 1
        print(f"  ok    {name}")


def sample(w=256, h=192, mode="RGB"):
    img = Image.new(mode, (w, h), (40, 90, 160) if mode == "RGB" else 128)
    d = __import__("PIL.ImageDraw", fromlist=["ImageDraw"]).Draw(img)
    d.ellipse([w // 5, h // 5, w * 4 // 5, h * 4 // 5], fill=(230, 200, 60))
    d.rectangle([10, 10, w // 3, h // 3], fill=(20, 20, 20))
    return img


def as_data_uri(img, fmt="PNG"):
    buf = io.BytesIO()
    img.save(buf, format=fmt)
    mime = "image/png" if fmt == "PNG" else "image/jpeg"
    return f"data:{mime};base64," + base64.b64encode(buf.getvalue()).decode()


def is_img(v):
    """An image port value must carry BOTH a readable file (for the REST API /
    gradio Image component) and a data URI (for the canvas and model nodes)."""
    assert isinstance(v, dict), f"image output must be a dict, got {type(v).__name__}: {str(v)[:70]}"
    assert set(v) >= {"path", "url"}, f"missing keys: {sorted(v)}"
    assert os.path.isfile(v["path"]), f"path does not exist: {v['path']}"
    Image.open(v["path"]).load()
    assert v["url"].startswith("data:image/"), f"url is not a data URI: {v['url'][:60]}"
    Image.open(io.BytesIO(base64.b64decode(v["url"].partition(",")[2]))).load()
    return v


TMP = tempfile.gettempdir()
img_path = os.path.join(TMP, "wf1111_test_src.png")
sample().save(img_path)
DATA_URI = as_data_uri(sample())

print("\n── image loading ─────────────────────────────────────────────")
check("load from data URI", lambda: N._load_image(DATA_URI).size)
check("load from plain path", lambda: N._load_image(img_path).size)
check("load from {'path'} dict", lambda: N._load_image({"path": img_path}).size)
check("load from model-node dict (url wins)",
      lambda: N._load_image({"path": img_path, "url": f"/gradio_api/file={img_path}",
                             "is_file": True}).size)
check("load from /gradio_api/file= string",
      lambda: N._load_image(f"/gradio_api/file={img_path}").size)
check("load from ImageSlider-style list",
      lambda: N._load_image([{"path": img_path}, {"path": img_path}]).size)
check("load from PIL Image", lambda: N._load_image(sample()).size)


def _rejects_empty():
    for bad in (None, "", {}):
        try:
            N._load_image(bad)
        except ValueError:
            continue
        raise AssertionError(f"should have rejected {bad!r}")


check("rejects empty input", _rejects_empty)

print("\n── coercion ──────────────────────────────────────────────────")
check("_num clamps + coerces strings",
      lambda: (N._num("7.5", 1, 1, 5) == 5.0 and N._num("junk", 3) == 3.0
               and N._num(None, 2, integer=True) == 2 and N._num(True, 0) == 1.0)
      or (_ for _ in ()).throw(AssertionError("bad _num")))
check("_flag parses truthy words",
      lambda: (N._flag("yes") and N._flag(True) and not N._flag("off")
               and N._flag("", True))
      or (_ for _ in ()).throw(AssertionError("bad _flag")))
check("_choice is lenient",
      lambda: (N._choice("photorealistic", list(N.STYLE_PRESETS), "None") == "Photorealistic"
               and N._choice("", list(N.STYLE_PRESETS), "None") == "None"
               and N._choice("zzz", list(N.STYLE_PRESETS), "None") == "None")
      or (_ for _ in ()).throw(AssertionError("bad _choice")))
check("_as_list handles json strings",
      lambda: (N._as_list('[{"a":1}]') == [{"a": 1}] and N._as_list(None) == []
               and N._as_list("garbage") == [])
      or (_ for _ in ()).throw(AssertionError("bad _as_list")))

print("\n── prompt / sampler ──────────────────────────────────────────")


def t_style():
    out = N.apply_style("a fox", "Cinematic", "golden hour", True)
    assert "a fox" in out and "cinematic film still" in out
    assert out.count("masterpiece") == 1
    # dedupe across sources
    dup = N.apply_style("sharp focus", "Photorealistic", "sharp focus", True)
    assert dup.lower().count("sharp focus") == 1, dup


check("apply_style composes + dedupes", t_style)


def t_style_blank():
    try:
        N.apply_style("   ", "None", "", False)
    except ValueError:
        return
    raise AssertionError("blank prompt should raise")


check("apply_style raises on blank", t_style_blank)


def t_neg():
    out = N.build_negative("ugly", "Anime", True, True)
    assert "ugly" in out and "lowres" in out and "photorealistic" in out and "nsfw" in out
    bare = N.build_negative("ugly", "None", False, False)
    assert bare == "ugly", bare


check("build_negative layers correctly", t_neg)


def t_sampler():
    steps, cfg, seed, w, h = N.sampler_settings(4, 1.0, 42, "Custom", 1000, 1000)
    assert (steps, cfg, seed) == (4, 1.0, 42)
    assert w % 16 == 0 and h % 16 == 0, (w, h)

    # cfg is clamped to >= 1.0 — fal-ai 422s below that
    assert N.sampler_settings(4, 0.0, 1, "Custom", 512, 512)[1] == 1.0
    assert N.sampler_settings(999, 99, 1, "Custom", 9999, 9999)[0] == 50

    # aspect preset overrides w/h
    assert N.sampler_settings(4, 1, 1, "16:9 Widescreen", 512, 512)[3:] == (1344, 768)

    # -1 randomizes into range
    s1 = N.sampler_settings(4, 1, -1, "Custom", 512, 512)[2]
    assert 0 <= s1 <= 2**31 - 1
    # blank/garbage seed falls back to the default (-1 → random), never crashes
    assert 0 <= N.sampler_settings(4, 1, "", "Custom", 512, 512)[2] <= 2**31 - 1


check("sampler_settings clamps/validates", t_sampler)


def t_info():
    txt = N.generation_info("a fox", "ugly", 4, 1.0, 42, 1024, 768, "m/x")
    assert txt.startswith("a fox")
    assert "Negative prompt: ugly" in txt
    assert "Seed: 42" in txt and "Size: 1024x768" in txt


check("generation_info matches A1111 format", t_info)


def t_matrix():
    p1, p2, p3, p4, labels = N.prompt_matrix("a cat", "in snow|on mars|at night|in a forest", "8k")
    assert p1.startswith("a cat") and "in snow" in p1 and "8k" in p1
    assert "on mars" in p2 and "at night" in p3 and "in a forest" in p4
    assert labels.count("|") == 3
    # fewer than four variants still fills four slots
    r = N.prompt_matrix("a cat", "red", "")
    assert all(isinstance(x, str) and x for x in r[:4])


check("prompt_matrix expands to 4", t_matrix)


def t_clean():
    messy = ('Sure! Here is your prompt:\n\n"a fox, golden hour, cinematic, bokeh"\n\n'
             "Let me know if you want changes.")
    out = N.clean_prompt(messy, 40)
    assert out.startswith("a fox"), out
    assert "Sure" not in out and '"' not in out and "Let me know" not in out
    assert len(N.clean_prompt("a, b, c, d, e, f", 3).split(",")) == 3
    assert N.clean_prompt("<think>hmm</think>\na fox, bokeh", 40).startswith("a fox")


check("clean_prompt strips LLM chatter", t_clean)


def t_labels():
    text, rows_json = N.top_labels(
        [{"label": "tiger", "score": 0.88}, {"label": "cat", "score": 0.10},
         {"label": "dog", "score": 0.001}], 2, 0.05)
    # second output is JSON *text* — a json port arrives as "[object Object]"
    rows = __import__("json").loads(rows_json)
    assert rows[0]["label"] == "tiger" and len(rows) == 2
    # and it must survive a round trip through a text port
    assert N.top_labels(rows_json, 2, 0.05)[0].startswith("tiger")
    assert "tiger" in text and "%" in text
    empty, _ = N.top_labels([], 5, 0.5)
    assert "No labels" in empty


check("top_labels formats + filters", t_labels)

print("\n── image operators ───────────────────────────────────────────")
check("postprocess full chain",
      lambda: is_img(N.postprocess(DATA_URI, 1.5, "Lanczos", 0.6, 1.2, 1.1, 1.0,
                                   0.35, 0.15, 0.02, "© test", "")))
check("postprocess neutral settings",
      lambda: is_img(N.postprocess(DATA_URI, 1, "Lanczos", 0, 1, 1, 1, 0, 0, 0, "", "")))
check("prep_image: Fit", lambda: is_img(N.prep_image(DATA_URI, 512, "Fit", True)))
check("prep_image: Pad to square",
      lambda: is_img(N.prep_image(DATA_URI, 384, "Pad to square", True)))
check("prep_image: Cover", lambda: is_img(N.prep_image(DATA_URI, 384, "Cover (crop)", True)))
check("prep_image: Stretch", lambda: is_img(N.prep_image(DATA_URI, 384, "Stretch", False)))


def t_prep_square():
    out = N.prep_image(DATA_URI, 320, "Pad to square", True)
    img = N._load_image(out)
    assert img.size == (320, 320), img.size


check("prep_image pads to exact square", t_prep_square)


def t_upscale():
    out, report = N.extras_upscale(DATA_URI, 2.0, "Lanczos", 0.5, True, True)
    is_img(out)
    img = N._load_image(out)
    assert img.size == (512, 384), img.size
    assert "512×384" in report and "Lanczos" in report


check("extras_upscale doubles + reports", t_upscale)

for mode in N.CONTROL_MODES:
    check(f"controlnet: {mode}",
          lambda m=mode: is_img(N.controlnet_preprocess(DATA_URI, m, 60, 160, False, 0)))
check("controlnet: inverted canny",
      lambda: is_img(N.controlnet_preprocess(DATA_URI, "Canny edges", 40, 120, True, 1.0)))


def t_canny_real():
    """Canny must actually find the disc/rect edges — and nothing in flat areas."""
    out = N.controlnet_preprocess(DATA_URI, "Canny edges", 50, 140, False, 0)
    arr = __import__("numpy").asarray(N._load_image(out).convert("L"))
    white = (arr > 200).mean()
    assert 0.001 < white < 0.30, f"implausible edge density {white:.4f}"
    flat = N.controlnet_preprocess(as_data_uri(Image.new("RGB", (128, 128), (90, 90, 90))),
                                   "Canny edges", 50, 140, False, 0)
    flat_arr = __import__("numpy").asarray(N._load_image(flat).convert("L"))
    assert (flat_arr > 200).mean() < 0.01, "flat image should yield no edges"


check("canny finds real edges, not noise", t_canny_real)

DETS = [
    {"box": {"xmin": 20, "ymin": 20, "xmax": 120, "ymax": 120}, "label": "cat", "score": 0.93},
    {"box": {"xmin": 60, "ymin": 40, "xmax": 200, "ymax": 150}, "label": "dog", "score": 0.61},
    {"box": {"xmin": 0, "ymin": 0, "xmax": 30, "ymax": 30}, "label": "bird", "score": 0.11},
]

check("draw_detections", lambda: is_img(N.draw_detections(DATA_URI, DETS, 0.5, True)[0]))
check("draw_detections from JSON string",
      lambda: is_img(N.draw_detections(DATA_URI, __import__("json").dumps(DETS), 0.5, True)[0]))


def t_det_summary():
    _, summary = N.draw_detections(DATA_URI, DETS, 0.5, True)
    assert "2 object(s)" in summary, summary
    _, none = N.draw_detections(DATA_URI, DETS, 0.99, True)
    assert "No objects" in none


check("draw_detections summary counts", t_det_summary)
check("mask_from_detections", lambda: is_img(N.mask_from_detections(
    DATA_URI, DETS, "", 0.5, 6, False, False)))
check("mask filtered by label", lambda: is_img(N.mask_from_detections(
    DATA_URI, DETS, "cat", 0.5, 4, False, False)))
check("mask preview overlay", lambda: is_img(N.mask_from_detections(
    DATA_URI, DETS, "", 0.5, 4, False, True)))


def t_mask_white():
    out = N.mask_from_detections(DATA_URI, DETS, "cat", 0.5, 0, False, False)
    arr = __import__("numpy").asarray(N._load_image(out).convert("L"))
    assert arr[70, 70] > 200, "inside the cat box should be white"
    assert arr[180, 240] < 60, "outside every box should be black"


check("mask geometry is correct", t_mask_white)


def t_mask_nohit():
    try:
        N.mask_from_detections(DATA_URI, DETS, "elephant", 0.5, 4, False, False)
    except ValueError:
        return
    raise AssertionError("unmatched label filter should raise")


check("mask raises when nothing matches", t_mask_nohit)

check("contact_sheet 4-up", lambda: is_img(N.contact_sheet(
    DATA_URI, DATA_URI, DATA_URI, DATA_URI, "a|b|c|d", 2, 14, "Prompt matrix")))
check("contact_sheet tolerates gaps", lambda: is_img(N.contact_sheet(
    DATA_URI, None, "", DATA_URI, "a|d", 2, 10, "")))
check("contact_sheet single column", lambda: is_img(N.contact_sheet(
    DATA_URI, None, None, None, "solo", 1, 0, "One")))


def t_sheet_empty():
    try:
        N.contact_sheet(None, None, "", None, "", 2, 10, "")
    except ValueError:
        return
    raise AssertionError("empty contact sheet should raise")


check("contact_sheet raises when empty", t_sheet_empty)

print("\n── PNG info round trip ───────────────────────────────────────")


def t_png_roundtrip():
    info = N.generation_info("a fox in snow", "ugly, blurry", 8, 3.5, 12345,
                             1024, 768, "black-forest-labs/FLUX.1-schnell")
    stamped = N.postprocess(DATA_URI, 1, "Lanczos", 0, 1, 1, 1, 0, 0, 0, "", info)
    assert stamped["url"].startswith("data:image/png"), "metadata must force PNG"

    report, fields_json = N.png_info(stamped)
    fields = __import__("json").loads(fields_json)
    assert "Generation parameters" in report
    assert fields["prompt"] == "a fox in snow", fields.get("prompt")
    assert fields["negative_prompt"] == "ugly, blurry", fields.get("negative_prompt")
    assert str(fields.get("seed")) == "12345", fields.get("seed")
    assert fields.get("size") == "1024x768", fields.get("size")
    assert str(fields.get("cfg_scale")) == "3.5", fields.get("cfg_scale")
    assert str(fields.get("steps")) == "8", fields.get("steps")


check("generation params survive the round trip", t_png_roundtrip)


def t_png_bare():
    report, fields_json = N.png_info(DATA_URI)
    fields = __import__("json").loads(fields_json)
    assert "No generation parameters" in report
    assert fields["width"] == 256 and fields["height"] == 192


check("png_info on a bare image", t_png_bare)

print("\n── output sizing ─────────────────────────────────────────────")


def t_jpeg_switch():
    big = Image.new("RGB", (2000, 1400), (100, 120, 140))
    assert N._emit(big)["url"].startswith("data:image/jpeg"), "large images should be JPEG"
    small = Image.new("RGB", (400, 400), (100, 120, 140))
    assert N._emit(small)["url"].startswith("data:image/png")
    alpha = Image.new("RGBA", (2000, 1400), (100, 120, 140, 128))
    assert N._emit(alpha)["url"].startswith("data:image/png"), "alpha must stay PNG"
    assert N._emit_uri(small).startswith("data:image/png")


check("emit picks PNG/JPEG sensibly", t_jpeg_switch)


def t_bind_complete():
    import inspect
    for name, fn in N.BIND.items():
        assert callable(fn), name
        assert getattr(fn, "__name__", None) == name, f"{name} bound to {fn}"
        assert not str(inspect.signature(fn)).startswith("(self"), name


check("BIND keys match function names", t_bind_complete)

print("\n" + "=" * 62)
print(f"  {PASS} passed, {FAIL} failed")
print("=" * 62)
if FAILURES:
    for name, tb in FAILURES:
        print(f"\n--- {name} ---\n{tb}")
sys.exit(1 if FAIL else 0)