| from __future__ import annotations |
|
|
| import asyncio |
| from copy import deepcopy |
|
|
| import pytest |
|
|
| from koolook_setup_runner import PublishedSetupRunner, SetupRunError |
| from koolook_setups import PublishedSetupRegistry, StaticSetupStorage |
| from tests.server.test_published_setup_registry import _valid_setup |
|
|
|
|
| class FakeComfyClient: |
| def __init__(self, prompt_response=None, *, history=None, queue=None, queue_error=None) -> None: |
| self.prompt_response = prompt_response or {"prompt_id": "comfy-prompt-1"} |
| self.submitted_prompts: list[dict] = [] |
| self.history = history or {} |
| self.queue = queue or {"queue_running": [], "queue_pending": []} |
| self.queue_error = queue_error |
|
|
| async def queue_prompt(self, prompt: dict) -> dict: |
| if self.queue_error is not None: |
| raise self.queue_error |
| self.submitted_prompts.append(prompt) |
| return self.prompt_response |
|
|
| async def get_history(self, prompt_id: str) -> dict: |
| return {prompt_id: self.history[prompt_id]} if prompt_id in self.history else {} |
|
|
| async def get_queue(self) -> dict: |
| return self.queue |
|
|
|
|
| def test_run_setup_injects_declared_inputs_into_cloned_prompt() -> None: |
| async def exercise() -> None: |
| setup = _valid_setup() |
| original_prompt = deepcopy(setup["apiPrompt"]) |
| registry = PublishedSetupRegistry(StaticSetupStorage([setup])) |
| comfy = FakeComfyClient() |
| runner = PublishedSetupRunner(registry, comfy) |
|
|
| result = await runner.runSetup("ltx-director-demo", {"prompt": "external prompt"}) |
|
|
| assert result == { |
| "runId": "run-000001", |
| "promptId": "comfy-prompt-1", |
| "status": "queued", |
| } |
| assert comfy.submitted_prompts == [ |
| {"12": {"class_type": "Text Multiline", "inputs": {"text": "external prompt"}}} |
| ] |
| assert setup["apiPrompt"] == original_prompt |
|
|
| asyncio.run(exercise()) |
|
|
|
|
| def test_run_setup_rejects_inputs_not_declared_by_contract() -> None: |
| async def exercise() -> None: |
| registry = PublishedSetupRegistry(StaticSetupStorage([_valid_setup()])) |
| runner = PublishedSetupRunner(registry, FakeComfyClient()) |
|
|
| with pytest.raises(SetupRunError) as exc_info: |
| await runner.runSetup("ltx-director-demo", {"prompt": "ok", "seed": 123}) |
|
|
| assert exc_info.value.status_code == 400 |
| assert exc_info.value.errors == ["input 'seed' is not declared by this setup"] |
|
|
| asyncio.run(exercise()) |
|
|
|
|
| def test_run_setup_rejects_non_callable_setup() -> None: |
| async def exercise() -> None: |
| setup = _valid_setup() |
| setup["apiPrompt"] = None |
| setup["validation"] = {"status": "draft", "diagnostics": ["API prompt conversion pending."]} |
| registry = PublishedSetupRegistry(StaticSetupStorage([setup])) |
| runner = PublishedSetupRunner(registry, FakeComfyClient()) |
|
|
| with pytest.raises(SetupRunError) as exc_info: |
| await runner.runSetup("ltx-director-demo", {"prompt": "external prompt"}) |
|
|
| assert exc_info.value.status_code == 400 |
| assert exc_info.value.errors == ["Published setup 'ltx-director-demo' is not callable."] |
|
|
| asyncio.run(exercise()) |
|
|
|
|
| def test_run_setup_translates_comfy_queue_failures() -> None: |
| async def exercise() -> None: |
| registry = PublishedSetupRegistry(StaticSetupStorage([_valid_setup()])) |
| runner = PublishedSetupRunner(registry, FakeComfyClient(queue_error=RuntimeError("server down"))) |
|
|
| with pytest.raises(SetupRunError) as exc_info: |
| await runner.runSetup("ltx-director-demo", {"prompt": "external prompt"}) |
|
|
| assert exc_info.value.status_code == 502 |
| assert exc_info.value.errors == ["ComfyUI queue request failed: server down"] |
|
|
| asyncio.run(exercise()) |
|
|
|
|
| def test_get_run_reports_running_state_from_comfy_queue() -> None: |
| async def exercise() -> None: |
| registry = PublishedSetupRegistry(StaticSetupStorage([_valid_setup()])) |
| runner = PublishedSetupRunner( |
| registry, |
| FakeComfyClient(queue={"queue_running": [[7, "comfy-prompt-1", {}]], "queue_pending": []}), |
| ) |
|
|
| queued = await runner.runSetup("ltx-director-demo", {"prompt": "external prompt"}) |
| status = await runner.getRun(queued["runId"]) |
|
|
| assert status == { |
| "runId": "run-000001", |
| "setupId": "ltx-director-demo", |
| "promptId": "comfy-prompt-1", |
| "status": "running", |
| "outputs": [], |
| } |
|
|
| asyncio.run(exercise()) |
|
|
|
|
| def test_get_run_reports_lost_when_prompt_is_missing_from_history_and_queue() -> None: |
| async def exercise() -> None: |
| registry = PublishedSetupRegistry(StaticSetupStorage([_valid_setup()])) |
| runner = PublishedSetupRunner( |
| registry, |
| FakeComfyClient(queue={"queue_running": [], "queue_pending": []}), |
| ) |
|
|
| queued = await runner.runSetup("ltx-director-demo", {"prompt": "external prompt"}) |
| status = await runner.getRun(queued["runId"]) |
|
|
| assert status == { |
| "runId": "run-000001", |
| "setupId": "ltx-director-demo", |
| "promptId": "comfy-prompt-1", |
| "status": "lost", |
| "outputs": [], |
| } |
|
|
| asyncio.run(exercise()) |
|
|
|
|
| def test_get_run_preserves_comfy_history_status_payload() -> None: |
| async def exercise() -> None: |
| registry = PublishedSetupRegistry(StaticSetupStorage([_valid_setup()])) |
| comfy_status = { |
| "completed": False, |
| "status_str": "error", |
| "messages": [["execution_error", {"node_id": "12", "exception_message": "missing file"}]], |
| } |
| runner = PublishedSetupRunner( |
| registry, |
| FakeComfyClient( |
| history={ |
| "comfy-prompt-1": { |
| "status": comfy_status, |
| "outputs": {}, |
| } |
| } |
| ), |
| ) |
|
|
| queued = await runner.runSetup("ltx-director-demo", {"prompt": "external prompt"}) |
| status = await runner.getRun(queued["runId"]) |
|
|
| assert status == { |
| "runId": "run-000001", |
| "setupId": "ltx-director-demo", |
| "promptId": "comfy-prompt-1", |
| "status": "failed", |
| "comfyStatus": comfy_status, |
| "outputs": [{"key": "video", "label": "Video", "type": "video", "items": []}], |
| } |
|
|
| asyncio.run(exercise()) |
|
|
|
|
| def test_run_setup_accepts_inputs_declared_by_app_surface_contract() -> None: |
| async def exercise() -> None: |
| setup = _valid_setup() |
| setup["inputContract"] = {"inputs": []} |
| setup["outputContract"] = {"outputs": []} |
| setup["visualGraph"]["groups"] = [ |
| {"title": "Koolook Input", "bounding": [0, 0, 300, 200]}, |
| {"title": "Koolook Output", "bounding": [400, 0, 300, 200]}, |
| ] |
| setup["visualGraph"]["nodes"][0]["pos"] = [20, 20] |
| setup["setupSurface"] = { |
| "sourceInputs": [ |
| { |
| "group": "Koolook Input", |
| "nodes": [{"id": "12", "type": "Text Multiline", "title": "Prompt"}], |
| } |
| ], |
| "outputs": [ |
| { |
| "group": "Koolook Output", |
| "nodes": [{"id": "20", "type": "Preview Image", "title": "Preview"}], |
| } |
| ], |
| "controls": [], |
| "app": { |
| "inputs": [ |
| { |
| "key": "prompt", |
| "label": "Prompt", |
| "visible": True, |
| "target": {"node": "12", "input": "text"}, |
| "default": "demo prompt", |
| } |
| ], |
| "outputs": [], |
| "results": [], |
| }, |
| } |
| registry = PublishedSetupRegistry(StaticSetupStorage([setup])) |
| comfy = FakeComfyClient() |
| runner = PublishedSetupRunner(registry, comfy) |
|
|
| await runner.runSetup("ltx-director-demo", {"prompt": "app contract prompt"}) |
|
|
| assert comfy.submitted_prompts == [ |
| {"12": {"class_type": "Text Multiline", "inputs": {"text": "app contract prompt"}}} |
| ] |
|
|
| asyncio.run(exercise()) |
|
|
|
|
| def test_run_setup_accepts_output_controls_declared_by_app_surface() -> None: |
| async def exercise() -> None: |
| setup = _valid_setup() |
| setup["inputContract"] = {"inputs": []} |
| setup["outputContract"] = {"outputs": []} |
| setup["visualGraph"] = { |
| "nodes": [ |
| { |
| "id": 200, |
| "type": "Koolook_PublishOutput", |
| "title": "Koolook Publish Output", |
| "pos": [20, 20], |
| "size": [320, 160], |
| "inputs": [], |
| "widgets_values": ["/shots/example/output", "demo", "1"], |
| } |
| ], |
| "links": [], |
| "groups": [ |
| {"title": "Koolook Input", "bounding": [0, 0, 80, 80]}, |
| {"title": "Koolook Output", "bounding": [0, 0, 380, 220]}, |
| ], |
| } |
| setup["apiPrompt"] = { |
| "200": { |
| "class_type": "Koolook_PublishOutput", |
| "inputs": { |
| "folder": "/shots/example/output", |
| "name": "demo", |
| "version": "1", |
| }, |
| } |
| } |
| setup["setupSurface"] = { |
| "sourceInputs": [{"group": "Koolook Input", "nodes": [{"id": "200", "type": "Koolook_PublishOutput", "title": "Koolook Publish Output"}]}], |
| "outputs": [{"group": "Koolook Output", "nodes": [{"id": "200", "type": "Koolook_PublishOutput", "title": "Koolook Publish Output"}]}], |
| "controls": [], |
| "app": { |
| "inputs": [], |
| "outputs": [ |
| { |
| "key": "folder", |
| "label": "Output folder", |
| "visible": True, |
| "target": {"node": "200", "input": "folder"}, |
| "default": "/shots/example/output", |
| } |
| ], |
| "results": [], |
| }, |
| } |
| registry = PublishedSetupRegistry(StaticSetupStorage([setup])) |
| comfy = FakeComfyClient() |
| runner = PublishedSetupRunner(registry, comfy) |
|
|
| await runner.runSetup("ltx-director-demo", {"folder": "/shots/custom-output"}) |
|
|
| assert comfy.submitted_prompts == [ |
| { |
| "200": { |
| "class_type": "Koolook_PublishOutput", |
| "inputs": { |
| "folder": "/shots/custom-output", |
| "name": "demo", |
| "version": "1", |
| }, |
| } |
| } |
| ] |
|
|
| asyncio.run(exercise()) |
|
|
|
|
| def test_get_run_reports_app_surface_results_without_output_contract() -> None: |
| async def exercise() -> None: |
| setup = _valid_setup() |
| setup["inputContract"] = {"inputs": []} |
| setup["outputContract"] = {"outputs": []} |
| setup["visualGraph"] = { |
| "nodes": [ |
| { |
| "id": 100, |
| "type": "Koolook_PublishInput", |
| "title": "Koolook Publish Input", |
| "pos": [20, 20], |
| "size": [320, 240], |
| "inputs": [], |
| "widgets_values": ["Img", "", "", "", ""], |
| }, |
| { |
| "id": 200, |
| "type": "Koolook_PublishOutput", |
| "title": "Koolook Publish Output", |
| "pos": [420, 20], |
| "size": [320, 160], |
| "inputs": [], |
| "widgets_values": ["/shots/example/output", "demo", "1"], |
| }, |
| { |
| "id": 300, |
| "type": "Koolook_PublishResult", |
| "title": "Koolook Publish Result", |
| "pos": [420, 220], |
| "size": [320, 120], |
| "inputs": [], |
| "widgets_values": ["/shots/example/output/default.mov"], |
| }, |
| ], |
| "links": [], |
| "groups": [ |
| {"title": "Koolook Input", "bounding": [0, 0, 360, 300]}, |
| {"title": "Koolook Output", "bounding": [400, 0, 380, 380]}, |
| ], |
| } |
| setup["apiPrompt"] = { |
| "100": { |
| "class_type": "Koolook_PublishInput", |
| "inputs": { |
| "mode": "Img", |
| "sequence_folder": "", |
| "qt_file": "", |
| "single_file": "", |
| "prompt": "", |
| }, |
| }, |
| "200": { |
| "class_type": "Koolook_PublishOutput", |
| "inputs": { |
| "folder": "/shots/example/output", |
| "name": "demo", |
| "version": "1", |
| }, |
| }, |
| "300": { |
| "class_type": "Koolook_PublishResult", |
| "inputs": {"result": "/shots/example/output/default.mov"}, |
| }, |
| } |
| setup["setupSurface"] = { |
| "sourceInputs": [ |
| { |
| "group": "Koolook Input", |
| "nodes": [{"id": "100", "type": "Koolook_PublishInput", "title": "Koolook Publish Input"}], |
| } |
| ], |
| "outputs": [ |
| { |
| "group": "Koolook Output", |
| "nodes": [ |
| {"id": "200", "type": "Koolook_PublishOutput", "title": "Koolook Publish Output"}, |
| {"id": "300", "type": "Koolook_PublishResult", "title": "Koolook Publish Result"}, |
| ], |
| } |
| ], |
| "controls": [], |
| "app": { |
| "inputs": [], |
| "outputs": [ |
| { |
| "key": "folder", |
| "label": "Output folder", |
| "visible": True, |
| "target": {"node": "200", "input": "folder"}, |
| "default": "/shots/example/output", |
| } |
| ], |
| "results": [ |
| { |
| "key": "result", |
| "label": "Result", |
| "visible": True, |
| "target": {"node": "300", "input": "result"}, |
| "default": "/shots/example/output/default.mov", |
| } |
| ], |
| }, |
| } |
| registry = PublishedSetupRegistry(StaticSetupStorage([setup])) |
| runner = PublishedSetupRunner( |
| registry, |
| FakeComfyClient( |
| history={ |
| "comfy-prompt-1": { |
| "status": {"completed": True, "status_str": "success"}, |
| "outputs": { |
| "300": { |
| "text": ["/shots/example/output/generated.mov"] |
| } |
| }, |
| } |
| } |
| ), |
| ) |
|
|
| queued = await runner.runSetup("ltx-director-demo", {}) |
| status = await runner.getRun(queued["runId"]) |
|
|
| assert status["outputs"] == [ |
| { |
| "key": "folder", |
| "label": "Output folder", |
| "type": "output", |
| "visible": True, |
| "target": {"node": "200", "input": "folder"}, |
| "default": "/shots/example/output", |
| "items": [], |
| }, |
| { |
| "key": "result", |
| "label": "Result", |
| "type": "result", |
| "visible": True, |
| "target": {"node": "300", "input": "result"}, |
| "default": "/shots/example/output/default.mov", |
| "items": [ |
| { |
| "nodeId": "300", |
| "kind": "text", |
| "value": "/shots/example/output/generated.mov", |
| } |
| ], |
| }, |
| ] |
|
|
| asyncio.run(exercise()) |
|
|
|
|
| def test_run_setup_uses_execution_map_router_to_keep_selected_writer() -> None: |
| async def exercise() -> None: |
| setup = _valid_setup() |
| setup["inputContract"] = {"inputs": []} |
| setup["outputContract"] = {"outputs": []} |
| setup["visualGraph"] = { |
| "nodes": [ |
| { |
| "id": 100, |
| "type": "Koolook_PublishInput", |
| "inputs": [ |
| {"name": "mode", "widget": {"name": "mode"}}, |
| {"name": "single_file", "widget": {"name": "single_file"}}, |
| ], |
| }, |
| {"id": 200, "type": "Koolook_PublishOutput", "inputs": []}, |
| {"id": 300, "type": "RMBG", "inputs": [{"name": "image", "link": 1}]}, |
| {"id": 400, "type": "Koolook_PublishRouter", "inputs": [{"name": "selector"}, {"name": "payload"}]}, |
| {"id": 301, "type": "EasyAIPipeline", "inputs": []}, |
| {"id": 302, "type": "EasyAIPipeline", "inputs": []}, |
| {"id": 303, "type": "EasyAIPipeline", "inputs": []}, |
| {"id": 311, "type": "SaveEXRFrames", "inputs": []}, |
| {"id": 312, "type": "Easy_VideoCombine", "inputs": []}, |
| {"id": 313, "type": "SaveImageAndPromptExact", "inputs": [{"name": "filepath", "widget": {"name": "filepath"}}]}, |
| ], |
| "links": [], |
| "groups": [], |
| } |
| setup["apiPrompt"] = { |
| "100": { |
| "class_type": "Koolook_PublishInput", |
| "inputs": { |
| "mode": "Img", |
| "sequence_folder": "", |
| "qt_file": "", |
| "single_file": "/shots/source/input.png", |
| "prompt": "", |
| }, |
| }, |
| "200": { |
| "class_type": "Koolook_PublishOutput", |
| "inputs": {"folder": "/shots/output", "name": "publish-OUT", "version": "2"}, |
| }, |
| "300": {"class_type": "RMBG", "inputs": {"image": ["100", 2]}}, |
| "301": {"class_type": "EasyAIPipeline", "inputs": {"base_directory_path": ["200", 0], "extension": ".exr"}}, |
| "302": {"class_type": "EasyAIPipeline", "inputs": {"base_directory_path": ["200", 0], "extension": ".mov"}}, |
| "303": {"class_type": "EasyAIPipeline", "inputs": {"base_directory_path": ["200", 0], "extension": ".png"}}, |
| "400": {"class_type": "Koolook_PublishRouter", "inputs": {"selector": ["100", 4], "payload": ["300", 0]}}, |
| "311": {"class_type": "SaveEXRFrames", "inputs": {"filepath": ["301", 0], "images": ["400", 0]}}, |
| "312": {"class_type": "Easy_VideoCombine", "inputs": {"filepath": ["302", 0], "images": ["400", 1]}}, |
| "313": {"class_type": "SaveImageAndPromptExact", "inputs": {"filepath": ["303", 0], "image": ["400", 2]}}, |
| } |
| setup["setupSurface"] = { |
| "sourceInputs": [ |
| { |
| "group": "Koolook Input", |
| "nodes": [{"id": "100", "type": "Koolook_PublishInput", "title": "Koolook Publish Input"}], |
| } |
| ], |
| "outputs": [ |
| { |
| "group": "Koolook Output", |
| "nodes": [ |
| {"id": "200", "type": "Koolook_PublishOutput", "title": "Koolook Publish Output"}, |
| {"id": "400", "type": "Koolook_PublishRouter", "title": "Koolook Publish Router"}, |
| {"id": "311", "type": "SaveEXRFrames", "title": "EXR writer"}, |
| {"id": "312", "type": "Easy_VideoCombine", "title": "QT writer"}, |
| {"id": "313", "type": "SaveImageAndPromptExact", "title": "Image writer"}, |
| ], |
| } |
| ], |
| "controls": [], |
| "app": { |
| "inputs": [ |
| { |
| "key": "single_file", |
| "label": "Single file", |
| "visible": True, |
| "target": {"node": "100", "input": "single_file"}, |
| "default": "/shots/source/input.png", |
| } |
| ], |
| "outputs": [ |
| { |
| "key": "folder", |
| "label": "Output folder", |
| "visible": True, |
| "target": {"node": "200", "input": "folder"}, |
| "default": "/shots/output", |
| }, |
| { |
| "key": "name", |
| "label": "Output name", |
| "visible": True, |
| "target": {"node": "200", "input": "name"}, |
| "default": "publish-OUT", |
| }, |
| { |
| "key": "version", |
| "label": "Version", |
| "visible": True, |
| "target": {"node": "200", "input": "version"}, |
| "default": "2", |
| }, |
| ], |
| "results": [], |
| "switch": { |
| "key": "switch", |
| "label": "Input type", |
| "visible": True, |
| "target": {"node": "100", "input": "mode"}, |
| "default": 2, |
| "options": [ |
| {"value": 0, "label": "EXR", "visible": False, "input": "single_file"}, |
| {"value": 1, "label": "QT", "visible": False, "input": "single_file"}, |
| {"value": 2, "label": "Img", "visible": True, "input": "single_file"}, |
| {"value": 3, "label": "Prompt", "visible": False, "input": "single_file"}, |
| ], |
| }, |
| }, |
| } |
| setup["executionMap"] = { |
| "version": 1, |
| "routers": [ |
| { |
| "node": "400", |
| "switchKey": "switch", |
| "selector": {"node": "100", "output": 4}, |
| "payload": {"node": "300", "output": 0}, |
| "branches": { |
| "0": {"label": "EXR", "output": 0, "writerNodes": ["311"]}, |
| "1": {"label": "QT", "output": 1, "writerNodes": ["312"]}, |
| "2": {"label": "Img", "output": 2, "writerNodes": ["313"]}, |
| "3": {"label": "Prompt", "output": 3, "writerNodes": []}, |
| }, |
| } |
| ], |
| } |
| registry = PublishedSetupRegistry(StaticSetupStorage([setup])) |
| comfy = FakeComfyClient() |
| runner = PublishedSetupRunner(registry, comfy) |
|
|
| await runner.runSetup( |
| "ltx-director-demo", |
| { |
| "switch": 2, |
| "single_file": "/shots/source/input.png", |
| "folder": "/shots/output", |
| "name": "external", |
| "version": "7", |
| }, |
| ) |
|
|
| assert set(comfy.submitted_prompts[0]) == {"100", "200", "300", "303", "400", "313"} |
| assert comfy.submitted_prompts[0]["100"]["inputs"]["mode"] == "Img" |
| assert comfy.submitted_prompts[0]["313"]["inputs"] == { |
| "filepath": ["303", 0], |
| "image": ["400", 2], |
| } |
|
|
| asyncio.run(exercise()) |
|
|
|
|
| @pytest.mark.parametrize("switch_value", [True, 99]) |
| def test_run_setup_rejects_invalid_switch_values_before_queueing(switch_value: object) -> None: |
| async def exercise() -> None: |
| setup = _valid_setup() |
| setup["inputContract"] = {"inputs": []} |
| setup["outputContract"] = {"outputs": []} |
| setup["apiPrompt"] = { |
| "100": { |
| "class_type": "Koolook_PublishInput", |
| "inputs": { |
| "mode": "Img", |
| "sequence_folder": "", |
| "qt_file": "", |
| "single_file": "/shots/source/input.png", |
| "prompt": "", |
| }, |
| }, |
| "300": {"class_type": "RMBG", "inputs": {"image": ["100", 2]}}, |
| "400": {"class_type": "Koolook_PublishRouter", "inputs": {"selector": ["100", 4], "payload": ["300", 0]}}, |
| "313": {"class_type": "SaveImageAndPromptExact", "inputs": {"image": ["400", 2]}}, |
| } |
| setup["setupSurface"] = { |
| "sourceInputs": [], |
| "outputs": [], |
| "controls": [], |
| "app": { |
| "inputs": [], |
| "outputs": [], |
| "results": [], |
| "switch": { |
| "key": "switch", |
| "label": "Input type", |
| "visible": True, |
| "target": {"node": "100", "input": "mode"}, |
| "default": 2, |
| "options": [ |
| {"value": 0, "label": "EXR", "visible": False, "input": "single_file"}, |
| {"value": 1, "label": "QT", "visible": False, "input": "single_file"}, |
| {"value": 2, "label": "Img", "visible": True, "input": "single_file"}, |
| ], |
| }, |
| }, |
| } |
| setup["executionMap"] = { |
| "version": 1, |
| "routers": [ |
| { |
| "node": "400", |
| "switchKey": "switch", |
| "selector": {"node": "100", "output": 4}, |
| "payload": {"node": "300", "output": 0}, |
| "branches": { |
| "0": {"label": "EXR", "output": 0, "writerNodes": []}, |
| "1": {"label": "QT", "output": 1, "writerNodes": []}, |
| "2": {"label": "Img", "output": 2, "writerNodes": ["313"]}, |
| }, |
| } |
| ], |
| } |
| comfy = FakeComfyClient() |
| runner = PublishedSetupRunner( |
| type("Registry", (), {"getSetup": lambda self, setup_id: setup if setup_id == "ltx-director-demo" else None})(), |
| comfy, |
| ) |
|
|
| with pytest.raises(SetupRunError) as exc_info: |
| await runner.runSetup("ltx-director-demo", {"switch": switch_value}) |
|
|
| assert exc_info.value.status_code == 400 |
| assert exc_info.value.errors == [ |
| "input 'switch' must be one of: 0 (EXR), 1 (QT), 2 (Img)" |
| ] |
| assert comfy.submitted_prompts == [] |
|
|
| asyncio.run(exercise()) |
|
|
|
|
| def test_run_setup_rejects_execution_map_branch_with_stale_writer_nodes() -> None: |
| async def exercise() -> None: |
| setup = _valid_setup() |
| setup["inputContract"] = {"inputs": []} |
| setup["outputContract"] = {"outputs": []} |
| setup["apiPrompt"] = { |
| "100": { |
| "class_type": "Koolook_PublishInput", |
| "inputs": { |
| "mode": "Img", |
| "sequence_folder": "", |
| "qt_file": "", |
| "single_file": "/shots/source/input.png", |
| "prompt": "", |
| }, |
| }, |
| "300": {"class_type": "RMBG", "inputs": {"image": ["100", 2]}}, |
| "400": {"class_type": "Koolook_PublishRouter", "inputs": {"selector": ["100", 4], "payload": ["300", 0]}}, |
| "313": {"class_type": "SaveImageAndPromptExact", "inputs": {"image": ["400", 2]}}, |
| } |
| setup["setupSurface"] = { |
| "sourceInputs": [], |
| "outputs": [], |
| "controls": [], |
| "app": { |
| "inputs": [], |
| "outputs": [], |
| "results": [], |
| "switch": { |
| "key": "switch", |
| "label": "Input type", |
| "visible": True, |
| "target": {"node": "100", "input": "mode"}, |
| "default": 2, |
| "options": [ |
| {"value": 0, "label": "EXR", "visible": False, "input": "single_file"}, |
| {"value": 1, "label": "QT", "visible": False, "input": "single_file"}, |
| {"value": 2, "label": "Img", "visible": True, "input": "single_file"}, |
| ], |
| }, |
| }, |
| } |
| setup["executionMap"] = { |
| "version": 1, |
| "routers": [ |
| { |
| "node": "400", |
| "switchKey": "switch", |
| "selector": {"node": "100", "output": 4}, |
| "payload": {"node": "300", "output": 0}, |
| "branches": { |
| "2": {"label": "Img", "output": 2, "writerNodes": ["999"]}, |
| }, |
| } |
| ], |
| } |
| comfy = FakeComfyClient() |
| runner = PublishedSetupRunner( |
| type("Registry", (), {"getSetup": lambda self, setup_id: setup if setup_id == "ltx-director-demo" else None})(), |
| comfy, |
| ) |
|
|
| with pytest.raises(SetupRunError) as exc_info: |
| await runner.runSetup("ltx-director-demo", {"switch": 2}) |
|
|
| assert exc_info.value.status_code == 400 |
| assert exc_info.value.errors == [ |
| "execution map branch 2 for switch 'switch' references writer node(s) not present in the prompt: 999" |
| ] |
| assert comfy.submitted_prompts == [] |
|
|
| asyncio.run(exercise()) |
|
|
|
|
| def test_get_run_reports_execution_map_writer_filepath_when_saver_history_is_empty( |
| monkeypatch: pytest.MonkeyPatch, |
| tmp_path, |
| ) -> None: |
| async def exercise() -> None: |
| import k_ai_pipeline |
|
|
| makedirs_calls = [] |
| monkeypatch.setattr( |
| k_ai_pipeline.os, |
| "makedirs", |
| lambda *args, **kwargs: makedirs_calls.append((args, kwargs)), |
| ) |
| setup = _valid_setup() |
| output_folder = tmp_path / "missing-output" |
| setup["inputContract"] = {"inputs": []} |
| setup["outputContract"] = {"outputs": []} |
| setup["apiPrompt"] = { |
| "100": { |
| "class_type": "Koolook_PublishInput", |
| "inputs": { |
| "mode": "Img", |
| "sequence_folder": "", |
| "qt_file": "", |
| "single_file": "/shots/source/input.png", |
| "prompt": "", |
| }, |
| }, |
| "200": { |
| "class_type": "Koolook_PublishOutput", |
| "inputs": {"folder": str(output_folder), "name": "publish-OUT", "version": "2"}, |
| }, |
| "300": {"class_type": "RMBG", "inputs": {"image": ["100", 2]}}, |
| "303": { |
| "class_type": "EasyAIPipeline", |
| "inputs": { |
| "shot_duration": 81, |
| "seed_value": 1, |
| "instruction": "", |
| "base_directory_path": ["200", 0], |
| "extension": ".png", |
| "shot_name": ["200", 1], |
| "ai_method": "", |
| "version": ["200", 2], |
| "disable_versioning": False, |
| "enable_overwrite": True, |
| "no_subfolders": True, |
| }, |
| }, |
| "400": {"class_type": "Koolook_PublishRouter", "inputs": {"selector": ["100", 4], "payload": ["300", 0]}}, |
| "313": {"class_type": "SaveImageAndPromptExact", "inputs": {"filepath": ["303", 0], "image": ["400", 2]}}, |
| "999": {"class_type": "Koolook_PublishResult", "inputs": {"result": ""}}, |
| } |
| setup["setupSurface"] = { |
| "sourceInputs": [ |
| { |
| "group": "Koolook Input", |
| "nodes": [{"id": "100", "type": "Koolook_PublishInput", "title": "Koolook Publish Input"}], |
| } |
| ], |
| "outputs": [ |
| { |
| "group": "Koolook Output", |
| "nodes": [ |
| {"id": "200", "type": "Koolook_PublishOutput", "title": "Koolook Publish Output"}, |
| {"id": "400", "type": "Koolook_PublishRouter", "title": "Koolook Publish Router"}, |
| {"id": "313", "type": "SaveImageAndPromptExact", "title": "Image writer"}, |
| {"id": "999", "type": "Koolook_PublishResult", "title": "Koolook Publish Result"}, |
| ], |
| } |
| ], |
| "controls": [], |
| "app": { |
| "inputs": [ |
| { |
| "key": "single_file", |
| "label": "Single file", |
| "visible": True, |
| "target": {"node": "100", "input": "single_file"}, |
| "default": "/shots/source/input.png", |
| } |
| ], |
| "outputs": [ |
| {"key": "folder", "label": "Output folder", "visible": True, "target": {"node": "200", "input": "folder"}, "default": str(output_folder)}, |
| {"key": "name", "label": "Output name", "visible": True, "target": {"node": "200", "input": "name"}, "default": "publish-OUT"}, |
| {"key": "version", "label": "Version", "visible": True, "target": {"node": "200", "input": "version"}, "default": "2"}, |
| ], |
| "results": [ |
| {"key": "result", "label": "Result", "visible": True, "target": {"node": "999", "input": "result"}, "default": ""} |
| ], |
| "switch": { |
| "key": "switch", |
| "label": "Input type", |
| "visible": True, |
| "target": {"node": "100", "input": "mode"}, |
| "default": 2, |
| "options": [ |
| {"value": 0, "label": "EXR", "visible": False, "input": "single_file"}, |
| {"value": 1, "label": "QT", "visible": False, "input": "single_file"}, |
| {"value": 2, "label": "Img", "visible": True, "input": "single_file"}, |
| {"value": 3, "label": "Prompt", "visible": False, "input": "single_file"}, |
| ], |
| }, |
| }, |
| } |
| setup["executionMap"] = { |
| "version": 1, |
| "routers": [ |
| { |
| "node": "400", |
| "switchKey": "switch", |
| "selector": {"node": "100", "output": 4}, |
| "payload": {"node": "300", "output": 0}, |
| "branches": { |
| "2": {"label": "Img", "output": 2, "writerNodes": ["313"]}, |
| }, |
| } |
| ], |
| } |
| comfy = FakeComfyClient() |
| runner = PublishedSetupRunner( |
| type("Registry", (), {"getSetup": lambda self, setup_id: setup if setup_id == "ltx-director-demo" else None})(), |
| comfy, |
| ) |
|
|
| queued = await runner.runSetup( |
| "ltx-director-demo", |
| {"switch": 2, "folder": str(output_folder), "name": "external", "version": "7"}, |
| ) |
| comfy.history[queued["promptId"]] = { |
| "status": {"completed": True, "status_str": "success"}, |
| "prompt": [2, queued["promptId"], comfy.submitted_prompts[0], {}], |
| "outputs": {"313": {"images": []}}, |
| } |
|
|
| status = await runner.getRun(queued["runId"]) |
|
|
| |
| |
| |
| canonical_folder = str(output_folder).replace("\\", "/") |
| assert status["outputs"][-1]["items"] == [ |
| { |
| "nodeId": "313", |
| "kind": "text", |
| "value": f"{canonical_folder}/v007/external_v007.png", |
| } |
| ] |
| assert makedirs_calls == [] |
| assert not output_folder.exists() |
|
|
| asyncio.run(exercise()) |
|
|
|
|
| @pytest.mark.parametrize( |
| ("switch_value", "mode_label", "branch_node", "source_node", "saver_node", "result_path"), |
| [ |
| (0, "EXR", "301", "401", "311", "/shots/output/publish-OUT_v002.exr"), |
| (1, "QT", "302", "402", "312", "/shots/output/publish-OUT_v002.mov"), |
| (2, "Img", "303", "403", "313", "/shots/output/publish-OUT_v002.png"), |
| ], |
| ) |
| def test_get_run_resolves_selected_switch_result_when_result_node_history_is_missing( |
| switch_value: int, |
| mode_label: str, |
| branch_node: str, |
| source_node: str, |
| saver_node: str, |
| result_path: str, |
| ) -> None: |
| async def exercise() -> None: |
| setup = _valid_setup() |
| setup["inputContract"] = {"inputs": []} |
| setup["outputContract"] = {"outputs": []} |
| setup["visualGraph"] = { |
| "nodes": [ |
| { |
| "id": 188, |
| "type": "Koolook_PublishInput", |
| "inputs": [ |
| {"name": "mode", "widget": {"name": "mode"}}, |
| {"name": "sequence_folder", "widget": {"name": "sequence_folder"}}, |
| {"name": "qt_file", "widget": {"name": "qt_file"}}, |
| {"name": "single_file", "widget": {"name": "single_file"}}, |
| ], |
| }, |
| { |
| "id": 205, |
| "type": "Koolook_PublishResult", |
| "inputs": [{"name": "result", "link": 900}], |
| }, |
| ], |
| "links": [[900, 179, 0, 205, 0, "STRING"]], |
| "groups": [], |
| } |
| setup["apiPrompt"] = { |
| "188": { |
| "class_type": "Koolook_PublishInput", |
| "inputs": { |
| "mode": "Img", |
| "sequence_folder": "", |
| "qt_file": "", |
| "single_file": "/shots/source/input.png", |
| "prompt": "", |
| }, |
| }, |
| "179": { |
| "class_type": "Switch any [Crystools]", |
| "inputs": { |
| "index": ["188", 4], |
| "value0": ["301", 0], |
| "value1": ["302", 0], |
| "value2": ["303", 0], |
| }, |
| }, |
| "180": { |
| "class_type": "Switch any [Crystools]", |
| "inputs": { |
| "index": ["188", 4], |
| "value0": ["401", 0], |
| "value1": ["402", 0], |
| "value2": ["403", 0], |
| }, |
| }, |
| "301": { |
| "class_type": "EasyAIPipeline", |
| "inputs": {"WRITE_file_path": "/shots/output/publish-OUT_v002.exr"}, |
| }, |
| "302": { |
| "class_type": "Easy_VideoCombine", |
| "inputs": {"video_path": "/shots/output/publish-OUT_v002.mov"}, |
| }, |
| "303": { |
| "class_type": "EasyAIPipeline", |
| "inputs": {"WRITE_file_path": "/shots/output/publish-OUT_v002.png"}, |
| }, |
| "311": { |
| "class_type": "SaveEXRFrames", |
| "inputs": {"filepath": ["301", 0], "images": ["180", 0]}, |
| }, |
| "312": { |
| "class_type": "Easy_VideoCombine", |
| "inputs": {"filepath": ["302", 0], "images": ["180", 0]}, |
| }, |
| "313": { |
| "class_type": "SaveImageAndPromptExact", |
| "inputs": {"filepath": ["303", 0], "images": ["180", 0]}, |
| }, |
| "401": { |
| "class_type": "LoadEXR", |
| "inputs": {"path": "/shots/source/input.exr"}, |
| }, |
| "402": { |
| "class_type": "VHS_LoadVideoPath", |
| "inputs": {"video": "/shots/source/input.mov"}, |
| }, |
| "403": { |
| "class_type": "VHS_LoadImagePath", |
| "inputs": {"image": "/shots/source/input.png"}, |
| }, |
| "205": { |
| "class_type": "Koolook_PublishResult", |
| "inputs": {"result": ["179", 0]}, |
| }, |
| "900": { |
| "class_type": "Preview Image", |
| "inputs": {"images": ["301", 0]}, |
| }, |
| } |
| setup["setupSurface"] = { |
| "sourceInputs": [ |
| { |
| "group": "Koolook Input", |
| "nodes": [{"id": "188", "type": "Koolook_PublishInput", "title": "Koolook Publish Input"}], |
| } |
| ], |
| "outputs": [ |
| { |
| "group": "Koolook Output", |
| "nodes": [ |
| {"id": "205", "type": "Koolook_PublishResult", "title": "Koolook Publish Result"}, |
| {"id": "301", "type": "EasyAIPipeline", "title": "EXR path"}, |
| {"id": "302", "type": "EasyAIPipeline", "title": "QT path"}, |
| {"id": "303", "type": "EasyAIPipeline", "title": "Img path"}, |
| {"id": "311", "type": "SaveEXRFrames", "title": "EXR saver"}, |
| {"id": "312", "type": "Easy_VideoCombine", "title": "QT saver"}, |
| {"id": "313", "type": "SaveImageAndPromptExact", "title": "Image saver"}, |
| ], |
| } |
| ], |
| "controls": [], |
| "app": { |
| "inputs": [ |
| { |
| "key": "sequence_folder", |
| "label": "Sequence folder", |
| "visible": True, |
| "target": {"node": "188", "input": "sequence_folder"}, |
| "default": "", |
| }, |
| { |
| "key": "qt_file", |
| "label": "QT file", |
| "visible": True, |
| "target": {"node": "188", "input": "qt_file"}, |
| "default": "", |
| }, |
| { |
| "key": "single_file", |
| "label": "Single file", |
| "visible": True, |
| "target": {"node": "188", "input": "single_file"}, |
| "default": "/shots/source/input.png", |
| } |
| ], |
| "outputs": [], |
| "results": [ |
| { |
| "key": "result", |
| "label": "Result", |
| "visible": True, |
| "target": {"node": "205", "input": "result"}, |
| "default": "", |
| } |
| ], |
| "switch": { |
| "key": "switch", |
| "label": "Input type", |
| "visible": True, |
| "target": {"node": "188", "input": "mode"}, |
| "default": 2, |
| "options": [ |
| {"value": 0, "label": "EXR", "visible": True, "input": "sequence_folder"}, |
| {"value": 1, "label": "QT", "visible": True, "input": "qt_file"}, |
| {"value": 2, "label": "Img", "visible": True, "input": "single_file"}, |
| ], |
| }, |
| }, |
| } |
| registry = PublishedSetupRegistry(StaticSetupStorage([setup])) |
| comfy = FakeComfyClient( |
| history={ |
| "comfy-prompt-1": { |
| "status": {"completed": True, "status_str": "success"}, |
| "outputs": { |
| branch_node: { |
| "text": [result_path], |
| } |
| }, |
| } |
| } |
| ) |
| runner = PublishedSetupRunner( |
| registry, |
| comfy, |
| ) |
|
|
| queued = await runner.runSetup("ltx-director-demo", {"switch": switch_value}) |
| status = await runner.getRun(queued["runId"]) |
|
|
| assert set(comfy.submitted_prompts[0]) == { |
| "188", |
| "179", |
| "180", |
| "205", |
| branch_node, |
| source_node, |
| saver_node, |
| } |
| assert comfy.submitted_prompts[0]["188"]["inputs"]["mode"] == mode_label |
| assert comfy.submitted_prompts[0]["179"]["inputs"] == { |
| "index": ["188", 4], |
| f"value{switch_value}": [branch_node, 0], |
| } |
| assert comfy.submitted_prompts[0]["180"]["inputs"] == { |
| "index": ["188", 4], |
| f"value{switch_value}": [source_node, 0], |
| } |
| assert status["outputs"] == [ |
| { |
| "key": "result", |
| "label": "Result", |
| "type": "result", |
| "visible": True, |
| "target": {"node": "205", "input": "result"}, |
| "default": "", |
| "items": [ |
| { |
| "nodeId": branch_node, |
| "kind": "text", |
| "value": result_path, |
| } |
| ], |
| } |
| ] |
|
|
| asyncio.run(exercise()) |
|
|
