154 lines
6.0 KiB
Python
154 lines
6.0 KiB
Python
"""Tests for utils.comfyui_integration."""
|
|
|
|
import os
|
|
import sys
|
|
import types
|
|
import unittest
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from utils.comfyui_integration import ComfyUIMetadataIntegration
|
|
|
|
|
|
class TestSaveImageNotPatched(unittest.TestCase):
|
|
"""Regression for #75: SaveImage must save each node's own prompt inputs.
|
|
|
|
The old SaveImage patch wrote the last-run PromptManager prompt into every
|
|
PromptManager node, so negative prompts were saved as copies of the positive.
|
|
"""
|
|
|
|
def setUp(self):
|
|
def save_images(
|
|
self_node,
|
|
images,
|
|
filename_prefix="ComfyUI",
|
|
prompt=None,
|
|
extra_pnginfo=None,
|
|
):
|
|
return {"prompt": prompt}
|
|
|
|
self.original_save_images = save_images
|
|
fake_nodes = types.ModuleType("nodes")
|
|
fake_nodes.SaveImage = type("SaveImage", (), {"save_images": save_images})
|
|
self._saved_nodes = sys.modules.get("nodes")
|
|
sys.modules["nodes"] = fake_nodes
|
|
self.fake_nodes = fake_nodes
|
|
ComfyUIMetadataIntegration._instance = None
|
|
|
|
def tearDown(self):
|
|
ComfyUIMetadataIntegration._instance = None
|
|
if self._saved_nodes is None:
|
|
sys.modules.pop("nodes", None)
|
|
else:
|
|
sys.modules["nodes"] = self._saved_nodes
|
|
|
|
def test_save_images_is_left_untouched(self):
|
|
ComfyUIMetadataIntegration()
|
|
self.assertIs(self.fake_nodes.SaveImage.save_images, self.original_save_images)
|
|
|
|
def test_registered_prompt_does_not_rewrite_saved_metadata(self):
|
|
integration = ComfyUIMetadataIntegration()
|
|
integration.register_prompt("pm_1", "positive final text", {})
|
|
prompt = {
|
|
"39": {"class_type": "PromptManager", "inputs": {"text": "a cat"}},
|
|
"40": {"class_type": "PromptManager", "inputs": {"text": "blurry"}},
|
|
}
|
|
|
|
result = self.fake_nodes.SaveImage().save_images([], prompt=prompt)
|
|
|
|
self.assertEqual(result["prompt"]["39"]["inputs"]["text"], "a cat")
|
|
self.assertEqual(result["prompt"]["40"]["inputs"]["text"], "blurry")
|
|
|
|
def test_prompt_registry_still_available_to_nodes(self):
|
|
integration = ComfyUIMetadataIntegration()
|
|
integration.register_prompt("pm_1", "hello", {})
|
|
self.assertEqual(integration.get_current_prompt_text("pm_1"), "hello")
|
|
|
|
|
|
import threading # noqa: E402
|
|
import unittest.mock as mock # noqa: E402
|
|
|
|
from utils import comfyui_integration # noqa: E402
|
|
from utils.comfyui_integration import get_comfyui_integration # noqa: E402
|
|
|
|
|
|
class IntegrationTestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
ComfyUIMetadataIntegration._instance = None
|
|
comfyui_integration._integration_instance = None
|
|
# Patch the module's own ``time`` name rather than ``time.time`` globally:
|
|
# background threads from other tests (image monitor, watchdog) would
|
|
# otherwise tick this clock and make fresh registrations look stale.
|
|
self.clock = mock.patch.object(
|
|
comfyui_integration, "time", types.SimpleNamespace(time=self._now)
|
|
)
|
|
self.clock.start()
|
|
self.addCleanup(self.clock.stop)
|
|
self._ticks = 1000.0
|
|
self.integration = get_comfyui_integration()
|
|
|
|
def tearDown(self):
|
|
ComfyUIMetadataIntegration._instance = None
|
|
comfyui_integration._integration_instance = None
|
|
|
|
def _now(self):
|
|
self._ticks += 1.0
|
|
return self._ticks
|
|
|
|
|
|
class TestSingletonAndRegistry(IntegrationTestCase):
|
|
def test_singleton_keeps_registrations_across_constructions(self):
|
|
self.assertIs(get_comfyui_integration(), self.integration)
|
|
self.assertIs(ComfyUIMetadataIntegration(), self.integration)
|
|
self.integration.register_prompt("n1", "hello", {"category": "x"})
|
|
ComfyUIMetadataIntegration() # __init__ runs again: must not reset state
|
|
self.assertEqual(self.integration.get_current_prompt_text("n1"), "hello")
|
|
|
|
def test_latest_prompt_of_this_thread_wins_without_a_node_id(self):
|
|
self.integration.register_prompt("n1", "first", {})
|
|
self.integration.register_prompt("n2", "second", {})
|
|
self.assertEqual(self.integration.get_current_prompt_text(), "second")
|
|
self.assertEqual(self.integration.get_current_prompt_text("missing"), "second")
|
|
self.assertEqual(self.integration.get_current_prompt_text("n1"), "first")
|
|
|
|
def test_nothing_registered_returns_none(self):
|
|
self.assertIsNone(self.integration.get_current_prompt_text())
|
|
self.assertIsNone(self.integration.get_current_prompt_text("n1"))
|
|
|
|
|
|
class TestCrossThreadLookup(IntegrationTestCase):
|
|
def _register_in_thread(self, text):
|
|
worker = threading.Thread(
|
|
target=self.integration.register_prompt, args=("w1", text, {})
|
|
)
|
|
worker.start()
|
|
worker.join()
|
|
|
|
def test_recent_prompt_from_another_thread_is_used(self):
|
|
self._register_in_thread("from worker")
|
|
self.assertEqual(self.integration.get_current_prompt_text(), "from worker")
|
|
|
|
def test_stale_prompt_from_another_thread_is_ignored(self):
|
|
self._register_in_thread("from worker")
|
|
self._ticks += 400.0
|
|
self.assertIsNone(self.integration.get_current_prompt_text())
|
|
|
|
def test_global_entry_of_this_thread_is_found_without_thread_local(self):
|
|
self.integration.register_prompt("n1", "mine", {})
|
|
self.integration._thread_local = threading.local()
|
|
self.assertEqual(self.integration.get_current_prompt_text(), "mine")
|
|
|
|
def test_cleanup_removes_old_registrations_only(self):
|
|
self._register_in_thread("old")
|
|
self._ticks += 700.0
|
|
self.integration.register_prompt("n2", "new", {})
|
|
self.integration.cleanup_old_prompts(max_age_seconds=600)
|
|
texts = sorted(v["text"] for v in self.integration._current_prompts.values())
|
|
self.assertEqual(texts, ["new"])
|
|
self.integration.cleanup_old_prompts(max_age_seconds=600) # nothing to do
|
|
self.assertEqual(len(self.integration._current_prompts), 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|