50 lines
2.1 KiB
Python
50 lines
2.1 KiB
Python
import copy
|
|
import unittest
|
|
|
|
from plugin_loader import plugin
|
|
import chicktools.nodes.chick_save as chick_save
|
|
|
|
|
|
class PrivateSaverTests(unittest.TestCase):
|
|
def test_schema_does_not_request_prompt_metadata(self):
|
|
schema = plugin.NODE_CLASS_MAPPINGS["ChickSaveImage"].INPUT_TYPES()
|
|
self.assertNotIn("prompt", schema.get("hidden", {}))
|
|
self.assertEqual(schema["hidden"]["extra_pnginfo"], "EXTRA_PNGINFO")
|
|
self.assertEqual(plugin.NODE_DISPLAY_NAME_MAPPINGS["ChickSaveImage"],
|
|
"Chick Private Save")
|
|
|
|
def test_allowlist_keeps_aigc_without_mutating_shared_metadata(self):
|
|
captured = {}
|
|
original = chick_save._ComfySaveImage.save_images
|
|
|
|
def host_save(node, images, filename_prefix="ComfyUI", prompt=None,
|
|
extra_pnginfo=None, **kwargs):
|
|
captured.update(images=images, filename_prefix=filename_prefix,
|
|
prompt=prompt, extra_pnginfo=extra_pnginfo, kwargs=kwargs)
|
|
return {"ui": {"images": [{"filename": "carrier.png"}]}}
|
|
|
|
chick_save._ComfySaveImage.save_images = host_save
|
|
try:
|
|
node = plugin.NODE_CLASS_MAPPINGS["ChickSaveImage"]()
|
|
prompt = {"2": {"class_type": "ChickHideNode",
|
|
"inputs": {"password": "secret-password"}}}
|
|
metadata = {
|
|
"workflow": {"nodes": [{"id": 2, "widgets_values": ["secret-password"]}]},
|
|
"AIGC": {"Label": "1"},
|
|
}
|
|
before_prompt, before_metadata = copy.deepcopy(prompt), copy.deepcopy(metadata)
|
|
result = node.save_images(["pixels"], "carrier", prompt, metadata)
|
|
finally:
|
|
chick_save._ComfySaveImage.save_images = original
|
|
|
|
self.assertEqual(result["ui"]["images"][0]["filename"], "carrier.png")
|
|
self.assertEqual(captured["prompt"], None)
|
|
self.assertEqual(captured["extra_pnginfo"], {"AIGC": {"Label": "1"}})
|
|
self.assertNotIn("secret-password", repr(captured["extra_pnginfo"]))
|
|
self.assertEqual(prompt, before_prompt)
|
|
self.assertEqual(metadata, before_metadata)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|