162 lines
8.3 KiB
Python
162 lines
8.3 KiB
Python
import copy
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
import importlib.util
|
|
import json
|
|
from pathlib import Path
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
|
|
from plugin_loader import plugin, ROOT
|
|
|
|
spec = importlib.util.spec_from_file_location("prepare_workflow", ROOT / "tools/prepare_workflow.py")
|
|
prepare = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(prepare)
|
|
|
|
|
|
def schema(save_outputs=None, save_names=None):
|
|
result = {}
|
|
for name, cls in plugin.NODE_CLASS_MAPPINGS.items():
|
|
result[name] = {"input": cls.INPUT_TYPES(), "output": list(cls.RETURN_TYPES),
|
|
"output_name": list(cls.RETURN_NAMES), "output_node": getattr(cls, "OUTPUT_NODE", False)}
|
|
result["EmptyImage"] = {"input": {"required": {key: ["INT"] for key in ("width", "height", "batch_size", "color")}},
|
|
"output": ["IMAGE"], "output_name": ["IMAGE"]}
|
|
result["ChickSaveImage"] = {"input": {"required": {"images": ["IMAGE"], "filename_prefix": ["STRING", {"default": "private-default"}]}},
|
|
"output": save_outputs or [], "output_name": save_names or [], "output_node": True}
|
|
result["PrimitiveStringMultiline"] = {"input": {"required": {"value": ["STRING", {"default": "private-default"}]}},
|
|
"output": ["STRING"], "output_name": ["STRING"]}
|
|
return result
|
|
|
|
|
|
class PreparationTests(unittest.TestCase):
|
|
def test_notifications_do_not_depend_on_save_outputs(self):
|
|
for types, names in (([], []), (["IMAGE"], ["images"]), (["STRING"], ["filename"])):
|
|
info = schema(types, names)
|
|
report = prepare.analyze(info)
|
|
self.assertTrue(report["carrier_workflow_ready"])
|
|
self.assertTrue(report["callback_wiring_ready"])
|
|
prompts = prepare.build_prompts(info, report)
|
|
self.assertEqual(set(prompts), {"carrier-smoke", "notify-empty-smoke", "notify-text-smoke"})
|
|
self.assertEqual(prompts["carrier-smoke"]["3"]["class_type"], "ChickSaveImage")
|
|
self.assertEqual(prompts["notify-empty-smoke"]["1"]["inputs"], {"callback_url": "", "id": "demo-task-001"})
|
|
|
|
def test_text_source_slots_and_canvas_preserve_wiring(self):
|
|
for types, names, text_slot in ((["STRING"], ["STRING"], 0),
|
|
(["INT", "STRING"], ["length", "text"], 1)):
|
|
info = schema()
|
|
info["PrimitiveStringMultiline"].update(output=types, output_name=names)
|
|
report = prepare.analyze(info)
|
|
prompt = prepare.build_prompts(info, report)["notify-text-smoke"]
|
|
self.assertEqual(prompt["2"]["inputs"], {
|
|
"callback_url": "", "id": "demo-task-001", "msg": ["1", text_slot], "passthrough": ["1", text_slot]})
|
|
canvas = prepare.build_canvas(prompt, info)
|
|
nodes = {node["id"]: node for node in canvas["nodes"]}
|
|
self.assertEqual(nodes[2]["widgets_values"], ["", "demo-task-001"])
|
|
for id, source, slot, target, port, kind in canvas["links"]:
|
|
self.assertIn(id, nodes[source]["outputs"][slot]["links"])
|
|
self.assertEqual(nodes[target]["inputs"][port]["link"], id)
|
|
self.assertEqual(nodes[source]["outputs"][slot]["type"], kind)
|
|
self.assertIn(nodes[target]["inputs"][port]["type"], (kind, "*"))
|
|
self.assertEqual(nodes[2]["outputs"], [{"name": "passthrough", "type": "*", "slot_index": 0, "links": []}])
|
|
empty = prepare.build_canvas(prepare.build_prompts(info, report)["notify-empty-smoke"], info)
|
|
self.assertEqual(empty["nodes"][0]["inputs"], [
|
|
{"name": "msg", "type": "STRING", "link": None},
|
|
{"name": "passthrough", "type": "*", "link": None}])
|
|
with tempfile.TemporaryDirectory() as temp:
|
|
prepare.write_preparation(info, Path(temp))
|
|
self.assertEqual(len(list(Path(temp).glob("*.json"))), 7)
|
|
for file in Path(temp).glob("*.json"):
|
|
self.assertNotIn("private-default", file.read_text())
|
|
|
|
def test_missing_media_does_not_block_independent_notifications(self):
|
|
for name in ("ChickHideNode", "ChickSaveImage", "EmptyImage"):
|
|
info = schema()
|
|
del info[name]
|
|
report = prepare.analyze(info)
|
|
self.assertFalse(report["carrier_workflow_ready"])
|
|
self.assertTrue(report["callback_wiring_ready"])
|
|
self.assertEqual(set(prepare.build_prompts(info, report)), {"notify-empty-smoke", "notify-text-smoke"})
|
|
info = schema()
|
|
info["ChickSaveImage"]["input"]["required"]["private_key"] = ["STRING"]
|
|
self.assertFalse(prepare.analyze(info)["carrier_workflow_ready"])
|
|
|
|
def test_legacy_or_incompatible_notification_schema_is_blocked(self):
|
|
base = schema()
|
|
bad = []
|
|
missing = copy.deepcopy(base)
|
|
del missing["NotifyMeNode"]
|
|
bad.append(missing)
|
|
for change in ("legacy", "required_msg", "wrong_type", "missing_id", "wrong_id", "missing_passthrough", "wrong_passthrough",
|
|
"list", "output", "nonterminal"):
|
|
info = copy.deepcopy(base)
|
|
node = info["NotifyMeNode"]
|
|
if change == "legacy":
|
|
node["input"]["required"].update(image=["IMAGE"], preview_url=["STRING"])
|
|
elif change == "required_msg":
|
|
node["input"]["required"]["msg"] = node["input"]["optional"].pop("msg")
|
|
elif change == "wrong_type":
|
|
node["input"]["optional"]["msg"] = ["IMAGE"]
|
|
elif change == "missing_id":
|
|
del node["input"]["optional"]["id"]
|
|
elif change == "wrong_id":
|
|
node["input"]["optional"]["id"] = ["INT"]
|
|
elif change == "missing_passthrough":
|
|
del node["input"]["optional"]["passthrough"]
|
|
elif change == "wrong_passthrough":
|
|
node["input"]["optional"]["passthrough"] = ["IMAGE"]
|
|
elif change == "list":
|
|
node["is_input_list"] = True
|
|
elif change == "output":
|
|
node["output"] = ["IMAGE"]
|
|
else:
|
|
node["output_node"] = False
|
|
bad.append(info)
|
|
for info in bad:
|
|
report = prepare.analyze(info)
|
|
self.assertTrue(report["carrier_workflow_ready"])
|
|
self.assertFalse(report["callback_wiring_ready"])
|
|
self.assertEqual(set(prepare.build_prompts(info, report)), {"carrier-smoke"})
|
|
|
|
def test_missing_or_ambiguous_text_source_still_allows_empty_message(self):
|
|
for change in ("missing", "multiple", "list", "extra"):
|
|
info = schema()
|
|
node = info["PrimitiveStringMultiline"]
|
|
if change == "missing":
|
|
del info["PrimitiveStringMultiline"]
|
|
elif change == "multiple":
|
|
node["output"] = ["STRING", "STRING"]
|
|
elif change == "list":
|
|
node["output_is_list"] = [True]
|
|
else:
|
|
node["input"]["required"]["private_key"] = ["STRING"]
|
|
report = prepare.analyze(info)
|
|
self.assertTrue(report["callback_wiring_ready"])
|
|
self.assertFalse(report["text_callback_workflow_ready"])
|
|
self.assertEqual(set(prepare.build_prompts(info, report)), {"carrier-smoke", "notify-empty-smoke"})
|
|
|
|
def test_live_probe_only_reads_schema(self):
|
|
requests = []
|
|
info = schema()
|
|
class Handler(BaseHTTPRequestHandler):
|
|
def log_message(self, *args):
|
|
pass
|
|
def do_GET(self):
|
|
requests.append(("GET", self.path))
|
|
self.send_response(200)
|
|
self.end_headers()
|
|
self.wfile.write(json.dumps(info).encode())
|
|
def do_POST(self):
|
|
requests.append(("POST", self.path))
|
|
self.send_error(405)
|
|
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
|
|
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
|
thread.start()
|
|
try:
|
|
loaded = prepare.read_info(f"http://127.0.0.1:{server.server_port}")
|
|
self.assertTrue(prepare.analyze(loaded)["carrier_workflow_ready"])
|
|
self.assertEqual(requests, [("GET", "/object_info")])
|
|
finally:
|
|
server.shutdown()
|
|
server.server_close()
|
|
thread.join()
|