"""Image linking picks the positive prompt, preferring the image's own metadata.""" import os import sys import tempfile import unittest import unittest.mock as mock sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from database.operations import PromptDatabase from utils.hashing import generate_prompt_hash from utils.image_monitor import ImageGenerationHandler from utils.prompt_tracker import PromptTracker LOADER = { "class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "m.safetensors"}, } def graph(positive_text="a cat", negative_text="blurry"): # Negative listed first: dict order must not decide which prompt an image gets return { "176": {"class_type": "PromptManager", "inputs": {"text": negative_text}}, "134": {"class_type": "PromptManager", "inputs": {"text": positive_text}}, "5": { "class_type": "KSampler", "inputs": { "positive": ["134", 0], "negative": ["176", 0], "model": ["20", 0], }, }, "20": LOADER, } class LinkingTestCase(unittest.TestCase): def setUp(self): fd, self.db_path = tempfile.mkstemp(suffix=".db") os.close(fd) fd, self.image = tempfile.mkstemp(suffix=".png") os.close(fd) self.db = PromptDatabase(self.db_path) self.tracker = PromptTracker(self.db) self.handler = ImageGenerationHandler(self.db, self.tracker) self.linked = [] self.handler.link_image_to_prompt = ( lambda path, prompt, meta: self.linked.append(prompt["id"]) ) def tearDown(self): self.db.close_all() # Windows cannot unlink an open database file for path in ( self.image, self.db_path, self.db_path + "-wal", self.db_path + "-shm", ): if os.path.exists(path): os.unlink(path) def _save(self, text): return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text)) def _process(self, metadata): with mock.patch.object( self.handler.metadata_extractor, "extract_metadata", return_value=metadata ): self.handler.process_new_image(self.image) return self.linked[-1] if self.linked else None class TestMetadataLinking(LinkingTestCase): def test_links_the_positive_prompt_not_the_first_prompt_manager_node(self): positive = self._save("a cat") self._save("blurry") self.assertEqual(self._process({"prompt": graph()}), positive) def test_image_metadata_wins_over_a_stale_queue_entry(self): positive = self._save("a cat") stale = self._save("previous run") self.tracker.set_current_prompt("previous run", {"prompt_id": stale}) self.assertEqual(self._process({"prompt": graph()}), positive) self.assertEqual( self.tracker.pop_next_prompt()["id"], stale, "queue left intact" ) def test_text_from_a_known_string_node_links_from_metadata(self): positive = self._save("a cat") linked_graph = graph() linked_graph["134"]["inputs"]["text"] = ["8", 0] linked_graph["8"] = { "class_type": "PrimitiveString", "inputs": {"value": "a cat"}, } self.assertEqual(self._process({"prompt": linked_graph}), positive) def test_batch_items_with_linked_text_use_the_queue(self): first, second = self._save("item one"), self._save("item two") batch_graph = graph() batch_graph["134"]["inputs"]["text"] = ["7", 0] batch_graph["7"] = {"class_type": "PromptSearchList", "inputs": {}} self.tracker.set_current_prompt("item one", {"prompt_id": first}) self.tracker.set_current_prompt("item two", {"prompt_id": second}) self.assertEqual(self._process({"prompt": batch_graph}), first) self.assertEqual(self._process({"prompt": batch_graph}), second) def test_api_graph_present_skips_the_role_unaware_workflow_fallback(self): negative = self._save("blurry") batch_graph = graph() batch_graph["134"]["inputs"]["text"] = ["7", 0] metadata = { "prompt": batch_graph, "text_encoder_nodes": [ {"type": "PromptManager", "inputs": [], "widgets_values": ["blurry"]} ], } self.assertEqual( self.handler._find_prompt_from_metadata(metadata), ("unknown", None) ) latest = self._save("most recent prompt") # recency fallback picks this self.assertEqual(self._process(metadata), latest) self.assertNotEqual(latest, negative) def test_workflow_fallback_still_used_without_an_api_graph(self): positive = self._save("from workflow only") metadata = { "text_encoder_nodes": [ { "type": "PromptManager", "inputs": [], "widgets_values": ["from workflow only"], } ] } state, prompt = self.handler._find_prompt_from_metadata(metadata) self.assertEqual(state, "linked") self.assertEqual(prompt["id"], positive) def foreign_graph(): """A workflow with no PromptManager node at all.""" return { "9": { "class_type": "CLIPTextEncode", "inputs": {"text": "a cat", "clip": ["20", 1]}, }, "5": { "class_type": "KSampler", "inputs": {"positive": ["9", 0], "negative": ["9", 0], "model": ["20", 0]}, }, "20": LOADER, } class TestForeignWorkflows(LinkingTestCase): """Images from workflows without a positive PromptManager node are not ours.""" def _queue_batch_item(self, text): prompt_id = self._save(text) self.tracker.set_current_prompt(text, {"prompt_id": prompt_id}) return prompt_id def test_foreign_image_leaves_the_batch_queue_aligned(self): first = self._queue_batch_item("item one") self.assertIsNone(self._process({"prompt": foreign_graph()})) self.assertEqual(self.linked, []) self.assertEqual(self.tracker.pop_next_prompt()["id"], first) self.assertIsNone(self.tracker.pop_next_prompt()) def test_foreign_image_is_not_linked_by_the_later_fallbacks_either(self): self._queue_batch_item("item one") # live tracker and recent prompt exist self.assertIsNone(self._process({"prompt": foreign_graph()})) self.assertEqual(self.linked, []) def test_negative_only_workflow_is_foreign(self): self._save("blurry") negative_only = foreign_graph() negative_only["176"] = { "class_type": "PromptManager", "inputs": {"text": "blurry"}, } negative_only["5"]["inputs"]["negative"] = ["176", 0] self.assertEqual( self.handler._find_prompt_from_metadata({"prompt": negative_only}), ("foreign", None), ) def test_workflow_only_metadata_without_prompt_manager_is_foreign(self): metadata = { "text_encoder_nodes": [ {"type": "CLIPTextEncode", "inputs": [], "widgets_values": ["a cat"]} ] } self.assertEqual( self.handler._find_prompt_from_metadata(metadata), ("foreign", None) ) def test_missing_metadata_is_unknown_and_uses_the_queue(self): first = self._queue_batch_item("item one") self.assertEqual( self.handler._find_prompt_from_metadata(None), ("unknown", None) ) self.assertEqual(self._process(None), first) def test_tri_state_for_linked_unknown_and_foreign(self): positive = self._save("a cat") state, prompt = self.handler._find_prompt_from_metadata({"prompt": graph()}) self.assertEqual((state, prompt["id"]), ("linked", positive)) self.assertEqual( self.handler._find_prompt_from_metadata({"prompt": foreign_graph()}), ("foreign", None), ) self.assertEqual( self.handler._find_prompt_from_metadata({"workflow": {}}), ("unknown", None), ) self.assertEqual( self.handler._find_prompt_from_metadata({"prompt": graph("not saved")}), ("unknown", None), ) def test_unknown_image_without_any_context_falls_back_to_the_latest_prompt(self): self._save("older") latest = self._save("latest") self.assertEqual(self._process(None), latest) def test_recent_prompt_fallback_handles_an_empty_database(self): self.assertIsNone(self.handler._get_fallback_prompt()) self.assertIsNone(self._process(None)) def test_metadata_lookup_errors_are_unknown(self): self._save("a cat") with mock.patch.object( self.db, "get_prompt_by_hash", side_effect=RuntimeError("locked") ): self.assertEqual( self.handler._find_prompt_from_metadata({"prompt": graph()}), ("unknown", None), ) if __name__ == "__main__": unittest.main()