comfyui

1"""Tests for run_workflow.py — focuses on logic that doesn't require a server."""2 3from __future__ import annotations4 5 6 7from extract_schema import extract_schema8from run_workflow import (9    ComfyRunner,10    download_outputs,11    inject_params,12    parse_input_image_arg,13)14 15 16class TestParseInputImageArg:17    def test_with_name(self, tmp_path):18        f = tmp_path / "x.png"19        f.write_text("x", encoding="utf-8")20        n, p = parse_input_image_arg(f"image={f}")21        assert n == "image"22        assert p == f23 24    def test_without_name_defaults(self, tmp_path):25        f = tmp_path / "x.png"26        f.write_text("x", encoding="utf-8")27        n, p = parse_input_image_arg(str(f))28        assert n == "image"29 30    def test_custom_name(self, tmp_path):31        f = tmp_path / "x.png"32        f.write_text("x", encoding="utf-8")33        n, p = parse_input_image_arg(f"mask_image={f}")34        assert n == "mask_image"35 36 37class TestInjectParams:38    def test_basic_injection(self, sd15_workflow):39        schema = extract_schema(sd15_workflow)40        wf, warnings = inject_params(sd15_workflow, schema, {41            "prompt": "new prompt",42            "seed": 999,43            "steps": 25,44        })45        assert wf["6"]["inputs"]["text"] == "new prompt"46        assert wf["3"]["inputs"]["seed"] == 99947        assert wf["3"]["inputs"]["steps"] == 2548        assert warnings == []49 50    def test_unknown_param_warns(self, sd15_workflow):51        schema = extract_schema(sd15_workflow)52        _, warnings = inject_params(sd15_workflow, schema, {"foobar": "x"})53        assert any("foobar" in w for w in warnings)54 55    def test_seed_minus_one_randomizes(self, sd15_workflow):56        schema = extract_schema(sd15_workflow)57        wf, warnings = inject_params(sd15_workflow, schema, {"seed": -1})58        assert wf["3"]["inputs"]["seed"] != -159        assert isinstance(wf["3"]["inputs"]["seed"], int)60        assert any("expanded" in w.lower() for w in warnings)61 62    def test_randomize_seed_when_unset(self, sd15_workflow):63        schema = extract_schema(sd15_workflow)64        original = sd15_workflow["3"]["inputs"]["seed"]65        wf, warnings = inject_params(sd15_workflow, schema, {}, randomize_seed_if_unset=True)66        assert wf["3"]["inputs"]["seed"] != original67        assert isinstance(wf["3"]["inputs"]["seed"], int)68 69    def test_does_not_mutate_original(self, sd15_workflow):70        schema = extract_schema(sd15_workflow)71        original_text = sd15_workflow["6"]["inputs"]["text"]72        inject_params(sd15_workflow, schema, {"prompt": "MUTATED"})73        assert sd15_workflow["6"]["inputs"]["text"] == original_text74 75    def test_refuses_to_overwrite_link(self):76        wf = {77            "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "x"}},78            "5": {"class_type": "EmptyLatentImage",79                  "inputs": {"width": 512, "height": 512, "batch_size": 1}},80            "6": {"class_type": "CLIPTextEncode",81                  "inputs": {"text": ["3", 0], "clip": ["1", 1]}},  # text is a link!82            "3": {"class_type": "KSampler",83                  "inputs": {"seed": 1, "steps": 20, "cfg": 7.5,84                             "sampler_name": "euler", "scheduler": "normal", "denoise": 1.0,85                             "model": ["1", 0], "positive": ["6", 0], "negative": ["6", 0],86                             "latent_image": ["5", 0]}},87            "9": {"class_type": "SaveImage", "inputs": {"filename_prefix": "x", "images": ["3", 0]}},88        }89        # Manually create a schema that has prompt pointing at 6.text90        schema = {91            "parameters": {92                "prompt": {"node_id": "6", "field": "text", "type": "string", "value": ""},93            }94        }95        wf2, warnings = inject_params(wf, schema, {"prompt": "literal value"})96        # The link should NOT have been overwritten97        assert wf2["6"]["inputs"]["text"] == ["3", 0]98        assert any("link" in w.lower() for w in warnings)99 100 101# =============================================================================102# Output download walk103# =============================================================================104 105class TestDownloadOutputsWalk:106    """Test that download_outputs walks the structure correctly."""107 108    def test_handles_videos_plural(self, tmp_path, monkeypatch):109        """Local ComfyUI uses 'videos'/'gifs' (plural) keys."""110        downloads = []111 112        class FakeRunner:113            def download_output(self, *, filename, subfolder, file_type, output_dir, preserve_subfolder, overwrite):114                downloads.append((filename, subfolder, file_type))115                p = output_dir / filename116                p.parent.mkdir(parents=True, exist_ok=True)117                p.write_bytes(b"x")118                return p119 120        outputs = {121            "9": {"images": [{"filename": "img1.png", "subfolder": "", "type": "output"}]},122            "10": {"videos": [{"filename": "vid1.mp4", "subfolder": "", "type": "output"}]},123            "11": {"gifs": [{"filename": "anim1.gif", "subfolder": "", "type": "output"}]},124        }125 126        result = download_outputs(FakeRunner(), outputs, tmp_path)127        files = sorted(d["filename"] for d in result)128        assert files == ["anim1.gif", "img1.png", "vid1.mp4"]129 130    def test_handles_video_singular_cloud(self, tmp_path):131        """Cloud uses 'video' (singular)."""132        class FakeRunner:133            def download_output(self, *, filename, subfolder, file_type, output_dir, preserve_subfolder, overwrite):134                p = output_dir / filename135                p.parent.mkdir(parents=True, exist_ok=True)136                p.write_bytes(b"x")137                return p138 139        outputs = {140            "10": {"video": [{"filename": "cloud.mp4", "subfolder": "", "type": "output"}]},141        }142        result = download_outputs(FakeRunner(), outputs, tmp_path)143        assert len(result) == 1144        assert result[0]["filename"] == "cloud.mp4"145 146    def test_preserves_subfolder(self, tmp_path):147        """When preserve_subfolder=True, server subfolder becomes local subdir."""148        class FakeRunner:149            def download_output(self, *, filename, subfolder, file_type, output_dir, preserve_subfolder, overwrite):150                if preserve_subfolder and subfolder:151                    p = output_dir / subfolder / filename152                else:153                    p = output_dir / filename154                p.parent.mkdir(parents=True, exist_ok=True)155                p.write_bytes(b"x")156                return p157 158        outputs = {159            "9": {"images": [160                {"filename": "img.png", "subfolder": "myrun", "type": "output"},161                {"filename": "img.png", "subfolder": "otherrun", "type": "output"},162            ]},163        }164        result = download_outputs(FakeRunner(), outputs, tmp_path, preserve_subfolder=True)165        files = [d["file"] for d in result]166        assert any("myrun" in f for f in files)167        assert any("otherrun" in f for f in files)168        # Both must exist (no collision)169        assert len({str(f) for f in files}) == 2170 171 172# =============================================================================173# ComfyRunner construction174# =============================================================================175 176class TestRunnerConstruction:177    def test_local_default(self):178        r = ComfyRunner()179        assert r.is_cloud is False180        assert r.host == "http://127.0.0.1:8188"181 182    def test_cloud_detection(self):183        r = ComfyRunner(host="https://cloud.comfy.org", api_key="abc")184        assert r.is_cloud is True185        assert "X-API-Key" in r.headers186 187    def test_cloud_subdomain_detected(self):188        r = ComfyRunner(host="https://staging.cloud.comfy.org", api_key="abc")189        assert r.is_cloud is True190 191    def test_partner_key_does_not_pollute_extra_data(self):192        r = ComfyRunner(host="https://cloud.comfy.org", api_key="auth-key")193        # No partner-key set → no extra_data should appear in submitted prompt194        # (This is a static check; runtime check happens in submit())195        assert r.partner_key is None196 197    def test_url_routing_local(self):198        r = ComfyRunner()199        url = r._url("/prompt")200        assert url == "http://127.0.0.1:8188/prompt"201 202    def test_url_routing_cloud(self):203        r = ComfyRunner(host="https://cloud.comfy.org", api_key="x")204        url = r._url("/prompt")205        assert url == "https://cloud.comfy.org/api/prompt"206 207    def test_url_routing_cloud_history_renamed(self):208        r = ComfyRunner(host="https://cloud.comfy.org", api_key="x")209        url = r._url("/history/abc-123")210        assert url == "https://cloud.comfy.org/api/history_v2/abc-123"211