comfyui

1"""Tests for extract_schema.py."""2 3from __future__ import annotations4 5 6from extract_schema import (7    extract_schema,8    find_negative_prompt_node,9    find_positive_prompt_node,10    trace_to_node,11)12 13 14# =============================================================================15# Connection tracing16# =============================================================================17 18class TestConnectionTracing:19    def test_direct_link(self):20        wf = {21            "1": {"class_type": "CLIPTextEncode", "inputs": {"text": "x"}},22            "2": {"class_type": "KSampler",23                  "inputs": {"positive": ["1", 0], "negative": ["1", 0]}},24        }25        assert trace_to_node(wf, ["1", 0]) == "1"26 27    def test_through_reroute(self):28        wf = {29            "1": {"class_type": "CLIPTextEncode", "inputs": {"text": "x"}},30            "2": {"class_type": "Reroute", "inputs": {"input": ["1", 0]}},31            "3": {"class_type": "Reroute", "inputs": {"input": ["2", 0]}},32        }33        assert trace_to_node(wf, ["3", 0]) == "1"34 35    def test_circular_safe(self):36        wf = {37            "1": {"class_type": "Reroute", "inputs": {"input": ["2", 0]}},38            "2": {"class_type": "Reroute", "inputs": {"input": ["1", 0]}},39        }40        # Should hit max_hops without infinite loop41        result = trace_to_node(wf, ["1", 0], max_hops=5)42        assert result in {"1", "2"}  # any node, just don't hang43 44 45class TestPositiveNegativeDetection:46    def test_basic(self, sd15_workflow):47        # In sd15_workflow.json node 6 is positive, node 7 is negative48        assert find_positive_prompt_node(sd15_workflow) == "6"49        assert find_negative_prompt_node(sd15_workflow) == "7"50 51    def test_swapped_order(self):52        wf = {53            "3": {"class_type": "KSampler",54                  "inputs": {55                      "positive": ["7", 0], "negative": ["6", 0],56                      "model": ["4", 0], "latent_image": ["5", 0],57                      "seed": 1, "steps": 20, "cfg": 7.5,58                      "sampler_name": "euler", "scheduler": "normal", "denoise": 1.0,59                  }},60            "4": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "x"}},61            "5": {"class_type": "EmptyLatentImage", "inputs": {"width": 512, "height": 512, "batch_size": 1}},62            "6": {"class_type": "CLIPTextEncode", "inputs": {"text": "ugly", "clip": ["4", 1]}},63            "7": {"class_type": "CLIPTextEncode", "inputs": {"text": "beautiful", "clip": ["4", 1]}},64        }65        # Now 7 is the positive (despite higher node ID)66        assert find_positive_prompt_node(wf) == "7"67        assert find_negative_prompt_node(wf) == "6"68 69 70# =============================================================================71# Schema extraction72# =============================================================================73 74class TestExtractSchema:75    def test_basic_sd15(self, sd15_workflow):76        schema = extract_schema(sd15_workflow)77        params = schema["parameters"]78        assert "prompt" in params79        assert "negative_prompt" in params80        assert "seed" in params81        assert "steps" in params82        assert "cfg" in params83        assert "width" in params84        assert "height" in params85 86    def test_prompt_value_correct(self, sd15_workflow):87        schema = extract_schema(sd15_workflow)88        # The positive prompt in the example is the landscape one89        assert "landscape" in schema["parameters"]["prompt"]["value"]90        assert "ugly" in schema["parameters"]["negative_prompt"]["value"]91 92    def test_model_dependencies(self, sd15_workflow):93        schema = extract_schema(sd15_workflow)94        deps = schema["model_dependencies"]95        ckpts = [d["value"] for d in deps if d["folder"] == "checkpoints"]96        assert "v1-5-pruned-emaonly.safetensors" in ckpts97 98    def test_output_nodes(self, sd15_workflow):99        schema = extract_schema(sd15_workflow)100        assert "9" in schema["output_nodes"]101 102    def test_summary(self, sd15_workflow):103        schema = extract_schema(sd15_workflow)104        s = schema["summary"]105        assert s["has_negative_prompt"] is True106        assert s["has_seed"] is True107        assert s["is_video_workflow"] is False108        assert s["parameter_count"] > 5109 110    def test_flux_workflow(self, flux_workflow):111        schema = extract_schema(flux_workflow)112        # Flux uses RandomNoise for seed113        assert schema["summary"]["has_seed"] is True114        # Flux has only positive prompt (no negative encoder)115        assert schema["summary"]["has_negative_prompt"] is False116 117    def test_video_detected(self, video_workflow):118        schema = extract_schema(video_workflow)119        assert schema["summary"]["is_video_workflow"] is True120 121 122class TestEmbeddingDeps:123    def test_extract_from_prompt(self):124        wf = {125            "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "x"}},126            "5": {"class_type": "EmptyLatentImage",127                  "inputs": {"width": 512, "height": 512, "batch_size": 1}},128            "6": {"class_type": "CLIPTextEncode",129                  "inputs": {130                      "text": "a cat, embedding:goodvibes, embedding:art:1.2",131                      "clip": ["1", 1]132                  }},133            "7": {"class_type": "CLIPTextEncode",134                  "inputs": {135                      "text": "ugly, embedding:badhands",136                      "clip": ["1", 1]137                  }},138            "3": {"class_type": "KSampler",139                  "inputs": {140                      "positive": ["6", 0], "negative": ["7", 0],141                      "model": ["1", 0], "latent_image": ["5", 0],142                      "seed": 1, "steps": 20, "cfg": 7.5,143                      "sampler_name": "euler", "scheduler": "normal", "denoise": 1.0,144                  }},145            "9": {"class_type": "SaveImage", "inputs": {"filename_prefix": "x", "images": ["3", 0]}},146        }147        schema = extract_schema(wf)148        names = [d["embedding_name"] for d in schema["embedding_dependencies"]]149        assert sorted(names) == ["art", "badhands", "goodvibes"]150 151 152class TestDuplicateDeduplication:153    def test_two_ksamplers_get_unique_names(self):154        wf = {155            "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "x"}},156            "5": {"class_type": "EmptyLatentImage",157                  "inputs": {"width": 512, "height": 512, "batch_size": 1}},158            "6": {"class_type": "CLIPTextEncode", "inputs": {"text": "a", "clip": ["1", 1]}},159            "7": {"class_type": "CLIPTextEncode", "inputs": {"text": "b", "clip": ["1", 1]}},160            "3": {"class_type": "KSampler",161                  "inputs": {162                      "positive": ["6", 0], "negative": ["7", 0],163                      "model": ["1", 0], "latent_image": ["5", 0],164                      "seed": 42, "steps": 20, "cfg": 7.5,165                      "sampler_name": "euler", "scheduler": "normal", "denoise": 1.0,166                  }},167            "4": {"class_type": "KSampler",168                  "inputs": {169                      "positive": ["6", 0], "negative": ["7", 0],170                      "model": ["1", 0], "latent_image": ["5", 0],171                      "seed": 99, "steps": 30, "cfg": 8.0,172                      "sampler_name": "euler", "scheduler": "normal", "denoise": 0.6,173                  }},174            "9": {"class_type": "SaveImage", "inputs": {"filename_prefix": "x", "images": ["3", 0]}},175        }176        schema = extract_schema(wf)177        params = schema["parameters"]178        # Both seeds present with disambiguated names179        seed_keys = [k for k in params if "seed" in k]180        # Symmetric: both renamed (no bare "seed")181        assert "seed" not in params182        assert "seed_3" in params and "seed_4" in params183        assert params["seed_3"]["value"] == 42184        assert params["seed_4"]["value"] == 99185