132 lines
5.1 KiB
Python
132 lines
5.1 KiB
Python
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import ModuleType, SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
|
|
class TestNodesPreview(unittest.TestCase):
|
|
def setUp(self):
|
|
class _NodeOutput:
|
|
def __init__(self, *result, ui=None):
|
|
self.result = result
|
|
self.ui = ui or {}
|
|
|
|
class _Port:
|
|
def __init__(self, kind, name, optional=False, **kwargs):
|
|
self.kind = kind
|
|
self.name = name
|
|
self.optional = optional
|
|
self.kwargs = kwargs
|
|
|
|
class _PortFactory:
|
|
def __init__(self, kind):
|
|
self.kind = kind
|
|
|
|
def Input(self, name, optional=False, **kwargs):
|
|
return _Port(self.kind, name, optional=optional, **kwargs)
|
|
|
|
def Output(self, name, **kwargs):
|
|
return _Port(self.kind, name, optional=False, **kwargs)
|
|
|
|
class _Schema:
|
|
def __init__(self, **kwargs):
|
|
self.__dict__.update(kwargs)
|
|
self.inputs = list(kwargs.get("inputs", []))
|
|
self.outputs = list(kwargs.get("outputs", []))
|
|
self.hidden = list(kwargs.get("hidden", []))
|
|
self.is_output_node = bool(kwargs.get("is_output_node", False))
|
|
|
|
class _ComfyNode:
|
|
OUTPUT_NODE = False
|
|
RETURN_TYPES = ()
|
|
|
|
def __init_subclass__(cls, **kwargs):
|
|
super().__init_subclass__(**kwargs)
|
|
schema = cls.define_schema() if hasattr(cls, "define_schema") else None
|
|
cls._schema = schema
|
|
cls.OUTPUT_NODE = bool(getattr(schema, "is_output_node", False))
|
|
cls.RETURN_TYPES = tuple(getattr(port, "kind", "") for port in getattr(schema, "outputs", []))
|
|
|
|
def run(self, **kwargs):
|
|
out = self.execute(**kwargs)
|
|
if isinstance(out, _NodeOutput):
|
|
return {"ui": out.ui, "result": out.result}
|
|
return {"ui": {}, "result": out if isinstance(out, tuple) else (out,)}
|
|
|
|
io_module = SimpleNamespace(
|
|
ComfyNode=_ComfyNode,
|
|
Schema=_Schema,
|
|
NodeOutput=_NodeOutput,
|
|
Combo=_PortFactory("COMBO"),
|
|
String=_PortFactory("STRING"),
|
|
Image=_PortFactory("IMAGE"),
|
|
Float=_PortFactory("FLOAT"),
|
|
Int=_PortFactory("INT"),
|
|
Mask=_PortFactory("MASK"),
|
|
Hidden=SimpleNamespace(unique_id="UNIQUE_ID"),
|
|
)
|
|
comfy_api_module = ModuleType("comfy_api")
|
|
comfy_api_latest_module = ModuleType("comfy_api.latest")
|
|
comfy_api_latest_module.io = io_module
|
|
|
|
self.mock_torch = MagicMock()
|
|
self.mock_torch.from_numpy = MagicMock(return_value=MagicMock())
|
|
self.mock_nodes = MagicMock()
|
|
self.mock_preview_image = MagicMock()
|
|
self.mock_nodes.PreviewImage = MagicMock(return_value=self.mock_preview_image)
|
|
self.mock_preview_image.save_images.return_value = {
|
|
"ui": {"images": [{"filename": "preview.png", "type": "temp"}]}
|
|
}
|
|
|
|
modules_patch = {
|
|
"torch": self.mock_torch,
|
|
"torch.nn": MagicMock(),
|
|
"torch.nn.functional": MagicMock(),
|
|
"numpy": MagicMock(),
|
|
"PIL": MagicMock(),
|
|
"PIL.Image": MagicMock(),
|
|
"cv2": MagicMock(),
|
|
"nodes": self.mock_nodes,
|
|
"comfy_api": comfy_api_module,
|
|
"comfy_api.latest": comfy_api_latest_module,
|
|
}
|
|
|
|
self.patcher = patch.dict(sys.modules, modules_patch)
|
|
self.patcher.start()
|
|
|
|
import comfyui_ty360_photo_edit.nodes
|
|
|
|
self.nodes_module = comfyui_ty360_photo_edit.nodes
|
|
|
|
def tearDown(self):
|
|
self.patcher.stop()
|
|
|
|
def test_preview_node_saves_preview(self):
|
|
node = self.nodes_module.TY360PreviewNode()
|
|
dummy_erp = MagicMock()
|
|
|
|
res = node.run(erp_image=dummy_erp)
|
|
|
|
self.mock_preview_image.save_images.assert_called_once()
|
|
self.assertIn("pano_input_images", res["ui"])
|
|
|
|
def test_node_mappings_and_roundtrip_output_order(self):
|
|
self.assertIn("TY360Preview", self.nodes_module.NODE_CLASS_MAPPINGS)
|
|
self.assertIn("TY360Roundtrip", self.nodes_module.NODE_CLASS_MAPPINGS)
|
|
self.assertIn("TY360Pasteback", self.nodes_module.NODE_CLASS_MAPPINGS)
|
|
outputs = [port.name for port in self.nodes_module.TY360RoundtripNode._schema.outputs]
|
|
self.assertEqual(outputs, ["erp_image", "rect_image", "roundtrip_state"])
|
|
|
|
def test_web_registration_strings_use_ty360_names(self):
|
|
repo_root = Path(__file__).resolve().parent.parent
|
|
editor_js = (repo_root / "web" / "pano_editor.js").read_text(encoding="utf-8")
|
|
preview_identity_js = (repo_root / "web" / "pano_preview_identity.js").read_text(encoding="utf-8")
|
|
self.assertIn('installEditorButton(nodeType, nodeData, "TY360Roundtrip", "Open Roundtrip Editor")', editor_js)
|
|
self.assertIn('modalTitle: "TY360 Preview"', editor_js)
|
|
self.assertIn('name === "TY360Preview" || name === "TY360 Preview"', preview_identity_js)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|