import importlib import json import os import sys import tempfile import time import types import unittest from unittest import mock class FakeImageArray: def __init__(self, width=8, height=6): self.shape = (height, width, 3) def cpu(self): return self def numpy(self): return self def astype(self, _dtype): return self def __rmul__(self, _value): return self class FakeTensorBatch: def __init__(self, width=8, height=6, count=1): self.image = FakeImageArray(width=width, height=height) self.shape = (count, height, width, 3) self.count = count def __getitem__(self, index): if isinstance(index, slice): return FakeTensorBatch( width=self.shape[2], height=self.shape[1], count=len(range(*index.indices(self.count))), ) if index != 0: raise IndexError(index) return self.image class FakePILImage: saved_paths = [] def save(self, path, *args, **kwargs): self.saved_paths.append((path, args, kwargs)) def convert(self, _mode): return self class DumasImageNodeTests(unittest.TestCase): @classmethod def setUpClass(cls): cls.temp_dir = tempfile.mkdtemp(prefix="dumas-image-node-") fake_numpy = types.SimpleNamespace( clip=lambda array, _low, _high: array, uint8="uint8", ) fake_pil_image_module = types.SimpleNamespace(fromarray=lambda _array: FakePILImage()) fake_pil_module = types.SimpleNamespace(Image=fake_pil_image_module) fake_folder_paths = types.SimpleNamespace( get_temp_directory=lambda: cls.temp_dir, get_output_directory=lambda: cls.temp_dir, get_save_image_path=lambda prefix, _out, _width, _height: ( cls.temp_dir, prefix, 1, "", prefix, ), ) cls._saved_modules = { name: sys.modules.get(name) for name in ("numpy", "PIL", "PIL.Image", "folder_paths") } sys.modules["numpy"] = fake_numpy sys.modules["PIL"] = fake_pil_module sys.modules["PIL.Image"] = fake_pil_image_module sys.modules["folder_paths"] = fake_folder_paths cls.image_nodes = importlib.import_module("dumas_image_nodes") @classmethod def tearDownClass(cls): for name, module in cls._saved_modules.items(): if module is None: sys.modules.pop(name, None) else: sys.modules[name] = module def setUp(self): FakePILImage.saved_paths = [] def test_compare_images_returns_second_input_as_new_image(self): node = self.image_nodes.DumasImageCompareNode() image1 = FakeTensorBatch() image2 = FakeTensorBatch() result = node.compare_images(image1=image1, image2=image2) self.assertIs(result["result"][0], image2) self.assertEqual([item["slot"] for item in result["ui"]["images"]], [1, 2]) self.assertEqual(len(FakePILImage.saved_paths), 2) def test_compare_images_falls_back_to_first_image(self): node = self.image_nodes.DumasImageCompareNode() image1 = FakeTensorBatch() result = node.compare_images(image1=image1) self.assertIs(result["result"][0], image1) self.assertEqual([item["slot"] for item in result["ui"]["images"]], [1]) def test_compare_images_handles_missing_inputs(self): node = self.image_nodes.DumasImageCompareNode() result = node.compare_images() self.assertIsNone(result["result"][0]) self.assertEqual(result["ui"]["images"], []) def test_saved_filenames_use_dumas_prefix(self): node = self.image_nodes.DumasImageCompareNode() image1 = FakeTensorBatch() result = node.compare_images(image1=image1) self.assertTrue(result["ui"]["images"][0]["filename"].startswith("dumas_compare")) self.assertTrue(os.path.basename(FakePILImage.saved_paths[0][0]).startswith("dumas_compare")) def test_save_image_uses_second_input_token(self): node = self.image_nodes.DumasSaveImageNode() image = FakeTensorBatch(width=10, height=12) node.save_images( images=image, folder=self.temp_dir, pattern="shot_%input%_%input2%_%counter%", format="png", quality=100, embed_workflow=False, save_on_run=True, name="alpha.png", name_2="beta/final", ) saved_name = os.path.basename(FakePILImage.saved_paths[0][0]) self.assertEqual(saved_name, "shot_alpha_beta_final_001.png") def test_save_image_resolves_date_size_and_batch_tokens(self): node = self.image_nodes.DumasSaveImageNode() image = FakeTensorBatch(width=10, height=12, count=2) fake_now = time.struct_time((2026, 8, 5, 13, 7, 9, 2, 217, -1)) with mock.patch.object(self.image_nodes.time, "localtime", return_value=fake_now): node.save_images( images=image, folder=self.temp_dir, pattern="asset_%date:yyyy-MM-dd%_%date:hh-mm-ss%_%width%x%height%_%batch_num%_%counter%", format="png", quality=100, embed_workflow=False, save_on_run=True, ) saved_names = [os.path.basename(path) for path, _args, _kwargs in FakePILImage.saved_paths] self.assertEqual( saved_names, [ "asset_2026-08-05_13-07-09_10x12_0_001.png", "asset_2026-08-05_13-07-09_10x12_1_001.png", ], ) def test_save_image_counter_increments_for_existing_files(self): node = self.image_nodes.DumasSaveImageNode() image = FakeTensorBatch() node.save_images( images=image, folder=self.temp_dir, pattern="counter_%counter%", format="png", quality=100, embed_workflow=False, save_on_run=True, ) open(FakePILImage.saved_paths[0][0], "a", encoding="utf-8").close() node.save_images( images=image, folder=self.temp_dir, pattern="counter_%counter%", format="png", quality=100, embed_workflow=False, save_on_run=True, ) saved_names = [os.path.basename(path) for path, _args, _kwargs in FakePILImage.saved_paths] self.assertEqual(saved_names, ["counter_001.png", "counter_002.png"]) def test_save_image_uses_same_counter_for_folder_and_filename(self): node = self.image_nodes.DumasSaveImageNode() image = FakeTensorBatch() node.save_images( images=image, folder=self.temp_dir, pattern="Char%counter%/Char%counter%", format="png", quality=100, embed_workflow=False, save_on_run=True, ) saved_path = FakePILImage.saved_paths[0][0].replace("\\", "/") self.assertTrue(saved_path.endswith("/Char001/Char001.png")) def test_save_image_creates_nested_directories_before_saving(self): node = self.image_nodes.DumasSaveImageNode() image = FakeTensorBatch() node.save_images( images=image, folder=self.temp_dir, pattern="Char_%counter%/Base", format="png", quality=100, embed_workflow=False, save_on_run=True, ) saved_path = FakePILImage.saved_paths[0][0] self.assertTrue(os.path.isdir(os.path.dirname(saved_path))) def test_character_reference_builds_structured_reference(self): node = self.image_nodes.DumasCharacterReferenceNode() image = FakeTensorBatch() result = node.build_reference( image=image, character_id="char_dave", name="Dave", alias="The Locksmith", gender="male", age="41", nationality="English", occupation="a detective", height_feet="6", height_inches="2", accent="English", description="Square jaw, tired eyes, cropped brown hair.", general="wears a long grey coat", wardrobe="weathered red flight jacket, grey cargo shorts, black boots", ) reference = result[0] self.assertIs(reference["image"], image) self.assertEqual( reference, { "kind": "character", "id": "char-dave", "name": "Dave", "aliases": ["The Locksmith"], "picture_id": None, "picture_label": "", "image": image, "summary": "Dave reference.", "description": "Square jaw, tired eyes, cropped brown hair.", "wardrobe": "weathered red flight jacket, grey cargo shorts, black boots", "general": "wears a long grey coat", "facts": { "gender": "male", "age": "41", "nationality": "English", "occupation": "a detective", "height_feet": "6", "height_inches": "2", "accent": "English", }, }, ) def test_character_reference_input_types_do_not_expose_picture_id(self): required = self.image_nodes.DumasCharacterReferenceNode.INPUT_TYPES()["required"] self.assertNotIn("picture_id", required) def test_character_reference_handles_missing_optional_fields(self): node = self.image_nodes.DumasCharacterReferenceNode() image = FakeTensorBatch() result = node.build_reference( image=image, character_id="", name="", alias="", gender="", age="unknown", nationality="", occupation="", height_feet="", height_inches="", accent="", description="", general="", wardrobe="", ) reference = result[0] self.assertEqual(reference["kind"], "character") self.assertIsNone(reference["picture_id"]) self.assertEqual(reference["picture_label"], "") self.assertEqual(reference["wardrobe"], "") self.assertEqual(reference["general"], "") self.assertEqual(reference["facts"]["age"], "") def test_location_reference_builds_structured_reference(self): node = self.image_nodes.DumasLocationReferenceNode() image = FakeTensorBatch() result = node.build_reference( image=image, location_id="coffee-shop-01", name="Coffee Shop", alias="Cafe Interior", description="Warm tungsten lighting, narrow counter, rainy front window.", general="Evening ambience, cramped but cozy.", ) self.assertEqual( result[0], { "kind": "location", "id": "coffee-shop-01", "name": "Coffee Shop", "aliases": ["Cafe Interior"], "picture_id": None, "picture_label": "", "image": image, "summary": "Coffee Shop reference.", "description": "Warm tungsten lighting, narrow counter, rainy front window.", "wardrobe": "", "general": "Evening ambience, cramped but cozy.", "facts": {}, }, ) def test_location_reference_input_types_do_not_expose_picture_id(self): required = self.image_nodes.DumasLocationReferenceNode.INPUT_TYPES()["required"] self.assertNotIn("picture_id", required) def test_character_helper_restores_image_and_text_outputs(self): node = self.image_nodes.DumasCharacterHelperNode() image1 = FakeTensorBatch() image2 = FakeTensorBatch() result = node.build_character_text( image1=image1, image2=image2, image1_picture_id="1", image2_picture_id="2", character_id="char_dave", name="Dave", alias="The Locksmith", gender="male", age="41", nationality="English", occupation="a detective", height_feet="6", height_inches="2", accent="English", general="Moves carefully and notices every exit", wardrobe="weathered red flight jacket, grey cargo shorts, black boots", ) self.assertIs(result[0], image1) self.assertIs(result[1], image2) self.assertIn(" and reference the same character", result[2]) self.assertIn("Dave is also known as The Locksmith", result[2]) self.assertIn("is 41 years old", result[2]) self.assertEqual(result[3], "Dave = weathered red flight jacket, grey cargo shorts, black boots") self.assertIs(result[4]["image"], image1) self.assertIs(result[5]["image"], image2) self.assertEqual(result[4]["id"], "char-dave") self.assertEqual(result[5]["id"], "char-dave") self.assertEqual(result[4]["name"], "Dave") self.assertEqual(result[4]["aliases"], ["The Locksmith"]) self.assertEqual(result[4]["facts"]["age"], "41") self.assertEqual(result[4]["facts"]["height_feet"], "6") self.assertEqual(result[4]["facts"]["height_inches"], "2") self.assertEqual(result[4]["wardrobe"], "weathered red flight jacket, grey cargo shorts, black boots") self.assertEqual(len(result), 6) def test_location_helper_matches_character_helper_shape_without_wardrobe(self): node = self.image_nodes.DumasLocationHelperNode() image1 = FakeTensorBatch() image2 = FakeTensorBatch() result = node.build_location_text( image1=image1, image2=image2, image1_picture_id="3", image2_picture_id="4", location_id="coffee-shop-01", name="Coffee Shop", alias="Cafe Interior", description="Warm tungsten lighting, narrow counter, rainy front window", general="Evening ambience, cramped but cozy", ) self.assertIs(result[0], image1) self.assertIs(result[1], image2) self.assertIn(" and reference the same location", result[2]) self.assertIn("Coffee Shop is also known as Cafe Interior", result[2]) self.assertIn("Warm tungsten lighting, narrow counter, rainy front window.", result[2]) self.assertIn("Evening ambience, cramped but cozy.", result[2]) self.assertEqual(len(result), 3) def test_soundscape_helper_defaults_to_selected_preset_description(self): node = self.image_nodes.DumasSoundscapeHelperNode() result = node.build_soundscape("rainy street", "") self.assertEqual(result[0], "steady rain, wet pavement, distant traffic hum") def test_h3_prompt_curator_compacts_named_references(self): node = self.image_nodes.DumasH3PromptCuratorNode() dave_image = FakeTensorBatch() cafe_image = FakeTensorBatch() van_image = FakeTensorBatch() dave = self.image_nodes.make_reference( kind="character", image=dave_image, name="Dave", aliases="The Locksmith", description="tired eyes, cropped brown hair", wardrobe="red flight jacket", ) cafe = self.image_nodes.make_reference( kind="location", image=cafe_image, name="Coffee Shop", description="warm tungsten lighting and rainy windows", ) van = self.image_nodes.make_reference( kind="location", image=van_image, name="Blue Van", description="scuffed blue delivery van", ) result = node.curate_prompt( action_prompt="Dave runs from the Coffee Shop into the rain.", anatomy_guard="auto", subject_count_guard="auto", anchor="grounded handheld thriller", soundscape="steady rain", ref_1=dave, ref_2=van, ref_3=cafe, ) prompt = result[0] self.assertIn(" Dave", prompt) self.assertIn(" Coffee Shop", prompt) self.assertIn("Action: Dave runs from the Coffee Shop into the rain.", prompt) self.assertIn("Anatomy guard:", prompt) self.assertIn("Subject count guard:", prompt) self.assertIn("exactly one named character: Dave", prompt) self.assertIs(result[1], dave_image) self.assertIs(result[2], cafe_image) self.assertIsNone(result[3]) self.assertEqual(result[10], 2) self.assertIn("input 3-> Coffee Shop", result[11]) def test_h3_prompt_curator_renumbers_explicit_reference_tags(self): node = self.image_nodes.DumasH3PromptCuratorNode() image1 = FakeTensorBatch() image3 = FakeTensorBatch() unused = FakeTensorBatch() first = self.image_nodes.make_reference(kind="character", image=image1, name="Maya") second = self.image_nodes.make_reference(kind="location", image=unused, name="Lobby") third = self.image_nodes.make_reference(kind="location", image=image3, name="Rooftop") result = node.curate_prompt( action_prompt=" Maya crosses to as the wind rises.", anatomy_guard="off", subject_count_guard="off", ref_1=first, ref_2=second, ref_3=third, ) prompt = result[0] self.assertIn(" Maya crosses to ", prompt) self.assertNotIn("", prompt) self.assertIs(result[1], image1) self.assertIs(result[2], image3) self.assertIsNone(result[3]) self.assertEqual(result[10], 2) def test_h3_prompt_curator_can_force_subject_count_without_character_refs(self): node = self.image_nodes.DumasH3PromptCuratorNode() result = node.curate_prompt( action_prompt="A locked-off shot of the empty corridor.", anatomy_guard="off", subject_count_guard="on", ) self.assertIn("Subject count guard:", result[0]) self.assertIn("Only include the people explicitly described", result[0]) self.assertEqual(result[10], 0) def test_h3_prompt_curator_treats_helper_image_pair_as_one_character(self): helper = self.image_nodes.DumasCharacterHelperNode() curator = self.image_nodes.DumasH3PromptCuratorNode() image1 = FakeTensorBatch() image2 = FakeTensorBatch() helper_result = helper.build_character_text( image1=image1, image2=image2, image1_picture_id="1", image2_picture_id="2", character_id="char_dave", name="Dave", alias="The Locksmith", gender="male", age="41", nationality="English", occupation="detective", height_feet="6", height_inches="2", accent="English", general="Tired eyes, cropped brown hair", wardrobe="weathered red flight jacket", ) result = curator.curate_prompt( action_prompt="Dave checks the locked door.", anatomy_guard="on", subject_count_guard="auto", ref_1=helper_result[4], ref_2=helper_result[5], ) self.assertIs(result[1], image1) self.assertIs(result[2], image2) self.assertEqual(result[10], 2) self.assertIn("Character facts for Dave", result[0]) self.assertIn("41 years old", result[0]) self.assertIn("6 foot 2 tall", result[0]) self.assertIn("exactly one named character: Dave", result[0]) self.assertNotIn("exactly 2 named characters", result[0]) def test_h3_prompt_curator_defaults_anatomy_guard_to_on(self): required = self.image_nodes.DumasH3PromptCuratorNode.INPUT_TYPES()["required"] self.assertEqual(required["anatomy_guard"][1]["default"], "on") def test_helper_node_mappings_use_general_purpose_helpers(self): mappings = self.image_nodes.NODE_CLASS_MAPPINGS display = self.image_nodes.NODE_DISPLAY_NAME_MAPPINGS self.assertIs(mappings["DumasCharacterHelper"], self.image_nodes.DumasCharacterHelperNode) self.assertIs(mappings["DumasLocationHelper"], self.image_nodes.DumasLocationHelperNode) self.assertIs(mappings["DumasSoundscapeHelper"], self.image_nodes.DumasSoundscapeHelperNode) self.assertIs(mappings["DumasH3PromptCurator"], self.image_nodes.DumasH3PromptCuratorNode) self.assertEqual(display["DumasCharacterHelper"], "Dumas Character Helper") self.assertEqual(display["DumasLocationHelper"], "Dumas Location Helper") self.assertEqual(display["DumasSoundscapeHelper"], "Dumas Soundscape Helper") self.assertEqual(display["DumasH3PromptCurator"], "Dumas H3 Prompt Curator") def test_h3_prompt_curator_uses_documented_reference_limits(self): node = self.image_nodes.DumasH3PromptCuratorNode() self.assertEqual(len(node.RETURN_TYPES), 12) self.assertEqual(node.RETURN_NAMES[1:10], tuple(f"ref_image_{i}" for i in range(1, 10))) def test_normalize_reference_upgrades_generic_summary_with_socket_picture_id(self): image = FakeTensorBatch() reference = self.image_nodes.normalize_reference( { "kind": "character", "name": "Dave", "image": image, "summary": "Dave reference.", }, picture_id=3, allow_image_fallback=False, ) self.assertEqual(reference["picture_id"], 3) self.assertEqual(reference["picture_label"], "") self.assertEqual(reference["summary"], "Dave shown in .") def test_normalize_reference_keeps_custom_summary_when_socket_picture_id_is_added(self): image = FakeTensorBatch() reference = self.image_nodes.normalize_reference( { "kind": "character", "name": "Dave", "image": image, "summary": "Primary hero look for the opening close-up.", }, picture_id=3, allow_image_fallback=False, ) self.assertEqual(reference["picture_id"], 3) self.assertEqual(reference["summary"], "Primary hero look for the opening close-up.") def test_anchor_style_node_exposes_requested_presets(self): input_types = self.image_nodes.DumasAnchorStyleNode.INPUT_TYPES() options = input_types["required"]["anchor_style"][0] self.assertGreaterEqual(len(options), 40) self.assertIn("cinematic action movie", options) self.assertIn("comedy", options) self.assertIn("found footage", options) self.assertIn("90s sitcom", options) self.assertIn("mobile/cell phone captured", options) self.assertIn("news broadcast", options) self.assertIn("mockumentary", options) self.assertIn("heist thriller", options) self.assertIn("cyberpunk neon", options) self.assertIn("nature documentary", options) self.assertIn("courtroom drama", options) def test_anchor_style_node_defaults_to_selected_preset_description(self): node = self.image_nodes.DumasAnchorStyleNode() result = node.build_anchor("found footage", "") self.assertIn("found-footage", result[0]) self.assertIn("real time", result[0]) self.assertNotIn("persistent camera language", result[0]) def test_anchor_style_node_prefers_manual_description_edits(self): node = self.image_nodes.DumasAnchorStyleNode() custom = "Lo-fi pirate broadcast with smeared highlights and anxious zoom corrections." result = node.build_anchor("news broadcast", custom) self.assertEqual(result[0], custom) def test_save_image_returns_ui_entries_for_output_folder(self): node = self.image_nodes.DumasSaveImageNode() image = FakeTensorBatch() result = node.save_images( images=image, folder="", pattern="result_%counter%", format="jpg", quality=90, embed_workflow=False, save_on_run=True, ) self.assertEqual(result["ui"]["images"][0]["type"], "output") self.assertTrue(result["ui"]["images"][0]["filename"].endswith(".jpg")) def test_save_image_returns_ui_entries_for_external_folder(self): node = self.image_nodes.DumasSaveImageNode() image = FakeTensorBatch() external_dir = tempfile.mkdtemp(prefix="dumas-image-node-external-") result = node.save_images( images=image, folder=external_dir, pattern="external_%counter%", format="png", quality=100, embed_workflow=False, save_on_run=True, ) self.assertEqual(result["ui"]["images"][0]["type"], "external") self.assertEqual(result["ui"]["images"][0]["subfolder"], external_dir.replace("\\", "/")) self.assertTrue(result["ui"]["images"][0]["token"]) def test_save_image_handles_windows_different_drive_paths(self): node = self.image_nodes.DumasSaveImageNode() image = FakeTensorBatch() with mock.patch.object( self.image_nodes.folder_paths, "get_output_directory", return_value="C:\\ComfyUI\\output", ), mock.patch.object( self.image_nodes.os.path, "abspath", return_value="D:\\renders", ): result = node.save_images( images=image, folder="D:\\renders", pattern="drive_%counter%", format="png", quality=100, embed_workflow=False, save_on_run=True, ) self.assertEqual(result["ui"]["images"][0]["type"], "external") self.assertEqual(result["ui"]["images"][0]["subfolder"], "D:/renders") self.assertTrue(result["ui"]["images"][0]["token"]) def test_save_image_skips_when_save_is_disabled(self): node = self.image_nodes.DumasSaveImageNode() image = FakeTensorBatch() result = node.save_images( images=image, folder=self.temp_dir, pattern="ignored_%counter%", format="png", quality=100, embed_workflow=False, save_on_run=False, ) self.assertEqual(result["ui"]["images"], []) self.assertEqual(FakePILImage.saved_paths, []) def test_load_images_folder_uses_manual_selected_files(self): node = self.image_nodes.DumasLoadImagesFolderNode() folder = tempfile.mkdtemp(prefix="dumas-load-folder-manual-") for name in ("b.png", "a.png", "notes.txt"): open(os.path.join(folder, name), "a", encoding="utf-8").close() state = json.dumps( { "folder": folder, "recursive": False, "sort": "name", "sort_dir": "asc", "selection_mode": "selected", "selected": ["b.png", "a.png"], } ) with mock.patch.object( self.image_nodes, "_load_folder_image", side_effect=lambda path: (f"image:{os.path.basename(path)}", f"mask:{os.path.basename(path)}", 32, 24), ): result = node.load(state) self.assertEqual(result[0], ["image:b.png", "image:a.png"]) self.assertEqual(result[1], ["mask:b.png", "mask:a.png"]) self.assertEqual(result[4], ["b", "a"]) self.assertEqual(result[5], [1, 2]) self.assertEqual(result[6], [2, 2]) def test_load_images_folder_first_n_uses_sorted_files(self): node = self.image_nodes.DumasLoadImagesFolderNode() folder = tempfile.mkdtemp(prefix="dumas-load-folder-first-") for name in ("c.png", "a.png", "b.png"): open(os.path.join(folder, name), "a", encoding="utf-8").close() state = json.dumps( { "folder": folder, "recursive": False, "sort": "name", "sort_dir": "asc", "selection_mode": "first_n", "first_n": 2, } ) with mock.patch.object( self.image_nodes, "_load_folder_image", side_effect=lambda path: (f"image:{os.path.basename(path)}", f"mask:{os.path.basename(path)}", 64, 48), ): result = node.load(state) self.assertEqual(result[0], ["image:a.png", "image:b.png"]) self.assertEqual(result[4], ["a", "b"]) self.assertEqual(result[6], [2, 2]) def test_load_images_folder_random_selects_one_visible_image(self): node = self.image_nodes.DumasLoadImagesFolderNode() folder = tempfile.mkdtemp(prefix="dumas-load-folder-random-") for name in ("a.png", "b.png", "c.png"): open(os.path.join(folder, name), "a", encoding="utf-8").close() state = json.dumps( { "folder": folder, "recursive": False, "sort": "name", "sort_dir": "asc", "selection_mode": "random", } ) with mock.patch.object( self.image_nodes.random, "choice", side_effect=lambda files: files[1], ), mock.patch.object( self.image_nodes, "_load_folder_image", side_effect=lambda path: (f"image:{os.path.basename(path)}", f"mask:{os.path.basename(path)}", 80, 60), ): result = node.load(state) self.assertEqual(result[0], ["image:b.png"]) self.assertEqual(result[4], ["b"]) self.assertEqual(result[5], [1]) self.assertEqual(result[6], [1]) def test_load_images_folder_random_is_changed_changes_each_run(self): node = self.image_nodes.DumasLoadImagesFolderNode() folder = tempfile.mkdtemp(prefix="dumas-load-folder-changed-") open(os.path.join(folder, "a.png"), "a", encoding="utf-8").close() state = json.dumps( { "folder": folder, "recursive": False, "sort": "name", "sort_dir": "asc", "selection_mode": "random", } ) with mock.patch.object(self.image_nodes.time, "time_ns", side_effect=[1, 2]): changed_1 = node.IS_CHANGED(state) changed_2 = node.IS_CHANGED(state) self.assertNotEqual(changed_1, changed_2) def test_h3_plan_scene_images_attach_and_extract(self): attach_node = self.image_nodes.DumasH3PlanAttachSceneImagesNode() extract_node = self.image_nodes.DumasH3PlanExtractSceneImagesNode() plan = {"shots": [{"id": "intro"}, {"id": "middle"}]} image_a = FakeTensorBatch() image_b = FakeTensorBatch(width=10, height=10) attached_plan, connected = attach_node.attach( plan=plan, scene_index=2, image1=image_a, image3=image_b, ) extracted = extract_node.extract(attached_plan, 2) self.assertEqual(connected, 2) self.assertEqual(attached_plan["_dumas_scene_image_bindings"]["scene_counts"], {"2": 2}) self.assertIs(extracted[1], image_a) self.assertIsNone(extracted[2]) self.assertIs(extracted[3], image_b) self.assertEqual(extracted[-1], 2) self.assertNotIn("_dumas_scene_image_bindings", plan) def test_h3_plan_scene_images_support_multiple_scenes(self): attach_node = self.image_nodes.DumasH3PlanAttachSceneImagesNode() extract_node = self.image_nodes.DumasH3PlanExtractSceneImagesNode() plan = {"shots": [{"id": "one"}, {"id": "two"}]} image_1 = FakeTensorBatch() image_2 = FakeTensorBatch(width=12, height=9) image_3 = FakeTensorBatch(width=8, height=8) image_4 = FakeTensorBatch(width=16, height=16) image_5 = FakeTensorBatch(width=20, height=12) plan_after_first, _connected = attach_node.attach(plan=plan, scene_index=1, image2=image_1) plan_after_second, _connected = attach_node.attach( plan=plan_after_first, scene_index=2, image6=image_2, image7=image_3, image8=image_4, image9=image_5, ) scene1 = extract_node.extract(plan_after_second, 1) scene2 = extract_node.extract(plan_after_second, 2) self.assertIs(scene1[2], image_1) self.assertEqual(scene1[-1], 1) self.assertIs(scene2[6], image_2) self.assertIs(scene2[7], image_3) self.assertIs(scene2[8], image_4) self.assertIs(scene2[9], image_5) self.assertEqual(scene2[-1], 4) self.assertEqual( plan_after_second["_dumas_scene_image_bindings"]["scene_counts"], {"1": 1, "2": 4}, ) def test_h3_plan_scene_images_metadata_is_json_serializable(self): attach_node = self.image_nodes.DumasH3PlanAttachSceneImagesNode() plan = {"shots": [{"id": "one"}]} attached_plan, connected = attach_node.attach( plan=plan, scene_index=1, image4=FakeTensorBatch(), ) self.assertEqual(connected, 1) json.dumps(attached_plan) def test_h3_plan_scene_images_support_nine_slots(self): attach_node = self.image_nodes.DumasH3PlanAttachSceneImagesNode() extract_node = self.image_nodes.DumasH3PlanExtractSceneImagesNode() plan = {"shots": [{"id": "one"}]} images = [FakeTensorBatch(width=8 + index, height=8 + index) for index in range(9)] attached_plan, connected = attach_node.attach( plan=plan, scene_index=1, image1=images[0], image2=images[1], image3=images[2], image4=images[3], image5=images[4], image6=images[5], image7=images[6], image8=images[7], image9=images[8], ) extracted = extract_node.extract(attached_plan, 1) self.assertEqual(connected, 9) for index, image in enumerate(images, start=1): self.assertIs(extracted[index], image) self.assertEqual(extracted[-1], 9) def test_h3_plan_scene_images_can_clear_a_scene_binding(self): attach_node = self.image_nodes.DumasH3PlanAttachSceneImagesNode() extract_node = self.image_nodes.DumasH3PlanExtractSceneImagesNode() plan = {"shots": [{"id": "one"}]} attached_plan, connected = attach_node.attach( plan=plan, scene_index=1, image1=FakeTensorBatch(), ) cleared_plan, cleared = attach_node.attach(plan=attached_plan, scene_index=1) extracted = extract_node.extract(cleared_plan, 1) self.assertEqual(connected, 1) self.assertEqual(cleared, 0) self.assertNotIn("_dumas_scene_image_bindings", cleared_plan) self.assertEqual(extracted[-1], 0) def test_h3_plan_scene_images_reject_invalid_scene_index(self): attach_node = self.image_nodes.DumasH3PlanAttachSceneImagesNode() plan = {"shots": [{"id": "one"}]} with self.assertRaises(ValueError): attach_node.attach(plan=plan, scene_index=2, image1=FakeTensorBatch()) if __name__ == "__main__": unittest.main()