test: cover prompt_manager_base, prompt_manager and prompt_search_list
Node base: init wiring, metadata update on re-save, run counting, database error paths, tag parsing, search and recent API wrappers, gallery start failure, LoRA trigger-word injection (disabled, missing modules, enabled). PromptManager: missing CLIP error, encoding of the combined text with a fake CLIP, tracking failure tolerance, empty text. PromptManagerText: tracking failure tolerance. PromptSearchList: filters, multipart and LoRA-only skipping, newline collapsing, preview, empty and error results.
This commit is contained in:
Vito Sansevero
@@ -138,5 +138,291 @@ class TestIsChangedPromptSearchList(unittest.TestCase):
|
||||
self.assertIsInstance(result, float)
|
||||
|
||||
|
||||
import logging # noqa: E402
|
||||
import tempfile # noqa: E402
|
||||
import types # noqa: E402
|
||||
import unittest.mock as mock # noqa: E402
|
||||
|
||||
from database.operations import PromptDatabase # noqa: E402
|
||||
from prompt_manager_base import PromptManagerBase # noqa: E402
|
||||
import prompt_manager_base # noqa: E402
|
||||
|
||||
|
||||
class FakeClip:
|
||||
"""The two CLIP calls the node makes, recorded."""
|
||||
|
||||
def __init__(self):
|
||||
self.tokenized = []
|
||||
|
||||
def tokenize(self, text):
|
||||
self.tokenized.append(text)
|
||||
return ("tokens", text)
|
||||
|
||||
def encode_from_tokens_scheduled(self, tokens):
|
||||
return [("cond", tokens[1])]
|
||||
|
||||
|
||||
class NodeTestCase(unittest.TestCase):
|
||||
"""A node with a real temporary database and no monitor or tracker."""
|
||||
|
||||
node_cls = PromptManagerBase
|
||||
|
||||
def setUp(self):
|
||||
fd, self.db_path = tempfile.mkstemp(suffix=".db")
|
||||
os.close(fd)
|
||||
self.db = PromptDatabase(self.db_path)
|
||||
self.node = self.node_cls.__new__(self.node_cls)
|
||||
self.node.db = self.db
|
||||
self.node.logger = logging.getLogger("test.node")
|
||||
self.node.prompt_tracker = mock.Mock()
|
||||
self.node.image_monitor = mock.Mock()
|
||||
self.node.comfyui_integration = mock.Mock()
|
||||
|
||||
def tearDown(self):
|
||||
for suffix in ("", "-wal", "-shm"):
|
||||
if os.path.exists(self.db_path + suffix):
|
||||
os.unlink(self.db_path + suffix)
|
||||
|
||||
|
||||
class TestBaseSaveAndSearch(NodeTestCase):
|
||||
def test_init_wires_database_tracker_monitor_and_integration(self):
|
||||
with (
|
||||
mock.patch.object(prompt_manager_base, "PromptDatabase") as db_cls,
|
||||
mock.patch.object(prompt_manager_base, "get_prompt_tracker") as tracker,
|
||||
mock.patch.object(prompt_manager_base, "get_image_monitor") as monitor,
|
||||
mock.patch.object(
|
||||
prompt_manager_base, "get_comfyui_integration"
|
||||
) as integration,
|
||||
):
|
||||
node = PromptManagerBase()
|
||||
self.assertIs(node.db, db_cls.return_value)
|
||||
tracker.assert_called_once_with(db_cls.return_value)
|
||||
monitor.assert_called_once_with(db_cls.return_value, tracker.return_value)
|
||||
self.assertIs(node.comfyui_integration, integration.return_value)
|
||||
monitor.return_value.start_monitoring.assert_called_once_with()
|
||||
|
||||
def test_saving_an_existing_prompt_updates_its_metadata(self):
|
||||
first = self.node._save_prompt_to_database("a cat")
|
||||
second = self.node._save_prompt_to_database(
|
||||
"a cat", category="animals", tags=["cute"]
|
||||
)
|
||||
self.assertEqual(first, second)
|
||||
saved = self.db.get_prompt_by_id(first)
|
||||
self.assertEqual(saved["category"], "animals")
|
||||
self.assertIn("cute", saved["tags"])
|
||||
self.assertEqual(saved["run_count"], 0)
|
||||
|
||||
def test_count_run_on_an_existing_prompt(self):
|
||||
prompt_id = self.node._save_prompt_to_database("a cat")
|
||||
self.node._save_prompt_to_database("a cat", count_run=True)
|
||||
self.assertEqual(self.db.get_prompt_by_id(prompt_id)["run_count"], 1)
|
||||
|
||||
def test_database_errors_are_logged_and_return_none(self):
|
||||
self.node.db = mock.Mock()
|
||||
self.node.db.get_prompt_by_hash.side_effect = RuntimeError("locked")
|
||||
with self.assertLogs("test.node", level="ERROR"):
|
||||
self.assertIsNone(self.node._save_prompt_to_database("a cat"))
|
||||
|
||||
def test_record_use_failure_is_a_warning(self):
|
||||
self.node.db = mock.Mock()
|
||||
self.node.db.record_prompt_use.side_effect = RuntimeError("locked")
|
||||
with self.assertLogs("test.node", level="WARNING"):
|
||||
self.node._record_use(1)
|
||||
|
||||
def test_parse_tags(self):
|
||||
self.assertIsNone(self.node._parse_tags(""))
|
||||
self.assertIsNone(self.node._parse_tags(" , ,"))
|
||||
self.assertEqual(self.node._parse_tags(" a, b ,,c"), ["a", "b", "c"])
|
||||
|
||||
def test_search_prompts_and_api_wrappers(self):
|
||||
self.node._save_prompt_to_database("a cat on a mat")
|
||||
self.assertEqual(self.node._search_prompts(""), [])
|
||||
self.assertEqual(self.node._search_prompts(" "), [])
|
||||
self.assertEqual(len(self.node.search_prompts_api("cat")), 1)
|
||||
self.assertEqual(len(self.node.get_recent_prompts_api(limit=5)["prompts"]), 1)
|
||||
|
||||
def test_search_and_recent_errors_return_empty_lists(self):
|
||||
self.node.db = mock.Mock()
|
||||
self.node.db.search_prompts.side_effect = RuntimeError("locked")
|
||||
self.node.db.get_recent_prompts.side_effect = RuntimeError("locked")
|
||||
with self.assertLogs("test.node", level="ERROR"):
|
||||
self.assertEqual(self.node._search_prompts("cat"), [])
|
||||
with self.assertLogs("test.node", level="ERROR"):
|
||||
self.assertEqual(self.node.get_recent_prompts_api(), [])
|
||||
|
||||
def test_gallery_system_start_failure_is_logged(self):
|
||||
self.node.image_monitor.start_monitoring.side_effect = OSError("no inotify")
|
||||
with self.assertLogs("test.node", level="ERROR"):
|
||||
self.node._start_gallery_system()
|
||||
self.node.image_monitor.get_status.return_value = {"running": False}
|
||||
self.node.prompt_tracker.get_status.return_value = {"active_prompts_count": 0}
|
||||
self.assertEqual(
|
||||
self.node.get_gallery_status(),
|
||||
{
|
||||
"image_monitor": {"running": False},
|
||||
"prompt_tracker": {"active_prompts_count": 0},
|
||||
},
|
||||
)
|
||||
self.node.cleanup_gallery_system()
|
||||
self.node.image_monitor.stop_monitoring.assert_not_called()
|
||||
|
||||
|
||||
def fake_integration_config(enabled, trigger_words=True, path="/loras"):
|
||||
module = types.ModuleType("py.config")
|
||||
module.IntegrationConfig = types.SimpleNamespace(
|
||||
LORA_MANAGER_ENABLED=enabled,
|
||||
LORA_TRIGGER_WORDS_ENABLED=trigger_words,
|
||||
LORA_MANAGER_PATH=path,
|
||||
)
|
||||
return module
|
||||
|
||||
|
||||
class TestLoraTriggerWords(NodeTestCase):
|
||||
def test_disabled_integration_returns_text_unchanged(self):
|
||||
with mock.patch.dict(
|
||||
sys.modules, {"py.config": fake_integration_config(False)}
|
||||
):
|
||||
self.assertEqual(self.node._inject_lora_trigger_words("a cat"), "a cat")
|
||||
with mock.patch.dict(
|
||||
sys.modules, {"py.config": fake_integration_config(True, False)}
|
||||
):
|
||||
self.assertEqual(self.node._inject_lora_trigger_words("a cat"), "a cat")
|
||||
|
||||
def test_missing_config_or_lora_utils_returns_text_unchanged(self):
|
||||
with mock.patch.dict(sys.modules, {"py.config": None}):
|
||||
self.assertEqual(self.node._inject_lora_trigger_words("a cat"), "a cat")
|
||||
with mock.patch.dict(
|
||||
sys.modules,
|
||||
{"py.config": fake_integration_config(True), "py.lora_utils": None},
|
||||
):
|
||||
self.assertEqual(self.node._inject_lora_trigger_words("a cat"), "a cat")
|
||||
|
||||
def test_enabled_integration_loads_the_cache_and_injects(self):
|
||||
cache = mock.Mock(is_loaded=False)
|
||||
lora_utils = types.ModuleType("py.lora_utils")
|
||||
lora_utils.get_trigger_cache = lambda: cache
|
||||
lora_utils.inject_trigger_words = lambda text, c: (text + " trig", ["trig"])
|
||||
with mock.patch.dict(
|
||||
sys.modules,
|
||||
{"py.config": fake_integration_config(True), "py.lora_utils": lora_utils},
|
||||
):
|
||||
self.assertEqual(
|
||||
self.node._inject_lora_trigger_words("a cat"), "a cat trig"
|
||||
)
|
||||
cache.load.assert_called_once_with("/loras")
|
||||
|
||||
|
||||
class TestPromptManagerEncoding(NodeTestCase):
|
||||
node_cls = PromptManager
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.node._inject_lora_trigger_words = lambda text: text
|
||||
|
||||
def test_missing_clip_is_a_clear_error(self):
|
||||
with self.assertRaisesRegex(RuntimeError, "clip input is invalid"):
|
||||
self.node.encode_prompt(None, "a cat")
|
||||
|
||||
def test_encodes_combined_text_and_saves_the_main_text(self):
|
||||
clip = FakeClip()
|
||||
conditioning, final = self.node.encode_prompt(
|
||||
clip, "a cat", prepend_text=" best ", append_text="8k ", tags="x"
|
||||
)
|
||||
self.assertEqual(final, "best a cat 8k")
|
||||
self.assertEqual(clip.tokenized, ["best a cat 8k"])
|
||||
self.assertEqual(conditioning, [("cond", "best a cat 8k")])
|
||||
saved = self.db.get_recent_prompts(limit=1)["prompts"][0]
|
||||
self.assertEqual(saved["text"], "a cat")
|
||||
self.assertIn("prepend:best", saved["tags"])
|
||||
self.assertIn("append:8k", saved["tags"])
|
||||
self.node.comfyui_integration.register_prompt.assert_called_once()
|
||||
|
||||
def test_tracking_failure_does_not_break_encoding(self):
|
||||
self.node._track_prompt_execution = mock.Mock(side_effect=RuntimeError("db"))
|
||||
with self.assertLogs("test.node", level="WARNING"):
|
||||
_, final = self.node.encode_prompt(FakeClip(), "a cat")
|
||||
self.assertEqual(final, "a cat")
|
||||
|
||||
def test_empty_text_is_encoded_but_not_saved(self):
|
||||
_, final = self.node.encode_prompt(FakeClip(), " ")
|
||||
self.assertEqual(final.strip(), "")
|
||||
self.assertEqual(self.db.get_recent_prompts(limit=1)["prompts"], [])
|
||||
|
||||
|
||||
class TestPromptManagerTextProcessing(NodeTestCase):
|
||||
node_cls = PromptManagerText
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.node._inject_lora_trigger_words = lambda text: text
|
||||
|
||||
def test_tracking_failure_does_not_break_processing(self):
|
||||
self.node._track_prompt_execution = mock.Mock(side_effect=RuntimeError("db"))
|
||||
with self.assertLogs("test.node", level="WARNING"):
|
||||
result = self.node.process_text("a castle", prepend_text="big")
|
||||
self.assertEqual(result[0], "big a castle")
|
||||
|
||||
|
||||
class TestPromptSearchListSearch(NodeTestCase):
|
||||
node_cls = PromptSearchList
|
||||
|
||||
def _save(self, text, **kwargs):
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
return self.db.save_prompt(
|
||||
text=text, prompt_hash=generate_prompt_hash(text), **kwargs
|
||||
)
|
||||
|
||||
def test_init_opens_the_database(self):
|
||||
with mock.patch("prompt_search_list.PromptDatabase") as db_cls:
|
||||
node = PromptSearchList()
|
||||
self.assertIs(node.db, db_cls.return_value)
|
||||
|
||||
def test_no_results_returns_a_single_empty_item(self):
|
||||
result = self.node.search(text="nothing here")
|
||||
self.assertEqual(result["result"], ([""], "No results found", 0))
|
||||
self.assertEqual(result["ui"], {"text": ["No results found"]})
|
||||
|
||||
def test_filters_by_text_category_tags_and_rating(self):
|
||||
self._save("a red cat", category="animals", tags=["cute"])
|
||||
self._save("a blue car", category="vehicles", tags=["fast"])
|
||||
self.assertEqual(self.node.search(text="cat")["result"][0], ["a red cat"])
|
||||
self.assertEqual(
|
||||
self.node.search(category="vehicles")["result"][0], ["a blue car"]
|
||||
)
|
||||
self.assertEqual(self.node.search(tags=" cute , ")["result"][0], ["a red cat"])
|
||||
self.assertEqual(self.node.search(min_rating=5)["result"][2], 0)
|
||||
self.assertEqual(self.node.search(limit=1)["result"][2], 1)
|
||||
|
||||
def test_multipart_and_lora_only_prompts_are_skipped(self):
|
||||
self._save("Clip_1: a cat Clip_2: a dog")
|
||||
self._save("<lora:style:0.8> <lora:other:1>")
|
||||
self._save("plain prompt")
|
||||
self.assertEqual(self.node.search()["result"][0], ["plain prompt"])
|
||||
result = self.node.search(skip_multipart=False)["result"]
|
||||
self.assertEqual(
|
||||
sorted(result[0]), ["Clip_1: a cat Clip_2: a dog", "plain prompt"]
|
||||
)
|
||||
|
||||
def test_newlines_collapse_and_preview_is_numbered_and_truncated(self):
|
||||
self._save("line one\n line two")
|
||||
self._save("x" * 200)
|
||||
result = self.node.search()
|
||||
prompts, preview, count = result["result"]
|
||||
self.assertEqual(count, 2)
|
||||
self.assertIn("line one line two", prompts)
|
||||
lines = preview.splitlines()
|
||||
self.assertTrue(lines[0].startswith("[1] "))
|
||||
self.assertTrue(any(line.endswith("...") and len(line) < 140 for line in lines))
|
||||
self.assertEqual(result["ui"], {"text": ["Found 2 prompts"]})
|
||||
|
||||
def test_database_errors_are_returned_not_raised(self):
|
||||
self.node.db = mock.Mock()
|
||||
self.node.db.search_prompts.side_effect = RuntimeError("locked")
|
||||
result = self.node.search(text="cat")
|
||||
self.assertEqual(result["result"], ([""], "Error: locked", 0))
|
||||
self.assertEqual(result["ui"], {"text": ["Search error: locked"]})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user