by-ae
fd468f5c2b
The prompt joined objects in a list with a comma and a space, so every element after the first started on the same line as the previous brace.
673 lines
28 KiB
Python
673 lines
28 KiB
Python
"""Unit tests for nodes/ideogram4_nodes.py prompt assembly and validation."""
|
|
|
|
import json
|
|
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
# conftest.py handles mocking and path setup
|
|
from nodes.ideogram4_nodes import (
|
|
sanitize_bg_guid, _norm_bbox, _palette, _hex_rgb, _dumps, _parse_json_list,
|
|
_caption_to_boxes, _coerce_layout_int, _coerce_opacity,
|
|
_elements_from_boxes, _preview_edges, _render_preview, ImageDraw,
|
|
save_bg_image_to_temp, _uint8_frame_kind, _rgb_image_from_uint8, Image,
|
|
_parse_import_caption, _caption_field, _resolve_style, _editor_caption,
|
|
Ideogram4PromptBuilderKJEnhancedByAE, io,
|
|
)
|
|
|
|
|
|
class _FakeFrame:
|
|
def __init__(self, shape):
|
|
self.shape = shape
|
|
self.reshaped = None
|
|
|
|
def reshape(self, *args):
|
|
self.reshaped = args
|
|
return ("flat", args)
|
|
|
|
|
|
class TestPromptHelpers(unittest.TestCase):
|
|
def test_sanitize_bg_guid_valid(self):
|
|
valid_guid = "12345678-1234-4abc-8abc-123456789012"
|
|
self.assertEqual(sanitize_bg_guid(valid_guid), valid_guid)
|
|
|
|
# Test uppercase
|
|
upper_guid = "12345678-1234-4ABC-8ABC-123456789012"
|
|
self.assertEqual(sanitize_bg_guid(upper_guid), upper_guid)
|
|
|
|
def test_sanitize_bg_guid_invalid(self):
|
|
with self.assertRaises(ValueError):
|
|
sanitize_bg_guid("not-a-guid")
|
|
|
|
with self.assertRaises(ValueError):
|
|
sanitize_bg_guid("")
|
|
|
|
with self.assertRaises(ValueError):
|
|
sanitize_bg_guid("12345678-1234-4abc-8abc") # Too short
|
|
|
|
# Path traversal attempt
|
|
with self.assertRaises(ValueError):
|
|
sanitize_bg_guid("../../etc/passwd")
|
|
|
|
def test_norm_bbox_basic(self):
|
|
box = {"x": 0.1, "y": 0.2, "w": 0.3, "h": 0.4}
|
|
result = _norm_bbox(box)
|
|
self.assertEqual(result, [200, 100, 600, 400])
|
|
|
|
def test_norm_bbox_swap_coords_if_inverted(self):
|
|
# Test ymin > ymax case
|
|
box = {"x": 0.1, "y": 0.6, "w": 0.3, "h": -0.4}
|
|
result = _norm_bbox(box)
|
|
# Should swap ymin and ymax
|
|
self.assertEqual(result[0], min(result[0], result[2]))
|
|
self.assertEqual(result[2], max(result[0], result[2]))
|
|
|
|
def test_norm_bbox_clamps_to_1000(self):
|
|
box = {"x": 0.0, "y": 0.0, "w": 2.0, "h": 2.0}
|
|
result = _norm_bbox(box)
|
|
self.assertTrue(all(0 <= v <= 1000 for v in result))
|
|
|
|
def test_norm_bbox_missing_fields(self):
|
|
# Should handle missing fields gracefully
|
|
box = {}
|
|
result = _norm_bbox(box)
|
|
self.assertEqual(result, [0, 0, 0, 0])
|
|
|
|
def test_norm_bbox_accepts_numeric_strings(self):
|
|
box = {"x": "0.1", "y": "0.2", "w": "0.3", "h": "0.4"}
|
|
self.assertEqual(_norm_bbox(box), [200, 100, 600, 400])
|
|
|
|
def test_norm_bbox_rejects_non_finite_coordinates(self):
|
|
self.assertIsNone(_norm_bbox({"x": float("nan"), "y": 0.2, "w": 0.3, "h": 0.4}))
|
|
self.assertIsNone(_norm_bbox({"x": "nope", "y": 0.2, "w": 0.3, "h": 0.4}))
|
|
self.assertIsNone(_norm_bbox({"x": True, "y": 0.0, "w": 0.2, "h": 0.2}))
|
|
self.assertIsNone(_norm_bbox({"x": float("inf"), "y": 0.0, "w": 0.2, "h": 0.2}))
|
|
|
|
def test_norm_bbox_huge_finite_values_do_not_raise(self):
|
|
result = _norm_bbox({"x": 1e308, "y": 0.0, "w": 1e308, "h": 0.2})
|
|
self.assertEqual(result, [0, 1000, 200, 1000])
|
|
|
|
def test_norm_bbox_half_steps_round_like_the_editor(self):
|
|
# 2.5 and 0.5. Copy JSON uses JS Math.round (half toward +inf).
|
|
# Python round() uses ties-to-even, which made the prompt one unit off.
|
|
box = {"x": 0.0025, "y": 0.0005, "w": 0.001, "h": 0.001}
|
|
self.assertEqual(_norm_bbox(box), [1, 3, 2, 4])
|
|
|
|
def test_elements_from_boxes_omits_unusable_bbox(self):
|
|
elements = _elements_from_boxes([
|
|
{"type": "obj", "desc": "bad", "x": "nope", "y": 0, "w": 0.2, "h": 0.2},
|
|
{
|
|
"type": "text", "text": "Hi", "desc": "ok",
|
|
"x": 0.1, "y": 0.2, "w": 0.3, "h": 0.4,
|
|
},
|
|
"skip",
|
|
{"type": "obj", "desc": "no box", "nobbox": True, "x": 0.1, "y": 0.1, "w": 0.2, "h": 0.2},
|
|
])
|
|
self.assertEqual(len(elements), 3)
|
|
self.assertNotIn("bbox", elements[0])
|
|
self.assertEqual(elements[0]["desc"], "bad")
|
|
self.assertEqual(elements[1]["bbox"], [200, 100, 600, 400])
|
|
self.assertEqual(elements[1]["text"], "Hi")
|
|
self.assertEqual(list(elements[1].keys()), ["type", "bbox", "text", "desc"])
|
|
self.assertNotIn("bbox", elements[2])
|
|
|
|
def test_preview_edges_skip_unusable_box(self):
|
|
self.assertIsNone(_preview_edges({"x": "nope", "y": 0, "w": 0.2, "h": 0.2}, 100, 50))
|
|
self.assertEqual(
|
|
_preview_edges({"x": 0.1, "y": 0.2, "w": 0.3, "h": 0.4}, 100, 50),
|
|
(10, 10, 40, 30),
|
|
)
|
|
|
|
def test_preview_edges_half_pixel_rounds_like_the_editor(self):
|
|
self.assertEqual(
|
|
_preview_edges({"x": 0.005, "y": 0, "w": 0.01, "h": 0.02}, 100, 100),
|
|
(1, 0, 2, 2),
|
|
)
|
|
|
|
def test_elements_keep_schema_types_when_fields_are_not_strings(self):
|
|
elements = _elements_from_boxes([
|
|
{
|
|
"type": "text", "text": 12, "desc": None, "palette": True,
|
|
"x": 0.1, "y": 0.2, "w": 0.3, "h": 0.4,
|
|
},
|
|
{
|
|
"type": "obj", "desc": 5, "palette": ["#abc", "#GG0000", 1],
|
|
"x": 0.0, "y": 0.0, "w": 0.2, "h": 0.2,
|
|
},
|
|
{
|
|
"type": "obj", "desc": "a circle", "palette": ["#ff00aa"],
|
|
"x": 0.5, "y": 0.5, "w": 0.2, "h": 0.2,
|
|
},
|
|
])
|
|
self.assertEqual(elements[0]["text"], "12")
|
|
self.assertEqual(elements[0]["desc"], "")
|
|
self.assertNotIn("color_palette", elements[0])
|
|
self.assertEqual(list(elements[0].keys()), ["type", "bbox", "text", "desc"])
|
|
self.assertEqual(elements[1]["desc"], "5")
|
|
self.assertEqual(elements[1]["color_palette"], ["#AABBCC"])
|
|
self.assertEqual(list(elements[1].keys()), ["type", "bbox", "desc", "color_palette"])
|
|
self.assertEqual(elements[2]["desc"], "a circle")
|
|
self.assertEqual(elements[2]["color_palette"], ["#FF00AA"])
|
|
|
|
def test_preview_does_not_raise_on_bad_region_fields(self):
|
|
ImageDraw.Draw.return_value.textlength.return_value = 0
|
|
preview = _render_preview([
|
|
{
|
|
"type": "text", "text": 7, "desc": 5,
|
|
"palette": ["#abc", "#GG0000", 1, " #00ff00 "],
|
|
"x": 0.1, "y": 0.1, "w": 0.4, "h": 0.4,
|
|
},
|
|
], 128, 128)
|
|
self.assertIsNotNone(preview)
|
|
|
|
def test_palette_filters_empty_strings(self):
|
|
colors = {"c1": "#FF0000", "c2": "", "c3": "#00FF00", "c4": None}
|
|
result = _palette(colors)
|
|
self.assertEqual(result, ["#FF0000", "#00FF00"])
|
|
|
|
def test_palette_uppercase(self):
|
|
colors = ["#ff0000", "#00ff00"]
|
|
result = _palette(colors)
|
|
self.assertEqual(result, ["#FF0000", "#00FF00"])
|
|
|
|
def test_palette_drops_invalid_and_expands_short_hex(self):
|
|
colors = [None, "", "red", "#fff", "#aabbcc", 1, True, "#12"]
|
|
self.assertEqual(_palette(colors), ["#FFFFFF", "#AABBCC"])
|
|
self.assertEqual(_palette(None), [])
|
|
|
|
def test_palette_respects_ideogram_limits(self):
|
|
colors = ["#{:06X}".format(i) for i in range(20)]
|
|
self.assertEqual(len(_palette(colors, 16)), 16)
|
|
self.assertEqual(_palette(colors, 5), ["#{:06X}".format(i) for i in range(5)])
|
|
|
|
def test_palette_ignores_values_that_are_not_color_lists(self):
|
|
self.assertEqual(_palette(True), [])
|
|
self.assertEqual(_palette(5), [])
|
|
self.assertEqual(_palette("#FF0000"), [])
|
|
self.assertEqual(_palette(None), [])
|
|
|
|
def test_hex_rgb_expands_short_hex_and_rejects_invalid(self):
|
|
self.assertEqual(_hex_rgb("#abc"), (0xAA, 0xBB, 0xCC))
|
|
self.assertEqual(_hex_rgb(" #aabbcc "), (0xAA, 0xBB, 0xCC))
|
|
self.assertEqual(_hex_rgb("#GG0000"), (255, 255, 255))
|
|
self.assertEqual(_hex_rgb(1), (255, 255, 255))
|
|
|
|
def test_dumps_string(self):
|
|
result = _dumps("hello")
|
|
self.assertEqual(result, '"hello"')
|
|
|
|
def test_dumps_list_simple(self):
|
|
result = _dumps([1, 2, 3])
|
|
self.assertEqual(result, "[1, 2, 3]")
|
|
|
|
def test_dumps_list_nested(self):
|
|
result = _dumps([{"a": 1}, {"b": 2}])
|
|
self.assertIn("[\n", result)
|
|
self.assertIn('"a": 1', result)
|
|
# A comma-space join left the next object on the same line as "}".
|
|
self.assertIn("},\n", result)
|
|
self.assertNotIn("}, ", result)
|
|
self.assertEqual(json.loads(result), [{"a": 1}, {"b": 2}])
|
|
|
|
def test_dumps_dict(self):
|
|
result = _dumps({"key": "value", "num": 123})
|
|
self.assertIn('"key": "value"', result)
|
|
self.assertIn('"num": 123', result)
|
|
|
|
def test_dumps_unicode(self):
|
|
result = _dumps({"text": "Hello 世界"})
|
|
self.assertIn("世界", result)
|
|
|
|
def test_parse_json_list_valid(self):
|
|
result = _parse_json_list('[{"x": 1}, {"y": 2}]')
|
|
self.assertEqual(result, [{"x": 1}, {"y": 2}])
|
|
|
|
def test_parse_json_list_invalid(self):
|
|
with self.assertLogs("byae.ideogram.prompt", level="WARNING") as cm:
|
|
result = _parse_json_list("not json", "region data")
|
|
self.assertEqual(result, [])
|
|
self.assertTrue(any("region data" in line and "not valid JSON" in line for line in cm.output))
|
|
|
|
def test_parse_json_list_not_list(self):
|
|
with self.assertLogs("byae.ideogram.prompt", level="WARNING") as cm:
|
|
result = _parse_json_list('{"key": "value"}', "style palette")
|
|
self.assertEqual(result, [])
|
|
self.assertTrue(any("style palette" in line and "not a JSON list" in line for line in cm.output))
|
|
|
|
def test_parse_json_list_empty(self):
|
|
with self.assertNoLogs("byae.ideogram.prompt", level="WARNING"):
|
|
self.assertEqual(_parse_json_list(""), [])
|
|
self.assertEqual(_parse_json_list(" "), [])
|
|
self.assertEqual(_parse_json_list(None), [])
|
|
self.assertEqual(_parse_json_list("\ufeff[1, 2]"), [1, 2])
|
|
|
|
def test_bad_region_data_warns_once_and_keeps_the_caption(self):
|
|
io.NodeOutput.reset_mock()
|
|
with self.assertLogs("byae.ideogram.prompt", level="WARNING") as cm:
|
|
Ideogram4PromptBuilderKJEnhancedByAE.execute(
|
|
1024, 1024, "room", {"style": "none"},
|
|
elements_data="{nope",
|
|
style_palette_data="{}",
|
|
)
|
|
region = [line for line in cm.output if "region data" in line]
|
|
palette = [line for line in cm.output if "style palette" in line]
|
|
self.assertEqual(len(region), 1)
|
|
self.assertEqual(len(palette), 1)
|
|
prompt = io.NodeOutput.call_args.args[0]
|
|
self.assertIn('"background": "room"', prompt)
|
|
self.assertIn('"elements": []', prompt)
|
|
|
|
def test_caption_to_boxes_valid(self):
|
|
caption = {
|
|
"compositional_deconstruction": {
|
|
"elements": [
|
|
{
|
|
"type": "text",
|
|
"text": "Hello",
|
|
"desc": "A greeting",
|
|
"bbox": [100, 200, 300, 400],
|
|
"color_palette": ["#FF0000"]
|
|
}
|
|
]
|
|
}
|
|
}
|
|
result = _caption_to_boxes(caption)
|
|
self.assertEqual(len(result), 1)
|
|
self.assertEqual(result[0]["type"], "text")
|
|
self.assertEqual(result[0]["text"], "Hello")
|
|
self.assertAlmostEqual(result[0]["x"], 0.2)
|
|
self.assertAlmostEqual(result[0]["y"], 0.1)
|
|
|
|
def test_caption_to_boxes_missing_bbox(self):
|
|
caption = {
|
|
"compositional_deconstruction": {
|
|
"elements": [
|
|
{"type": "obj", "desc": "No bbox"}
|
|
]
|
|
}
|
|
}
|
|
result = _caption_to_boxes(caption)
|
|
self.assertEqual(len(result), 1)
|
|
self.assertTrue(result[0]["nobbox"])
|
|
|
|
def test_caption_to_boxes_inverted_bbox_is_normalized(self):
|
|
caption = {
|
|
"compositional_deconstruction": {
|
|
"elements": [
|
|
{"type": "obj", "desc": "flipped", "bbox": [400, 300, 100, 200]}
|
|
]
|
|
}
|
|
}
|
|
result = _caption_to_boxes(caption)
|
|
self.assertNotIn("nobbox", result[0])
|
|
self.assertAlmostEqual(result[0]["x"], 0.2)
|
|
self.assertAlmostEqual(result[0]["y"], 0.1)
|
|
self.assertAlmostEqual(result[0]["w"], 0.1)
|
|
self.assertAlmostEqual(result[0]["h"], 0.3)
|
|
|
|
def test_caption_to_boxes_out_of_range_bbox_is_clamped(self):
|
|
caption = {
|
|
"compositional_deconstruction": {
|
|
"elements": [
|
|
{"type": "obj", "desc": "huge", "bbox": [-100, -50, 1500, 2000]}
|
|
]
|
|
}
|
|
}
|
|
result = _caption_to_boxes(caption)
|
|
self.assertNotIn("nobbox", result[0])
|
|
self.assertAlmostEqual(result[0]["x"], 0.0)
|
|
self.assertAlmostEqual(result[0]["y"], 0.0)
|
|
self.assertAlmostEqual(result[0]["w"], 1.0)
|
|
self.assertAlmostEqual(result[0]["h"], 1.0)
|
|
|
|
def test_caption_to_boxes_numeric_strings(self):
|
|
caption = {
|
|
"compositional_deconstruction": {
|
|
"elements": [
|
|
{"type": "text", "text": "Hi", "desc": "d", "bbox": ["100", "200", "300", "400"]}
|
|
]
|
|
}
|
|
}
|
|
result = _caption_to_boxes(caption)
|
|
self.assertNotIn("nobbox", result[0])
|
|
self.assertAlmostEqual(result[0]["x"], 0.2)
|
|
self.assertAlmostEqual(result[0]["y"], 0.1)
|
|
self.assertAlmostEqual(result[0]["w"], 0.2)
|
|
self.assertAlmostEqual(result[0]["h"], 0.2)
|
|
|
|
def test_caption_to_boxes_non_numeric_bbox_keeps_grid_index(self):
|
|
caption = {
|
|
"compositional_deconstruction": {
|
|
"elements": [
|
|
{"type": "obj", "desc": "ok", "bbox": [0, 0, 100, 100]},
|
|
{"type": "obj", "desc": "bad", "bbox": ["nope", 0, 10, 10]},
|
|
{"type": "obj", "desc": "bool", "bbox": [True, 0, 10, 10]},
|
|
]
|
|
}
|
|
}
|
|
result = _caption_to_boxes(caption)
|
|
self.assertNotIn("nobbox", result[0])
|
|
self.assertTrue(result[1]["nobbox"])
|
|
self.assertTrue(result[2]["nobbox"])
|
|
self.assertAlmostEqual(result[1]["x"], 0.02 + 0.25)
|
|
self.assertAlmostEqual(result[1]["y"], 0.02)
|
|
self.assertAlmostEqual(result[2]["x"], 0.02 + 0.50)
|
|
self.assertAlmostEqual(result[2]["y"], 0.02)
|
|
|
|
def test_caption_to_boxes_invalid_bbox(self):
|
|
caption = {
|
|
"compositional_deconstruction": {
|
|
"elements": [
|
|
{"type": "obj", "desc": "Bad bbox", "bbox": "not a list"}
|
|
]
|
|
}
|
|
}
|
|
result = _caption_to_boxes(caption)
|
|
self.assertEqual(len(result), 1)
|
|
self.assertTrue(result[0]["nobbox"])
|
|
|
|
def test_caption_to_boxes_many_missing_bbox_no_overlap(self):
|
|
"""Test that many elements without bbox don't overlap (arrange in rows)."""
|
|
caption = {
|
|
"compositional_deconstruction": {
|
|
"elements": [{"type": "obj", "desc": f"element_{i}"} for i in range(12)]
|
|
}
|
|
}
|
|
boxes = _caption_to_boxes(caption)
|
|
self.assertEqual(len(boxes), 12)
|
|
|
|
# Check that boxes are arranged in rows of 4 and don't overlap
|
|
for i, box in enumerate(boxes):
|
|
self.assertTrue(box.get("nobbox"))
|
|
expected_row = i // 4
|
|
expected_col = i % 4
|
|
# Row 0: y=0.02, Row 1: y=0.19, Row 2: y=0.36
|
|
expected_y = 0.02 + expected_row * 0.17
|
|
# Col 0: x=0.02, Col 1: x=0.27, Col 2: x=0.52, Col 3: x=0.77
|
|
expected_x = 0.02 + expected_col * 0.25
|
|
self.assertAlmostEqual(box["y"], expected_y, places=5)
|
|
self.assertAlmostEqual(box["x"], expected_x, places=5)
|
|
|
|
# Verify no overlaps by checking each pair
|
|
for i in range(len(boxes)):
|
|
for j in range(i + 1, len(boxes)):
|
|
b1, b2 = boxes[i], boxes[j]
|
|
# Check if boxes overlap (with small margin for floating point)
|
|
overlap_x = (b1["x"] < b2["x"] + b2["w"] - 0.01 and
|
|
b1["x"] + b1["w"] > b2["x"] + 0.01)
|
|
overlap_y = (b1["y"] < b2["y"] + b2["h"] - 0.01 and
|
|
b1["y"] + b1["h"] > b2["y"] + 0.01)
|
|
if overlap_x and overlap_y:
|
|
self.fail(f"Box {i} overlaps with box {j}: "
|
|
f"b{i}=({b1['x']:.3f},{b1['y']:.3f},{b1['w']:.3f},{b1['h']:.3f}) "
|
|
f"b{j}=({b2['x']:.3f},{b2['y']:.3f},{b2['w']:.3f},{b2['h']:.3f})")
|
|
|
|
def test_caption_to_boxes_empty(self):
|
|
result = _caption_to_boxes({})
|
|
self.assertEqual(result, [])
|
|
|
|
def test_coerce_layout_int_valid(self):
|
|
self.assertEqual(_coerce_layout_int(512), 512)
|
|
self.assertEqual(_coerce_layout_int("1024"), 1024)
|
|
self.assertEqual(_coerce_layout_int(2048.7), 2048)
|
|
|
|
def test_coerce_layout_int_invalid(self):
|
|
self.assertEqual(_coerce_layout_int(0), 1024)
|
|
self.assertEqual(_coerce_layout_int(-100), 1024)
|
|
self.assertEqual(_coerce_layout_int("not a number"), 1024)
|
|
self.assertEqual(_coerce_layout_int(None), 1024)
|
|
|
|
def test_coerce_opacity_valid(self):
|
|
self.assertEqual(_coerce_opacity(0.5), 0.5)
|
|
self.assertEqual(_coerce_opacity(0), 0.0)
|
|
self.assertEqual(_coerce_opacity(1), 1.0)
|
|
self.assertEqual(_coerce_opacity("0.75"), 0.75)
|
|
|
|
def test_coerce_opacity_clamping(self):
|
|
self.assertEqual(_coerce_opacity(-1), 0.0)
|
|
self.assertEqual(_coerce_opacity(2), 1.0)
|
|
self.assertEqual(_coerce_opacity(1.5), 1.0)
|
|
|
|
def test_coerce_opacity_invalid(self):
|
|
self.assertEqual(_coerce_opacity("not a number"), 0.5)
|
|
self.assertEqual(_coerce_opacity(None), 0.5)
|
|
self.assertEqual(_coerce_opacity(float('nan')), 0.5)
|
|
|
|
def test_save_bg_image_warns_when_image_missing(self):
|
|
with self.assertLogs("byae.ideogram.bg", level="WARNING") as cm:
|
|
result = save_bg_image_to_temp(
|
|
"12345678-1234-4abc-8abc-123456789012", None
|
|
)
|
|
self.assertIsNone(result)
|
|
self.assertTrue(any("no image" in line.lower() for line in cm.output))
|
|
|
|
def test_save_bg_image_warns_when_guid_empty(self):
|
|
with self.assertLogs("byae.ideogram.bg", level="WARNING") as cm:
|
|
result = save_bg_image_to_temp(" ", object())
|
|
self.assertIsNone(result)
|
|
self.assertTrue(any("guid is empty" in line.lower() for line in cm.output))
|
|
|
|
def test_save_bg_image_warns_when_guid_invalid(self):
|
|
with self.assertLogs("byae.ideogram.bg", level="WARNING") as cm:
|
|
result = save_bg_image_to_temp("../../etc/passwd", object())
|
|
self.assertIsNone(result)
|
|
self.assertTrue(any("invalid background guid" in line.lower() for line in cm.output))
|
|
|
|
def test_uint8_frame_kind_matches_channel_count(self):
|
|
self.assertEqual(_uint8_frame_kind((8, 8, 3)), "rgb")
|
|
self.assertEqual(_uint8_frame_kind((8, 8, 4)), "rgba")
|
|
self.assertEqual(_uint8_frame_kind((8, 8, 1)), "gray1")
|
|
self.assertEqual(_uint8_frame_kind((8, 8)), "gray")
|
|
# Anything else keeps the RGB attempt, which reports a conversion failure.
|
|
self.assertEqual(_uint8_frame_kind((8, 8, 2)), "rgb")
|
|
self.assertEqual(_uint8_frame_kind((0, 8, 3)), "rgb")
|
|
self.assertEqual(_uint8_frame_kind(None), "rgb")
|
|
# bool is an int, and True == 1, so a boolean channel count must not
|
|
# take the single-channel path.
|
|
self.assertEqual(_uint8_frame_kind((8, 8, True)), "rgb")
|
|
|
|
def test_rgba_frame_is_composited_onto_black(self):
|
|
arr = _FakeFrame((2, 3, 4))
|
|
Image.fromarray.reset_mock()
|
|
Image.new.reset_mock()
|
|
Image.alpha_composite.reset_mock()
|
|
_rgb_image_from_uint8(arr)
|
|
Image.fromarray.assert_called_once_with(arr, mode="RGBA")
|
|
Image.new.assert_called_once()
|
|
args, _kwargs = Image.new.call_args
|
|
self.assertEqual(args[0], "RGBA")
|
|
self.assertEqual(args[2], (0, 0, 0, 255))
|
|
Image.alpha_composite.assert_called_once()
|
|
Image.alpha_composite.return_value.convert.assert_called_once_with("RGB")
|
|
|
|
def test_single_channel_frame_is_opened_as_grayscale(self):
|
|
gray = _FakeFrame((4, 5))
|
|
Image.fromarray.reset_mock()
|
|
_rgb_image_from_uint8(gray)
|
|
Image.fromarray.assert_called_once_with(gray, mode="L")
|
|
Image.fromarray.return_value.convert.assert_called_once_with("RGB")
|
|
|
|
column = _FakeFrame((4, 5, 1))
|
|
Image.fromarray.reset_mock()
|
|
_rgb_image_from_uint8(column)
|
|
self.assertEqual(column.reshaped, (4, 5))
|
|
Image.fromarray.assert_called_once_with(("flat", (4, 5)), mode="L")
|
|
Image.fromarray.return_value.convert.assert_called_once_with("RGB")
|
|
|
|
def test_save_bg_image_warns_when_conversion_fails(self):
|
|
guid = "12345678-1234-4abc-8abc-123456789012"
|
|
image_mod = save_bg_image_to_temp.__globals__["Image"]
|
|
with patch.object(image_mod, "fromarray", side_effect=ValueError("bad shape")):
|
|
with self.assertLogs("byae.ideogram.bg", level="WARNING") as cm:
|
|
result = save_bg_image_to_temp(guid, MagicMock())
|
|
self.assertIsNone(result)
|
|
self.assertTrue(any("could not convert" in line.lower() for line in cm.output))
|
|
self.assertTrue(any("bad shape" in line for line in cm.output))
|
|
|
|
|
|
class TestImportCaption(unittest.TestCase):
|
|
def test_blank_import_keeps_the_editor_quietly(self):
|
|
with self.assertNoLogs("byae.ideogram.prompt", level="WARNING"):
|
|
self.assertIsNone(_parse_import_caption(""))
|
|
self.assertIsNone(_parse_import_caption(" "))
|
|
self.assertIsNone(_parse_import_caption(None))
|
|
|
|
def test_invalid_json_and_non_objects_are_ignored(self):
|
|
with self.assertLogs("byae.ideogram.prompt", level="WARNING") as cm:
|
|
self.assertIsNone(_parse_import_caption("{not json"))
|
|
self.assertTrue(any("not valid JSON" in line for line in cm.output))
|
|
self.assertTrue(any("editor caption" in line for line in cm.output))
|
|
|
|
with self.assertLogs("byae.ideogram.prompt", level="WARNING") as cm:
|
|
self.assertIsNone(_parse_import_caption("[1, 2]"))
|
|
self.assertTrue(any("not a caption object" in line for line in cm.output))
|
|
|
|
def test_bom_object_is_a_caption(self):
|
|
raw = '\ufeff{"compositional_deconstruction": {"background": "sky", "elements": []}}'
|
|
with self.assertNoLogs("byae.ideogram.prompt", level="WARNING"):
|
|
cap = _parse_import_caption(raw)
|
|
self.assertEqual(cap["compositional_deconstruction"]["background"], "sky")
|
|
|
|
def test_incomplete_caption_object_is_kept_and_reported(self):
|
|
with self.assertLogs("byae.ideogram.prompt", level="WARNING") as cm:
|
|
cap = _parse_import_caption("{}")
|
|
self.assertEqual(cap, {})
|
|
self.assertTrue(any("compositional_deconstruction" in line for line in cm.output))
|
|
self.assertTrue(any("still used" in line for line in cm.output))
|
|
|
|
with self.assertLogs("byae.ideogram.prompt", level="WARNING") as cm:
|
|
_parse_import_caption('{"compositional_deconstruction": {}}')
|
|
self.assertTrue(any("background must be a string" in line for line in cm.output))
|
|
|
|
with self.assertLogs("byae.ideogram.prompt", level="WARNING") as cm:
|
|
_parse_import_caption(
|
|
'{"compositional_deconstruction": {"background": "sky", "elements": {}}}'
|
|
)
|
|
self.assertTrue(any("elements must be a list" in line for line in cm.output))
|
|
|
|
|
|
class TestStyleChoice(unittest.TestCase):
|
|
def test_photo_dict_keeps_its_text(self):
|
|
with self.assertNoLogs("byae.ideogram.prompt", level="WARNING"):
|
|
kind, photo, art = _resolve_style({
|
|
"style": "photo",
|
|
"photo": "35mm",
|
|
"art_style": "ignored when photo is selected",
|
|
})
|
|
self.assertEqual(kind, "photo")
|
|
self.assertEqual(photo, "35mm")
|
|
self.assertEqual(art, "ignored when photo is selected")
|
|
|
|
def test_missing_choice_is_none(self):
|
|
with self.assertNoLogs("byae.ideogram.prompt", level="WARNING"):
|
|
self.assertEqual(_resolve_style({})[0], "none")
|
|
self.assertEqual(_resolve_style({"style": "none"})[0], "none")
|
|
self.assertEqual(_resolve_style("none"), ("none", "", ""))
|
|
|
|
def test_bare_style_string_does_not_fail(self):
|
|
with self.assertLogs("byae.ideogram.prompt", level="WARNING") as cm:
|
|
kind, photo, art = _resolve_style("art_style")
|
|
self.assertEqual((kind, photo, art), ("art_style", "", ""))
|
|
self.assertTrue(any("string" in line and "art_style" in line for line in cm.output))
|
|
|
|
def test_unusable_style_omits_the_block(self):
|
|
with self.assertLogs("byae.ideogram.prompt", level="WARNING") as cm:
|
|
self.assertEqual(_resolve_style(None), ("none", "", ""))
|
|
self.assertTrue(any("ignored" in line for line in cm.output))
|
|
with self.assertLogs("byae.ideogram.prompt", level="WARNING") as cm:
|
|
self.assertEqual(_resolve_style({"style": "watercolor"})[0], "none")
|
|
self.assertTrue(any("watercolor" in line for line in cm.output))
|
|
|
|
def test_caption_field_keeps_strings(self):
|
|
self.assertEqual(_caption_field(" hi "), " hi ")
|
|
self.assertEqual(_caption_field(""), "")
|
|
self.assertEqual(_caption_field(None), "")
|
|
self.assertEqual(_caption_field(12), "12")
|
|
self.assertEqual(_caption_field(True), "")
|
|
|
|
def test_editor_caption_survives_unusable_fields(self):
|
|
with self.assertLogs("byae.ideogram.prompt", level="WARNING"):
|
|
caption = _editor_caption(None, None, high_level_description=None)
|
|
self.assertEqual(caption["compositional_deconstruction"]["background"], "")
|
|
self.assertEqual(caption["compositional_deconstruction"]["elements"], [])
|
|
self.assertNotIn("style_description", caption)
|
|
self.assertNotIn("high_level_description", caption)
|
|
|
|
with self.assertNoLogs("byae.ideogram.prompt", level="WARNING"):
|
|
caption = _editor_caption(
|
|
"room",
|
|
{"style": "photo", "photo": "35mm"},
|
|
high_level_description=" overview ",
|
|
aesthetics="moody",
|
|
lighting="dim",
|
|
medium="photograph",
|
|
style_palette_data='["#fff"]',
|
|
)
|
|
self.assertEqual(caption["high_level_description"], " overview ")
|
|
style = caption["style_description"]
|
|
self.assertEqual(
|
|
list(style.keys()),
|
|
["aesthetics", "lighting", "photo", "medium", "color_palette"],
|
|
)
|
|
self.assertEqual(style["photo"], "35mm")
|
|
self.assertEqual(style["color_palette"], ["#FFFFFF"])
|
|
|
|
|
|
class TestPromptAssembly(unittest.TestCase):
|
|
"""Test full prompt assembly logic without ComfyUI runtime."""
|
|
|
|
def test_build_minimal_caption(self):
|
|
# Simulate minimal inputs
|
|
boxes = []
|
|
background = "A simple scene"
|
|
|
|
caption = {
|
|
"compositional_deconstruction": {
|
|
"background": background,
|
|
"elements": []
|
|
}
|
|
}
|
|
|
|
json_str = _dumps(caption)
|
|
self.assertIn('"background":', json_str)
|
|
self.assertIn('"A simple scene"', json_str)
|
|
|
|
def test_build_caption_with_elements(self):
|
|
boxes = [
|
|
{
|
|
"type": "text",
|
|
"text": "HELLO",
|
|
"desc": "Large text",
|
|
"x": 0.1, "y": 0.2, "w": 0.3, "h": 0.4,
|
|
"palette": ["#FF0000", "#00FF00"]
|
|
},
|
|
{
|
|
"type": "obj",
|
|
"desc": "A circle",
|
|
"x": 0.5, "y": 0.5, "w": 0.2, "h": 0.2,
|
|
"palette": []
|
|
}
|
|
]
|
|
|
|
elements = []
|
|
for box in boxes:
|
|
elem = {"type": box["type"]}
|
|
if "nobbox" not in box:
|
|
elem["bbox"] = _norm_bbox(box)
|
|
if box["type"] == "text":
|
|
elem["text"] = box["text"]
|
|
elem["desc"] = box["desc"]
|
|
pal = _palette(box["palette"])
|
|
if pal:
|
|
elem["color_palette"] = pal
|
|
elements.append(elem)
|
|
|
|
self.assertEqual(len(elements), 2)
|
|
self.assertEqual(elements[0]["type"], "text")
|
|
self.assertEqual(elements[0]["text"], "HELLO")
|
|
self.assertEqual(elements[1]["type"], "obj")
|
|
self.assertEqual(len(elements[0]["color_palette"]), 2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|