62 lines
2.7 KiB
Python
62 lines
2.7 KiB
Python
"""Focused shares encoding and tracking while exposing composition sockets."""
|
|
|
|
import logging
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
from prompt_manager import PromptManager, PromptManagerFocused
|
|
from utils.prompt_graph import positive_prompt_nodes, run_prompt_nodes
|
|
|
|
|
|
class TestFocusedNode(unittest.TestCase):
|
|
def test_focused_keeps_the_execution_contract_and_visible_string_inputs(self):
|
|
self.assertTrue(issubclass(PromptManagerFocused, PromptManager))
|
|
self.assertEqual(PromptManagerFocused.RETURN_TYPES, PromptManager.RETURN_TYPES)
|
|
self.assertEqual(PromptManagerFocused.FUNCTION, PromptManager.FUNCTION)
|
|
focused = PromptManagerFocused.INPUT_TYPES()
|
|
for name in ("prepend_text", "append_text"):
|
|
self.assertTrue(focused["optional"][name][1]["forceInput"])
|
|
self.assertNotIn(
|
|
"forceInput", PromptManager.INPUT_TYPES()["optional"][name][1]
|
|
)
|
|
self.assertEqual(focused["hidden"]["unique_id"], "UNIQUE_ID")
|
|
|
|
def test_focused_encodes_composition_and_tracks_main_text_and_graph(self):
|
|
node = PromptManagerFocused.__new__(PromptManagerFocused)
|
|
node.logger = logging.getLogger("test.focused")
|
|
node.comfyui_integration = mock.Mock()
|
|
node._inject_lora_trigger_words = lambda text: text
|
|
node._track_prompt_execution = mock.Mock(return_value=42)
|
|
clip = mock.Mock()
|
|
graph = {"1": {"class_type": "PromptManagerFocused", "inputs": {"text": "cat"}}}
|
|
conditioning, text = node.encode_prompt(
|
|
clip,
|
|
"cat",
|
|
prepend_text="trigger",
|
|
append_text="detailed",
|
|
prompt=graph,
|
|
unique_id="1",
|
|
)
|
|
self.assertEqual(text, "trigger cat detailed")
|
|
clip.tokenize.assert_called_once_with(text)
|
|
self.assertIs(conditioning, clip.encode_from_tokens_scheduled.return_value)
|
|
tracked = node._track_prompt_execution.call_args.kwargs
|
|
self.assertEqual(tracked["text"], "cat")
|
|
self.assertEqual(tracked["prompt_graph"], graph)
|
|
|
|
def test_focused_is_traced_as_positive_and_not_as_the_negative_run(self):
|
|
graph = {
|
|
"1": {"class_type": "PromptManagerFocused", "inputs": {"text": "cat"}},
|
|
"2": {"class_type": "PromptManagerFocused", "inputs": {"text": "blur"}},
|
|
"3": {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"model": ["4", 0],
|
|
"positive": ["1", 0],
|
|
"negative": ["2", 0],
|
|
},
|
|
},
|
|
}
|
|
self.assertEqual(positive_prompt_nodes(graph), ["1"])
|
|
self.assertEqual(run_prompt_nodes(graph), ["1"])
|