Vito Sansevero
729190309c
The on-prompt hook runs before ComfyUI validates the request, so counting there bumped run_count for prompts that were rejected or deleted from the queue. The hook now only resolves the positive PromptManager texts, mints a UUID4 prompt_id when the client sent none, and stashes the hashes under that id (QueuedRuns, bounded to 1000 entries). PromptQueue.get is wrapped idempotently; when a prompt is dequeued for execution its stashed runs are counted (or remembered as pending first use). Servers without a prompt_queue fall back to counting at submit time with one warning. __init__ defines _global_db before the monitoring block and passes a factory that raises RuntimeError while the database is unavailable; the package is now imported under test with server and folder_paths stubbed.
567 lines
22 KiB
Python
567 lines
22 KiB
Python
"""Run counting and node roles: only positive PromptManager nodes count and link.
|
|
|
|
Runs are resolved when a prompt is queued (ComfyUI's on-prompt hook, which runs
|
|
before validation) and counted when ComfyUI dequeues the prompt for execution,
|
|
because ComfyUI skips unchanged nodes, so node execution can't be relied on to
|
|
count, and a rejected or cancelled queue entry must not count either.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
import unittest.mock as mock
|
|
import uuid
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from database.operations import PromptDatabase
|
|
from prompt_manager_base import PromptManagerBase
|
|
from utils.hashing import generate_prompt_hash
|
|
from utils.prompt_tracker import PromptTracker
|
|
from utils import usage_tracking
|
|
from utils.usage_tracking import (
|
|
PendingFirstUse,
|
|
QueuedRuns,
|
|
_wrap_queue_get,
|
|
count_queued_prompt,
|
|
handle_queued_prompt,
|
|
on_dequeue,
|
|
register_queue_hook,
|
|
)
|
|
|
|
LOADER = {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {"ckpt_name": "m.safetensors"},
|
|
}
|
|
|
|
|
|
def workflow(positive_text="a cat", negative_text="blurry"):
|
|
"""API graph with a positive (#134) and a negative (#176) PromptManager node."""
|
|
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 FakePromptQueue:
|
|
"""Stand-in for execution.PromptQueue: get() pops preset results."""
|
|
|
|
def __init__(self):
|
|
self.results = []
|
|
self.get_calls = []
|
|
|
|
def get(self, timeout=None):
|
|
self.get_calls.append(timeout)
|
|
return self.results.pop(0) if self.results else None
|
|
|
|
|
|
class FakeServer:
|
|
"""Stand-in for PromptServer: on-prompt handlers plus an optional queue."""
|
|
|
|
def __init__(self, with_queue=True):
|
|
self.handlers = []
|
|
if with_queue:
|
|
self.prompt_queue = FakePromptQueue()
|
|
|
|
def add_on_prompt_handler(self, handler):
|
|
self.handlers.append(handler)
|
|
|
|
def on_prompt(self, json_data):
|
|
for handler in self.handlers:
|
|
json_data = handler(json_data)
|
|
return json_data
|
|
|
|
|
|
def queue_item(prompt_id, graph):
|
|
"""What PromptQueue.get returns for a prompt that passed validation."""
|
|
return ((0, prompt_id, graph, {}, [], False), 1)
|
|
|
|
|
|
class DbTestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
fd, self.path = tempfile.mkstemp(suffix=".db")
|
|
os.close(fd)
|
|
self.db = PromptDatabase(self.path)
|
|
self.pending = PendingFirstUse()
|
|
self.queued = QueuedRuns()
|
|
|
|
def tearDown(self):
|
|
close = getattr(self.db, "close", None)
|
|
if callable(close):
|
|
close()
|
|
for suffix in ("", "-wal", "-shm"):
|
|
if os.path.exists(self.path + suffix):
|
|
os.unlink(self.path + suffix)
|
|
|
|
def _save(self, text):
|
|
return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
|
|
|
|
def _runs(self, prompt_id):
|
|
return self.db.get_prompt_by_id(prompt_id)["run_count"]
|
|
|
|
def _queue(self, graph):
|
|
"""Queue a graph and let ComfyUI dequeue it for execution."""
|
|
request = handle_queued_prompt({"prompt": graph, "client_id": "x"}, self.queued)
|
|
on_dequeue(
|
|
request.get("prompt_id", ""), self.queued, lambda: self.db, self.pending
|
|
)
|
|
return request
|
|
|
|
|
|
class TestQueueHook(DbTestCase):
|
|
def test_counts_the_positive_prompt_only(self):
|
|
positive, negative = self._save("a cat"), self._save("blurry")
|
|
self._queue(workflow())
|
|
self.assertEqual(self._runs(positive), 1)
|
|
self.assertEqual(self._runs(negative), 0)
|
|
|
|
def test_requeue_counts_again_even_when_comfyui_skips_the_cached_node(self):
|
|
positive = self._save("a cat")
|
|
self._queue(workflow())
|
|
self._queue(workflow()) # no node execution in between
|
|
self.assertEqual(self._runs(positive), 2)
|
|
|
|
def test_new_prompt_is_counted_when_the_node_first_saves_it(self):
|
|
self._queue(workflow(positive_text="a brand new prompt"))
|
|
self.assertTrue(
|
|
self.pending.consume(generate_prompt_hash("a brand new prompt"))
|
|
)
|
|
self.assertFalse(
|
|
self.pending.consume(generate_prompt_hash("a brand new prompt"))
|
|
)
|
|
|
|
def test_same_text_in_two_positive_nodes_counts_once(self):
|
|
positive = self._save("a cat")
|
|
graph = workflow()
|
|
graph["135"] = {"class_type": "PromptManager", "inputs": {"text": "a cat"}}
|
|
graph["4"] = {
|
|
"class_type": "ConditioningCombine",
|
|
"inputs": {"conditioning_1": ["134", 0], "conditioning_2": ["135", 0]},
|
|
}
|
|
graph["5"]["inputs"]["positive"] = ["4", 0]
|
|
self._queue(graph)
|
|
self.assertEqual(self._runs(positive), 1)
|
|
|
|
def test_text_from_a_known_string_node_is_counted_at_queue_time(self):
|
|
positive = self._save("a cat")
|
|
graph = workflow()
|
|
graph["134"]["inputs"]["text"] = ["8", 0]
|
|
graph["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "a cat"}}
|
|
self._queue(graph)
|
|
self._queue(graph) # cached re-run: the node won't execute
|
|
self.assertEqual(self._runs(positive), 2)
|
|
|
|
def test_new_prompt_queued_twice_before_first_run_counts_twice(self):
|
|
self._queue(workflow(positive_text="queued twice"))
|
|
self._queue(workflow(positive_text="queued twice"))
|
|
self.assertEqual(self.pending.consume(generate_prompt_hash("queued twice")), 2)
|
|
|
|
def test_linked_text_is_left_to_node_execution(self):
|
|
positive = self._save("from a batch")
|
|
graph = workflow()
|
|
graph["134"]["inputs"]["text"] = ["7", 0]
|
|
graph["7"] = {"class_type": "PromptSearchList", "inputs": {}}
|
|
self._queue(graph)
|
|
self.assertEqual(self._runs(positive), 0)
|
|
|
|
def test_returns_the_request_and_never_raises(self):
|
|
request = {"prompt": workflow(), "client_id": "x", "extra_data": {}}
|
|
self.assertIs(handle_queued_prompt(request, self.queued), request)
|
|
for bad in (None, {}, {"prompt": "x"}, {"prompt": None}, []):
|
|
self.assertIs(handle_queued_prompt(bad, self.queued), bad)
|
|
broken_db = mock.Mock()
|
|
broken_db.get_prompt_by_hash.side_effect = RuntimeError("db locked")
|
|
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"):
|
|
on_dequeue(
|
|
request["prompt_id"], self.queued, lambda: broken_db, self.pending
|
|
)
|
|
|
|
def test_graph_without_prompt_manager_nodes_is_not_stashed_or_tagged(self):
|
|
request = {"prompt": {"20": LOADER}}
|
|
self.assertIs(handle_queued_prompt(request, self.queued), request)
|
|
self.assertNotIn("prompt_id", request)
|
|
self.assertEqual(len(self.queued), 0)
|
|
|
|
|
|
class TestCountAtDequeue(DbTestCase):
|
|
"""End to end through the registered hook and the wrapped queue."""
|
|
|
|
def setUp(self):
|
|
super().setUp()
|
|
self.server = FakeServer()
|
|
self._register(self.server)
|
|
|
|
def _register(self, server):
|
|
with mock.patch.object(usage_tracking, "_hook_registered", threading.Event()):
|
|
return register_queue_hook(
|
|
server, lambda: self.db, pending=self.pending, queued=self.queued
|
|
)
|
|
|
|
def _submit(self, graph, prompt_id=None):
|
|
request = {"prompt": graph, "client_id": "x"}
|
|
if prompt_id is not None:
|
|
request["prompt_id"] = prompt_id
|
|
return self.server.on_prompt(request)
|
|
|
|
def _dequeue(self, request):
|
|
self.server.prompt_queue.results.append(
|
|
queue_item(request["prompt_id"], request["prompt"])
|
|
)
|
|
return self.server.prompt_queue.get(timeout=1.0)
|
|
|
|
def test_hook_adds_a_lowercase_uuid4_prompt_id_when_absent(self):
|
|
request = self._submit(workflow())
|
|
prompt_id = request["prompt_id"]
|
|
self.assertEqual(str(uuid.UUID(prompt_id, version=4)), prompt_id)
|
|
self.assertEqual(prompt_id, prompt_id.lower())
|
|
|
|
def test_hook_keeps_an_existing_prompt_id(self):
|
|
request = self._submit(workflow(), prompt_id="client-chosen")
|
|
self.assertEqual(request["prompt_id"], "client-chosen")
|
|
|
|
def test_hook_alone_does_not_count(self):
|
|
positive = self._save("a cat")
|
|
self._submit(workflow())
|
|
self.assertEqual(self._runs(positive), 0)
|
|
|
|
def test_dequeue_counts_exactly_once(self):
|
|
positive = self._save("a cat")
|
|
request = self._submit(workflow())
|
|
result = self._dequeue(request)
|
|
self.assertEqual(result[0][1], request["prompt_id"])
|
|
self.assertEqual(self.server.prompt_queue.get_calls, [1.0])
|
|
self.assertEqual(self._runs(positive), 1)
|
|
self._dequeue(request) # a second dequeue of the same id has nothing left
|
|
self.assertEqual(self._runs(positive), 1)
|
|
|
|
def test_two_queued_prompts_are_each_counted_once(self):
|
|
positive = self._save("a cat")
|
|
first = self._submit(workflow())
|
|
second = self._submit(workflow())
|
|
self.assertNotEqual(first["prompt_id"], second["prompt_id"])
|
|
self._dequeue(first)
|
|
self.assertEqual(self._runs(positive), 1)
|
|
self._dequeue(second)
|
|
self.assertEqual(self._runs(positive), 2)
|
|
|
|
def test_rejected_prompt_is_never_counted(self):
|
|
positive = self._save("a cat")
|
|
self._submit(workflow()) # validation fails: never dequeued
|
|
other = self._submit(workflow())
|
|
self._dequeue(other)
|
|
self.assertEqual(self._runs(positive), 1)
|
|
|
|
def test_get_returning_none_is_a_noop(self):
|
|
positive = self._save("a cat")
|
|
self._submit(workflow())
|
|
self.assertIsNone(self.server.prompt_queue.get(timeout=0.1))
|
|
self.assertEqual(self._runs(positive), 0)
|
|
|
|
def test_unknown_text_is_pending_only_after_dequeue(self):
|
|
prompt_hash = generate_prompt_hash("never saved")
|
|
request = self._submit(workflow(positive_text="never saved"))
|
|
self.assertEqual(self.pending.consume(prompt_hash), 0)
|
|
self._dequeue(request)
|
|
self.assertEqual(self.pending.consume(prompt_hash), 1)
|
|
|
|
def test_shorter_and_longer_queue_items_are_tolerated(self):
|
|
positive = self._save("a cat")
|
|
request = self._submit(workflow())
|
|
self.server.prompt_queue.results.append(
|
|
((0, request["prompt_id"], request["prompt"], {}, []), 1)
|
|
)
|
|
self.server.prompt_queue.get()
|
|
self.assertEqual(self._runs(positive), 1)
|
|
request = self._submit(workflow())
|
|
self.server.prompt_queue.results.append(
|
|
((0, request["prompt_id"], request["prompt"], {}, [], False, "x"), 2)
|
|
)
|
|
self.server.prompt_queue.get()
|
|
self.assertEqual(self._runs(positive), 2)
|
|
|
|
def test_malformed_queue_item_is_logged_and_returned(self):
|
|
item = (("not", "a", "tuple"),)
|
|
self.server.prompt_queue.results.append(item)
|
|
self.server.prompt_queue.results.append("garbage")
|
|
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"):
|
|
self.assertEqual(self.server.prompt_queue.get(), item)
|
|
self.assertEqual(self.server.prompt_queue.get(), "garbage")
|
|
|
|
def test_db_factory_raising_at_dequeue_is_logged_not_raised(self):
|
|
def broken():
|
|
raise RuntimeError("PromptManager database unavailable")
|
|
|
|
queued = QueuedRuns()
|
|
server = FakeServer()
|
|
with mock.patch.object(usage_tracking, "_hook_registered", threading.Event()):
|
|
register_queue_hook(server, broken, pending=self.pending, queued=queued)
|
|
request = server.on_prompt({"prompt": workflow()})
|
|
server.prompt_queue.results.append(queue_item(request["prompt_id"], {}))
|
|
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING") as logs:
|
|
self.assertIsNotNone(server.prompt_queue.get())
|
|
self.assertIn("database unavailable", "\n".join(logs.output))
|
|
|
|
def test_server_without_prompt_queue_counts_at_queue_time_and_warns_once(self):
|
|
positive = self._save("a cat")
|
|
server = FakeServer(with_queue=False)
|
|
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING") as logs:
|
|
self.assertTrue(self._register(server))
|
|
warnings = [line for line in logs.output if line.startswith("WARNING")]
|
|
self.assertEqual(len(warnings), 1)
|
|
server.on_prompt({"prompt": workflow()})
|
|
self.assertEqual(self._runs(positive), 1)
|
|
server.on_prompt({"prompt": workflow(positive_text="new text")})
|
|
self.assertEqual(self.pending.consume(generate_prompt_hash("new text")), 1)
|
|
|
|
def test_server_whose_queue_has_no_get_uses_legacy_counting(self):
|
|
positive = self._save("a cat")
|
|
server = FakeServer(with_queue=False)
|
|
server.prompt_queue = object()
|
|
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"):
|
|
self._register(server)
|
|
server.on_prompt({"prompt": workflow()})
|
|
self.assertEqual(self._runs(positive), 1)
|
|
|
|
def test_legacy_counting_calls_db_factory_inside_the_guard(self):
|
|
def broken():
|
|
raise RuntimeError("PromptManager database unavailable")
|
|
|
|
request = {"prompt": workflow()}
|
|
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"):
|
|
self.assertIs(count_queued_prompt(request, broken, self.pending), request)
|
|
|
|
def test_wrapping_twice_does_not_double_count(self):
|
|
positive = self._save("a cat")
|
|
prompt_queue = self.server.prompt_queue
|
|
wrapped_get = prompt_queue.get
|
|
calls = []
|
|
self.assertFalse(_wrap_queue_get(prompt_queue, calls.append))
|
|
self.assertIs(prompt_queue.get, wrapped_get)
|
|
request = self._submit(workflow())
|
|
self._dequeue(request)
|
|
self.assertEqual(self._runs(positive), 1)
|
|
self.assertEqual(calls, [])
|
|
|
|
def test_wrap_rejects_queues_without_a_callable_get(self):
|
|
self.assertFalse(_wrap_queue_get(object(), lambda prompt_id: None))
|
|
self.assertFalse(_wrap_queue_get(None, lambda prompt_id: None))
|
|
|
|
def test_wrapper_reports_the_dequeued_prompt_id(self):
|
|
prompt_queue = FakePromptQueue()
|
|
seen = []
|
|
self.assertTrue(_wrap_queue_get(prompt_queue, seen.append))
|
|
prompt_queue.results.append(queue_item("abc", {}))
|
|
self.assertEqual(prompt_queue.get(timeout=2.5), queue_item("abc", {}))
|
|
self.assertEqual(prompt_queue.get_calls, [2.5])
|
|
self.assertEqual(seen, ["abc"])
|
|
|
|
def test_dequeue_callback_errors_are_logged_not_raised(self):
|
|
prompt_queue = FakePromptQueue()
|
|
|
|
def explode(prompt_id):
|
|
raise ValueError("boom")
|
|
|
|
_wrap_queue_get(prompt_queue, explode)
|
|
prompt_queue.results.append(queue_item("abc", {}))
|
|
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"):
|
|
self.assertIsNotNone(prompt_queue.get())
|
|
|
|
|
|
class TestQueuedRuns(unittest.TestCase):
|
|
def test_pop_returns_runs_once(self):
|
|
queued = QueuedRuns()
|
|
queued.put("p1", [("h1", "text one")])
|
|
self.assertEqual(queued.pop("p1"), [("h1", "text one")])
|
|
self.assertEqual(queued.pop("p1"), [])
|
|
self.assertEqual(queued.pop("missing"), [])
|
|
|
|
def test_bound_evicts_oldest(self):
|
|
queued = QueuedRuns(max_entries=3)
|
|
for i in range(4):
|
|
queued.put(f"p{i}", [(f"h{i}", f"text {i}")])
|
|
self.assertEqual(len(queued), 3)
|
|
self.assertEqual(queued.pop("p0"), [])
|
|
self.assertEqual(queued.pop("p3"), [("h3", "text 3")])
|
|
|
|
def test_default_bound_is_one_thousand(self):
|
|
queued = QueuedRuns()
|
|
for i in range(1001):
|
|
queued.put(str(i), [("h", "t")])
|
|
self.assertEqual(len(queued), 1000)
|
|
self.assertEqual(queued.pop("0"), [])
|
|
self.assertEqual(queued.pop("1000"), [("h", "t")])
|
|
|
|
def test_requeue_of_the_same_id_replaces_the_runs(self):
|
|
queued = QueuedRuns()
|
|
queued.put("p1", [("h1", "one")])
|
|
queued.put("p1", [("h2", "two")])
|
|
self.assertEqual(queued.pop("p1"), [("h2", "two")])
|
|
|
|
def test_empty_runs_are_not_stored(self):
|
|
queued = QueuedRuns()
|
|
queued.put("p1", [])
|
|
self.assertEqual(len(queued), 0)
|
|
|
|
|
|
class TestPendingFirstUse(unittest.TestCase):
|
|
def test_entries_expire(self):
|
|
pending = PendingFirstUse(ttl_seconds=10)
|
|
with mock.patch("utils.usage_tracking.time.monotonic", return_value=100.0):
|
|
pending.add("h")
|
|
with mock.patch("utils.usage_tracking.time.monotonic", return_value=111.0):
|
|
self.assertFalse(pending.consume("h"))
|
|
|
|
def test_consume_returns_how_many_runs_were_pending(self):
|
|
pending = PendingFirstUse()
|
|
pending.add("h")
|
|
pending.add("h")
|
|
pending.add("h")
|
|
self.assertEqual(pending.consume("h"), 3)
|
|
self.assertEqual(pending.consume("h"), 0)
|
|
|
|
def test_size_is_capped(self):
|
|
pending = PendingFirstUse(max_entries=3)
|
|
for h in ("a", "b", "c", "d"):
|
|
pending.add(h)
|
|
self.assertFalse(pending.consume("a")) # oldest evicted
|
|
self.assertTrue(pending.consume("d"))
|
|
|
|
|
|
class FakeTracker:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
def set_current_prompt(self, prompt_text, additional_data=None, push_to_queue=True):
|
|
self.calls.append({"text": prompt_text, "push_to_queue": push_to_queue})
|
|
return "exec-1"
|
|
|
|
|
|
class TestNodeRoles(DbTestCase):
|
|
def _node(self):
|
|
node = PromptManagerBase.__new__(PromptManagerBase)
|
|
node.db = self.db
|
|
node.prompt_tracker = FakeTracker()
|
|
node.logger = __import__("logging").getLogger("test.usage_tracking")
|
|
node.pending_first_use = self.pending
|
|
return node
|
|
|
|
def _run(self, node, graph, unique_id, text):
|
|
return node._track_prompt_execution(
|
|
text=text,
|
|
encoding_text=text,
|
|
category=None,
|
|
tags=None,
|
|
additional_data={},
|
|
prompt_graph=graph,
|
|
unique_id=unique_id,
|
|
)
|
|
|
|
def test_negative_node_saves_but_never_becomes_the_current_prompt(self):
|
|
node = self._node()
|
|
prompt_id = self._run(node, workflow(), "176", "blurry")
|
|
self.assertIsNotNone(self.db.get_prompt_by_id(prompt_id))
|
|
self.assertEqual(node.prompt_tracker.calls, [])
|
|
self.assertEqual(self._runs(prompt_id), 0)
|
|
|
|
def test_positive_typed_text_is_current_without_queue_push(self):
|
|
node = self._node()
|
|
self._run(node, workflow(), "134", "a cat")
|
|
self.assertEqual(
|
|
node.prompt_tracker.calls, [{"text": "a cat", "push_to_queue": False}]
|
|
)
|
|
|
|
def test_first_save_after_queue_hook_counts_one_run(self):
|
|
self._queue(workflow(positive_text="never seen before"))
|
|
prompt_id = self._run(
|
|
self._node(),
|
|
workflow(positive_text="never seen before"),
|
|
"134",
|
|
"never seen before",
|
|
)
|
|
self.assertEqual(self._runs(prompt_id), 1)
|
|
|
|
def test_new_prompt_queued_twice_counts_both_runs_on_first_save(self):
|
|
graph = workflow(positive_text="double queued")
|
|
self._queue(graph)
|
|
self._queue(graph)
|
|
prompt_id = self._run(self._node(), graph, "134", "double queued")
|
|
self.assertEqual(self._runs(prompt_id), 2)
|
|
|
|
def test_text_resolved_at_queue_time_is_not_counted_again_or_queued(self):
|
|
graph = workflow()
|
|
graph["134"]["inputs"]["text"] = ["8", 0]
|
|
graph["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "a cat"}}
|
|
positive = self._save("a cat")
|
|
self._queue(graph)
|
|
node = self._node()
|
|
self._run(node, graph, "134", "a cat")
|
|
self.assertEqual(self._runs(positive), 1)
|
|
self.assertEqual(
|
|
node.prompt_tracker.calls, [{"text": "a cat", "push_to_queue": False}]
|
|
)
|
|
|
|
def test_resolution_mismatch_falls_back_to_counting_and_queueing(self):
|
|
# Hook guessed "a cat" but the node received different text: trust the node
|
|
graph = workflow()
|
|
graph["134"]["inputs"]["text"] = ["8", 0]
|
|
graph["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "a cat"}}
|
|
node = self._node()
|
|
prompt_id = self._run(node, graph, "134", "actually a dog")
|
|
self.assertEqual(self._runs(prompt_id), 1)
|
|
self.assertEqual(
|
|
node.prompt_tracker.calls,
|
|
[{"text": "actually a dog", "push_to_queue": True}],
|
|
)
|
|
|
|
def test_existing_prompt_is_not_counted_twice_by_hook_and_node(self):
|
|
positive = self._save("a cat")
|
|
self._queue(workflow())
|
|
self._run(self._node(), workflow(), "134", "a cat")
|
|
self.assertEqual(self._runs(positive), 1)
|
|
|
|
def test_positive_linked_text_pushes_queue_and_counts_each_item(self):
|
|
graph = workflow()
|
|
graph["134"]["inputs"]["text"] = ["7", 0]
|
|
graph["7"] = {"class_type": "PromptSearchList", "inputs": {}}
|
|
node = self._node()
|
|
first = self._run(node, graph, "134", "batch item one")
|
|
self._run(node, graph, "134", "batch item one")
|
|
self.assertEqual(self._runs(first), 2)
|
|
self.assertTrue(all(c["push_to_queue"] for c in node.prompt_tracker.calls))
|
|
|
|
def test_without_graph_context_keeps_legacy_behaviour(self):
|
|
node = self._node()
|
|
self._run(node, None, None, "standalone")
|
|
self.assertEqual(
|
|
node.prompt_tracker.calls, [{"text": "standalone", "push_to_queue": True}]
|
|
)
|
|
|
|
|
|
class TestTrackerQueueFlag(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tracker = PromptTracker(mock.Mock())
|
|
|
|
def test_push_to_queue_false_keeps_the_batch_queue_empty(self):
|
|
self.tracker.set_current_prompt("a cat", {"prompt_id": 1}, push_to_queue=False)
|
|
self.assertIsNone(self.tracker.pop_next_prompt())
|
|
self.assertEqual(self.tracker.get_current_prompt()["id"], 1)
|
|
|
|
def test_default_still_pushes(self):
|
|
self.tracker.set_current_prompt("a cat", {"prompt_id": 1})
|
|
self.assertEqual(self.tracker.pop_next_prompt()["id"], 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|