test_server_runtime.py•4.08 kB
import json
from pathlib import Path
from typing import Any
import anyio
import pytest
from comfyui_mcp.config import ComfyUISettings, DirectorySettings, FeatureToggles
from comfyui_mcp.runtime import create_mcp_app
from comfyui_mcp.server import ComfyUIMCPServer
def _unwrap_tool_result(result: Any) -> Any:
if isinstance(result, tuple) and len(result) == 2 and isinstance(result[1], dict):
payload = result[1]
return payload.get("result", payload)
return result
def test_watch_workflows_refreshes(tmp_path: Path) -> None:
workflows = tmp_path / "workflows"
workflows.mkdir()
settings = ComfyUISettings(
workflows_path=workflows,
feature_toggles=FeatureToggles(watch_workflows=True),
directories=DirectorySettings(),
)
async def run() -> None:
server = ComfyUIMCPServer(settings)
try:
initial = await server.list_workflows()
assert initial == []
payload = {
"last_node_id": 1,
"last_link_id": 0,
"nodes": [
{
"id": 1,
"type": "CLIPTextEncode",
"class_type": "CLIPTextEncode",
"inputs": {"text": "Prompt", "clip": [2, "CLIP"]},
}
],
"links": [],
"extra": {"metadata": {"name": "watch_test"}},
}
(workflows / "watch.json").write_text(json.dumps(payload))
refreshed = await server.list_workflows()
names = [item["name"] for item in refreshed]
assert "watch_test" in names
finally:
await server.close()
anyio.run(run)
def test_runtime_tools_expose_workflow_metadata(tmp_path: Path) -> None:
workflows = tmp_path / "flows"
workflows.mkdir()
workflow = {
"last_node_id": 3,
"last_link_id": 2,
"nodes": [
{
"id": 1,
"type": "CLIPTextEncode",
"class_type": "CLIPTextEncode",
"inputs": {"text": "A test prompt", "clip": [2, "CLIP"]},
},
{
"id": 2,
"type": "CheckpointLoaderSimple",
"class_type": "CheckpointLoaderSimple",
"inputs": {"ckpt_name": "base_model.safetensors"},
},
{
"id": 3,
"type": "KSampler",
"class_type": "KSampler",
"inputs": {
"model": [2, "MODEL"],
"positive": [1, "CONDITIONING"],
"negative": [1, "CONDITIONING"],
"cfg": 7.5,
"steps": 20,
"seed": 13,
},
},
],
"links": [
{"id": 1, "from_node": 2, "from_slot": 0, "to_node": 3, "to_slot": 0},
{"id": 2, "from_node": 1, "from_slot": 0, "to_node": 3, "to_slot": 1},
],
"extra": {"metadata": {"name": "runtime_demo"}},
}
(workflows / "runtime.json").write_text(json.dumps(workflow))
settings = ComfyUISettings(workflows_path=workflows, directories=DirectorySettings())
app = create_mcp_app(settings=settings)
async def run() -> None:
async with app.settings.lifespan(app):
listed = _unwrap_tool_result(await app.call_tool("list_workflows", {}))
assert listed and listed[0]["name"] == "runtime_demo"
described = _unwrap_tool_result(await app.call_tool("describe_workflow", {"name": "runtime_demo"}))
assert described["name"] == "runtime_demo"
assert described["graph"]["nodes"]
customised = _unwrap_tool_result(
await app.call_tool(
"customize_workflow",
{"name": "runtime_demo", "changes": {"prompts": {"prompt_positive": "Updated"}}},
)
)
assert customised["diff"]
anyio.run(run)