"""Paired node encoding preserves recipe storage and output roles.""" import logging import os import tempfile import unittest from unittest.mock import Mock try: from tests.open_handles import assert_closed except ImportError: # discovered with tests/ as the top-level directory from open_handles import assert_closed from database.operations import PromptDatabase from prompt_manager_pair import PromptManagerPair from utils.hashing import generate_prompt_hash from utils.usage_tracking import PendingFirstUse def pair_graph(text="cat", negative="blur", positive_slot=0, negative_slot=1): return { "1": { "class_type": "PromptManagerPair", "inputs": {"text": text, "negative_text": negative}, }, "5": { "class_type": "KSampler", "inputs": { "model": ["9", 0], "positive": ["1", positive_slot], "negative": ["1", negative_slot], }, }, } class TestPromptPairNode(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.addCleanup(self.tmp.cleanup) self.addCleanup(assert_closed, self, self.tmp.name) self.db = PromptDatabase(os.path.join(self.tmp.name, "pairs.db")) self.addCleanup(self.db.close_all) self.node = PromptManagerPair.__new__(PromptManagerPair) self.node.db = self.db self.node.logger = logging.getLogger("test.pair_node") self.node.prompt_tracker = Mock() self.node._inject_lora_trigger_words = Mock( side_effect=lambda text: text + " LORA" ) self.node.pending_first_use = PendingFirstUse() self.node.pending_pair_first_use = PendingFirstUse() self.clip = Mock() self.clip.encode_from_tokens_scheduled.side_effect = lambda tokens: ( "conditioning", tokens, ) self.clip.tokenize.side_effect = lambda text: text def test_schema_has_positive_composition_and_four_ordered_outputs(self): schema = PromptManagerPair.INPUT_TYPES() self.assertEqual(list(schema["required"]), ["text", "negative_text", "clip"]) self.assertTrue(schema["required"]["negative_text"][1]["multiline"]) self.assertEqual(schema["required"]["negative_text"][1]["default"], "") for name in ("prepend_text", "append_text"): self.assertTrue(schema["optional"][name][1]["forceInput"]) self.assertEqual( PromptManagerPair.RETURN_TYPES, ("CONDITIONING", "CONDITIONING", "STRING", "STRING"), ) self.assertNotEqual( PromptManagerPair.IS_CHANGED(None, "cat", negative_text="blur"), PromptManagerPair.IS_CHANGED(None, "cat", negative_text="watermark"), ) def test_encodes_each_side_once_positive_composition_and_lora_only(self): graph = pair_graph() result = self.node.encode_prompt( self.clip, "cat", negative_text=" blur ", prepend_text=" trigger ", append_text=" detail ", category="Animals", tags="studio", prompt=graph, unique_id="1", ) self.assertEqual( result, ( ("conditioning", "trigger cat detail LORA"), ("conditioning", " blur "), "trigger cat detail LORA", " blur ", ), ) self.node._inject_lora_trigger_words.assert_called_once_with( "trigger cat detail" ) self.assertEqual(self.clip.tokenize.call_count, 2) pair = self.db.list_prompt_pairs()["pairs"][0] self.assertEqual( ( pair["positive_text"], pair["negative_text"], pair["category"], pair["tags"], ), ("cat", "blur", "Animals", ["studio"]), ) self.assertEqual( self.db.get_prompt_by_id(pair["positive_prompt_id"])["category"], "Animals", ) tracked = self.node.prompt_tracker.set_current_prompt.call_args.kwargs self.assertEqual(tracked["additional_data"]["pair_id"], pair["id"]) self.assertEqual(tracked["additional_data"]["negative_text"], " blur ") self.assertEqual(tracked["additional_data"]["positive_storage_text"], "cat") def save(self, text="cat", negative="blur", **widgets): self.node.encode_prompt( self.clip, text, negative_text=negative, prompt=pair_graph(text, negative), unique_id="1", **widgets, ) hash_of = generate_prompt_hash return ( self.db.get_prompt_by_hash(hash_of(text)), self.db.get_prompt_by_hash(hash_of(negative)), ) def test_first_save_gives_the_positive_its_category_and_tags(self): positive, negative = self.save(category=" Animals ", tags="studio, macro") self.assertEqual(positive["category"], "Animals") self.assertCountEqual(positive["tags"], ["studio", "macro"]) self.assertIsNone(negative["category"]) self.assertEqual(negative["tags"], []) def test_a_later_save_never_changes_a_category_or_tags_already_set(self): self.save(category="Animals", tags="studio") positive, _ = self.save(negative="noise", category="Pets", tags="other") self.assertEqual(positive["category"], "Animals") self.assertEqual(positive["tags"], ["studio"]) def test_a_save_keeps_metadata_a_person_set_on_the_prompt(self): text = "cat" self.db.save_prompt( text, category="Mine", tags=["kept"], prompt_hash=generate_prompt_hash(text), ) positive, _ = self.save(category="Animals", tags="studio") self.assertEqual(positive["category"], "Mine") self.assertEqual(positive["tags"], ["kept"]) def test_a_run_on_an_existing_empty_prompt_changes_nothing(self): self.db.save_prompt("cat", prompt_hash=generate_prompt_hash("cat")) before = self.db.get_prompt_by_hash(generate_prompt_hash("cat"))["updated_at"] positive, _ = self.save(category="Animals", tags="studio") self.assertIsNone(positive["category"]) self.assertEqual(positive["tags"], []) self.assertEqual(positive["updated_at"], before) def test_a_pair_saved_without_metadata_stays_empty_on_a_later_run(self): self.save() positive, _ = self.save(negative="noise", category="Animals", tags="studio") self.assertIsNone(positive["category"]) self.assertEqual(positive["tags"], []) def test_a_cleared_category_and_tags_stay_cleared(self): positive, _ = self.save(category="Animals", tags="studio") self.db.update_prompt_metadata(positive["id"], category="", tags=[]) positive, _ = self.save(negative="noise", category="Animals", tags="studio") self.assertEqual(positive["category"], "") self.assertEqual(positive["tags"], []) def test_a_prompt_first_seen_as_a_negative_stays_empty(self): self.save(text="dog", negative="cat") positive, _ = self.save(text="cat", negative="blur", category="Animals") self.assertIsNone(positive["category"]) def test_blank_widgets_leave_the_prompt_empty(self): positive, _ = self.save(category=" ", tags=" ") self.assertIsNone(positive["category"]) self.assertEqual(positive["tags"], []) def test_a_filled_prompt_does_not_change_the_saved_pair(self): self.save(category="Animals", tags="studio") positive, _ = self.save(category="Pets", tags="other") pairs = self.db.list_prompt_pairs()["pairs"] self.assertEqual( [(pair["category"], pair["tags"]) for pair in pairs], [("Animals", ["studio"])], ) self.assertEqual(positive["category"], "Animals") def test_library_shows_filled_metadata_without_repeating_a_tag(self): self.save(category="Animals", tags="studio") item = self.db.list_library()["items"][0] self.assertEqual(item["category"], "Animals") self.assertEqual(item["tags"], ["studio"]) tag = self.db.list_vocabulary_tags(limit=5)["tags"][0] self.assertEqual(tag["prompt_count"], 1) def test_empty_negative_encodes_empty_text_and_stores_null_component(self): result = self.node.encode_prompt(self.clip, "cat", negative_text="") self.assertEqual(result[3], "") self.assertIsNone(self.db.list_prompt_pairs()["pairs"][0]["negative_prompt_id"]) def test_invalid_clip_cannot_save_or_count(self): with self.assertRaises(RuntimeError): self.node.encode_prompt(None, "cat", negative_text="blur") self.assertEqual(self.db.list_prompt_pairs()["total"], 0) self.node.prompt_tracker.set_current_prompt.assert_not_called() def test_persistence_failure_does_not_block_conditioning(self): self.node.db = Mock() self.node.db.save_prompt_pair.side_effect = RuntimeError("unavailable") result = self.node.encode_prompt(self.clip, "cat", negative_text="blur") self.assertEqual(result[2:], ("cat LORA", "blur")) self.node.prompt_tracker.set_current_prompt.assert_not_called() def test_negative_output_used_as_positive_does_not_claim_positive_usage(self): self.node.encode_prompt( self.clip, "cat", negative_text="blur", prompt=pair_graph(positive_slot=1, negative_slot=0), unique_id="1", ) pair = self.db.list_prompt_pairs()["pairs"][0] self.assertEqual(pair["run_count"], 0) self.assertEqual( self.db.get_prompt_by_id(pair["positive_prompt_id"])["run_count"], 0 ) self.node.prompt_tracker.set_current_prompt.assert_not_called() def test_invalid_text_does_not_write(self): with self.assertRaises(ValueError): self.node.encode_prompt(self.clip, [], negative_text="blur") result = self.node.encode_prompt(self.clip, "cat", negative_text="blur") self.assertEqual(result[2:], ("cat LORA", "blur")) def test_failed_negative_encoding_cannot_publish_positive_runtime_proof(self): from contextlib import ExitStack import sys import types from utils.execution_context_capture import ContextCapture from utils.execution_context_budget import ContextBudget from utils.execution_history_events import RunStart from utils.execution_history_writer import HistoryWriterStatus from unittest.mock import patch identity = "11111111-1111-4111-8111-111111111111" state = HistoryWriterStatus(identity, True, 0, 0, True, 1) budget = ContextBudget() manager = ContextCapture(1, lambda: state, budget) self.addCleanup(manager.revoke_all) cap = manager.begin( RunStart(identity, identity, 1, identity, identity, 0), 1, pair_graph(), None, ) host = types.ModuleType("comfy_execution.utils") host.get_executing_context = lambda: types.SimpleNamespace( prompt_id=identity, node_id="1", list_index=0 ) self.clip.encode_from_tokens_scheduled.side_effect = [ ("conditioning", "cat LORA"), RuntimeError("negative failed"), ] with ExitStack() as stack: stack.enter_context( patch.dict(sys.modules, {"comfy_execution.utils": host}) ) with self.assertRaisesRegex(RuntimeError, "negative failed"): self.node.encode_prompt(self.clip, "cat", negative_text="blur") sealed = manager.seal(cap) if sealed.lease: self.addCleanup(sealed.lease.release) self.assertIsNone(sealed.bundle.contexts[0].positive.effective.value) self.assertIsNone(sealed.bundle.contexts[0].negative.effective.value) self.assertEqual(self.db.list_prompt_pairs()["total"], 0)