Vito
4c9cb7f5ea
* refactor: remove dead and test-only code and the write-only prompt registry (#620) Removes helpers nothing calls, test-only wrappers and their tests, the unused config writer, three unused images, and the prompt registry the nodes wrote to but nothing read, which grew with every run. * fix: network-share readers, model unload off the loop, bad save requests (#620) Every read-only SQLite open now uses one URI helper, so a library on a Windows network share works in backups, restores, moves and history reads. Unloading a model no longer blocks the server. A malformed /save body gets a 400, not a 500. * feat(library): add Recently used, Most used and Name Z-A sorting (#620) The old Dashboard's sorts the community asked for, on the Prompt library, in saved views, and with paging through the whole list. * feat(workspace): stop linking the old pages (#620) They show a removal notice since 3.5.1. System no longer links them, the load failure offers a reload and ComfyUI instead, and the README says they are going away. The pages still work at their URLs.
83 lines
3.2 KiB
Python
83 lines
3.2 KiB
Python
"""Nodes pass ComfyUI's queued graph to role-aware tracking; the hook registers once."""
|
|
|
|
import os
|
|
import sys
|
|
import threading
|
|
import unittest
|
|
import unittest.mock as mock
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from prompt_manager import PromptManager
|
|
from prompt_manager_text import PromptManagerText
|
|
from utils import usage_tracking
|
|
|
|
GRAPH = {"134": {"class_type": "PromptManager", "inputs": {"text": "a cat"}}}
|
|
|
|
|
|
def bare(node_cls):
|
|
"""A node without its database, tracker or gallery side effects."""
|
|
node = node_cls.__new__(node_cls)
|
|
node.logger = __import__("logging").getLogger("test.wiring")
|
|
node._inject_lora_trigger_words = lambda text: text
|
|
node._track_prompt_execution = mock.Mock(return_value=42)
|
|
return node
|
|
|
|
|
|
class TestHiddenInputs(unittest.TestCase):
|
|
def test_both_nodes_request_the_queued_graph_and_their_id(self):
|
|
for node_cls in (PromptManager, PromptManagerText):
|
|
hidden = node_cls.INPUT_TYPES().get("hidden", {})
|
|
self.assertEqual(hidden.get("prompt"), "PROMPT", node_cls.__name__)
|
|
self.assertEqual(hidden.get("unique_id"), "UNIQUE_ID", node_cls.__name__)
|
|
|
|
|
|
class TestNodesForwardGraph(unittest.TestCase):
|
|
def test_prompt_manager_forwards_graph_and_texts(self):
|
|
node = bare(PromptManager)
|
|
clip = mock.Mock()
|
|
node.encode_prompt(
|
|
clip, "a cat", prepend_text="best,", prompt=GRAPH, unique_id="134"
|
|
)
|
|
|
|
kwargs = node._track_prompt_execution.call_args.kwargs
|
|
self.assertEqual(kwargs["prompt_graph"], GRAPH)
|
|
self.assertEqual(kwargs["unique_id"], "134")
|
|
self.assertEqual(kwargs["text"], "a cat")
|
|
self.assertEqual(kwargs["encoding_text"], "best, a cat")
|
|
self.assertEqual(kwargs["additional_data"]["final_text"], "best, a cat")
|
|
|
|
def test_prompt_manager_text_forwards_graph(self):
|
|
node = bare(PromptManagerText)
|
|
node.process_text("a castle", prompt=GRAPH, unique_id="7")
|
|
|
|
kwargs = node._track_prompt_execution.call_args.kwargs
|
|
self.assertEqual((kwargs["prompt_graph"], kwargs["unique_id"]), (GRAPH, "7"))
|
|
self.assertEqual(kwargs["text"], "a castle")
|
|
|
|
def test_nodes_still_work_without_graph_context(self):
|
|
node = bare(PromptManager)
|
|
node.encode_prompt(mock.Mock(), "a cat")
|
|
kwargs = node._track_prompt_execution.call_args.kwargs
|
|
self.assertIsNone(kwargs["prompt_graph"])
|
|
self.assertIsNone(kwargs["unique_id"])
|
|
|
|
|
|
class TestQueueHookRegistration(unittest.TestCase):
|
|
def test_registers_once_and_handler_passes_requests_through(self):
|
|
server = mock.Mock()
|
|
db = mock.Mock()
|
|
db.get_prompt_by_hash.return_value = None
|
|
with mock.patch.object(usage_tracking, "_hook_registered", threading.Event()):
|
|
self.assertTrue(usage_tracking.register_queue_hook(server, lambda: db))
|
|
self.assertFalse(usage_tracking.register_queue_hook(server, lambda: db))
|
|
|
|
server.add_on_prompt_handler.assert_called_once()
|
|
handler = server.add_on_prompt_handler.call_args.args[0]
|
|
request = {"prompt": GRAPH}
|
|
self.assertIs(handler(request), request)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|