comfyui

1"""Unit tests for _common.py — pure logic only, no network."""2 3from __future__ import annotations4 5 6import pytest7 8from _common import (9    EMBEDDING_REGEX,10    cloud_endpoint,11    coerce_seed,12    folder_aliases_for,13    is_api_format,14    is_cloud_host,15    is_link,16    iter_embedding_refs,17    iter_model_deps,18    iter_nodes,19    looks_like_video_workflow,20    media_type_from_filename,21    parse_model_list,22    resolve_url,23    safe_path_join,24    unwrap_workflow,25)26 27 28# =============================================================================29# Cloud detection / URL routing30# =============================================================================31 32class TestCloudDetection:33    def test_cloud_host_exact(self):34        assert is_cloud_host("https://cloud.comfy.org") is True35        assert is_cloud_host("https://cloud.comfy.org/foo/bar") is True36 37    def test_cloud_host_subdomain(self):38        assert is_cloud_host("https://staging.cloud.comfy.org") is True39        assert is_cloud_host("https://api.cloud.comfy.org") is True40 41    def test_local_not_cloud(self):42        assert is_cloud_host("http://127.0.0.1:8188") is False43        assert is_cloud_host("http://localhost:8188") is False44        assert is_cloud_host("http://my-server.local:8188") is False45 46    def test_no_scheme(self):47        # Defaults to http://48        assert is_cloud_host("cloud.comfy.org") is True49        assert is_cloud_host("127.0.0.1:8188") is False50 51 52class TestCloudEndpointRename:53    def test_history_renamed(self):54        assert cloud_endpoint("/history") == "/history_v2"55        assert cloud_endpoint("/history/abc-123") == "/history_v2/abc-123"56 57    def test_history_v2_preserved(self):58        assert cloud_endpoint("/history_v2") == "/history_v2"59 60    def test_models_renamed(self):61        assert cloud_endpoint("/models") == "/experiment/models"62        assert cloud_endpoint("/models/checkpoints") == "/experiment/models/checkpoints"63        assert cloud_endpoint("/models/loras") == "/experiment/models/loras"64 65    def test_other_paths_unchanged(self):66        assert cloud_endpoint("/prompt") == "/prompt"67        assert cloud_endpoint("/queue") == "/queue"68 69 70class TestResolveURL:71    def test_local_no_prefix(self):72        assert resolve_url("http://127.0.0.1:8188", "/prompt") == "http://127.0.0.1:8188/prompt"73 74    def test_cloud_adds_api_prefix(self):75        assert resolve_url("https://cloud.comfy.org", "/prompt") == "https://cloud.comfy.org/api/prompt"76 77    def test_cloud_history_renamed(self):78        assert resolve_url("https://cloud.comfy.org", "/history/abc") == "https://cloud.comfy.org/api/history_v2/abc"79 80    def test_cloud_models_renamed(self):81        assert resolve_url("https://cloud.comfy.org", "/models/loras") == "https://cloud.comfy.org/api/experiment/models/loras"82 83    def test_cloud_already_has_api(self):84        # Don't double-prefix85        assert resolve_url("https://cloud.comfy.org", "/api/prompt") == "https://cloud.comfy.org/api/prompt"86 87    def test_trailing_slash_stripped(self):88        assert resolve_url("http://127.0.0.1:8188/", "/prompt") == "http://127.0.0.1:8188/prompt"89 90 91# =============================================================================92# Workflow validation93# =============================================================================94 95class TestAPIFormatDetection:96    def test_valid_api(self, sd15_workflow):97        assert is_api_format(sd15_workflow) is True98 99    def test_editor_format_rejected(self):100        editor = {"nodes": [], "links": [], "version": 0.4}101        assert is_api_format(editor) is False102 103    def test_empty_dict(self):104        assert is_api_format({}) is False105 106    def test_non_dict(self):107        assert is_api_format([]) is False108        assert is_api_format(None) is False109        assert is_api_format("string") is False110 111    def test_node_with_class_type(self):112        wf = {"3": {"class_type": "KSampler", "inputs": {}}}113        assert is_api_format(wf) is True114 115 116class TestUnwrapWorkflow:117    def test_passthrough_api_format(self, sd15_workflow):118        result = unwrap_workflow(sd15_workflow)119        assert result is sd15_workflow120 121    def test_unwrap_prompt_key(self, sd15_workflow):122        wrapped = {"prompt": sd15_workflow, "client_id": "abc"}123        result = unwrap_workflow(wrapped)124        assert result is sd15_workflow125 126    def test_editor_format_raises(self):127        with pytest.raises(ValueError, match="editor format"):128            unwrap_workflow({"nodes": [], "links": []})129 130    def test_garbage_raises(self):131        with pytest.raises(ValueError):132            unwrap_workflow({"foo": "bar"})133 134 135class TestIsLink:136    def test_valid_link(self):137        assert is_link(["3", 0]) is True138        assert is_link(["10", 1]) is True139 140    def test_non_link(self):141        assert is_link("string") is False142        assert is_link(42) is False143        assert is_link([]) is False144        assert is_link(["3"]) is False  # missing slot145        assert is_link(["3", "0"]) is False  # slot must be int146        assert is_link([3, 0]) is False  # node_id must be string147 148 149# =============================================================================150# Workflow iterators151# =============================================================================152 153class TestIterators:154    def test_iter_nodes(self, sd15_workflow):155        nodes = dict(iter_nodes(sd15_workflow))156        assert "3" in nodes157        assert nodes["3"]["class_type"] == "KSampler"158 159    def test_iter_nodes_skips_comments(self, sd15_workflow):160        # _comment is not a node161        nodes = dict(iter_nodes(sd15_workflow))162        assert "_comment" not in nodes163 164    def test_iter_model_deps(self, sd15_workflow):165        deps = list(iter_model_deps(sd15_workflow))166        names = [d["value"] for d in deps]167        assert "v1-5-pruned-emaonly.safetensors" in names168 169    def test_iter_model_deps_flux(self, flux_workflow):170        deps = list(iter_model_deps(flux_workflow))171        names = {d["value"]: d["folder"] for d in deps}172        assert names["flux1-dev.safetensors"] == "unet"173        assert names["t5xxl_fp16.safetensors"] == "clip"174        assert names["clip_l.safetensors"] == "clip"175        assert names["ae.safetensors"] == "vae"176 177 178# =============================================================================179# Embedding extraction180# =============================================================================181 182class TestEmbeddingRegex:183    def test_basic_embedding(self):184        m = EMBEDDING_REGEX.search("a cat, embedding:goodvibes, more text")185        assert m is not None186        assert m.group(1) == "goodvibes"187 188    def test_embedding_with_strength(self):189        m = EMBEDDING_REGEX.search("embedding:bad-hands-5:1.2")190        assert m is not None191        assert m.group(1) == "bad-hands-5"192 193    def test_embedding_with_extension(self):194        # Strips .pt / .safetensors / .bin195        m = EMBEDDING_REGEX.search("embedding:my-emb.pt")196        assert m is not None197        assert m.group(1) == "my-emb"198 199    def test_embedding_in_parens(self):200        m = EMBEDDING_REGEX.search("(embedding:foo:0.8)")201        assert m is not None202        assert m.group(1) == "foo"203 204    def test_multiple_in_one_string(self):205        text = "a cat, embedding:foo:1.2, and embedding:bar"206        matches = [m.group(1) for m in EMBEDDING_REGEX.finditer(text)]207        assert matches == ["foo", "bar"]208 209    def test_no_false_positive_on_word_embedding(self):210        # "embedding " (with space, no colon) should not match211        m = EMBEDDING_REGEX.search("the embedding is great")212        assert m is None213 214 215class TestIterEmbeddingRefs:216    def test_finds_in_clip_text_encode(self):217        wf = {218            "1": {"class_type": "CLIPTextEncode",219                  "inputs": {"text": "embedding:foo, embedding:bar:0.5", "clip": ["2", 0]}},220            "2": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "x"}},221        }222        refs = list(iter_embedding_refs(wf))223        names = [name for _, name in refs]224        assert names == ["foo", "bar"]225 226    def test_ignores_non_prompt_fields(self):227        wf = {228            "1": {"class_type": "CheckpointLoaderSimple",229                  "inputs": {"ckpt_name": "embedding:foo.safetensors"}},230        }231        refs = list(iter_embedding_refs(wf))232        # ckpt_name is not a prompt field — ignored233        assert refs == []234 235 236# =============================================================================237# Path safety238# =============================================================================239 240class TestSafePathJoin:241    def test_normal_join(self, tmp_path):242        p = safe_path_join(tmp_path, "subdir", "file.png")243        assert p.is_relative_to(tmp_path)244 245    def test_blocks_traversal(self, tmp_path):246        with pytest.raises(ValueError, match="path traversal"):247            safe_path_join(tmp_path, "..", "..", "etc", "passwd")248 249    def test_blocks_absolute(self, tmp_path):250        with pytest.raises(ValueError):251            safe_path_join(tmp_path, "/etc/passwd")252 253    def test_subfolder_with_filename(self, tmp_path):254        p = safe_path_join(tmp_path, "outputs", "img.png")255        assert p.name == "img.png"256        assert p.parent.name == "outputs"257 258 259# =============================================================================260# Seed coercion261# =============================================================================262 263class TestCoerceSeed:264    def test_explicit_int(self):265        assert coerce_seed(42) == 42266        assert coerce_seed(0) == 0267 268    def test_minus_one_randomizes(self):269        s = coerce_seed(-1)270        assert isinstance(s, int)271        assert 0 <= s < 2**63272 273    def test_none_randomizes(self):274        s = coerce_seed(None)275        assert isinstance(s, int)276 277    def test_string_int(self):278        # str() that converts cleanly is allowed (relaxed)279        assert coerce_seed("12345") == 12345280 281    def test_string_minus_one_randomizes(self):282        # CLI / JSON sometimes carries seed as a string.283        s = coerce_seed("-1")284        assert isinstance(s, int)285        assert 0 <= s < 2**63286        # And whitespace tolerated287        s2 = coerce_seed(" -1 ")288        assert isinstance(s2, int)289        assert 0 <= s2 < 2**63290 291 292# =============================================================================293# Model list normalization (cloud format)294# =============================================================================295 296class TestParseModelList:297    def test_local_format_strings(self):298        result = parse_model_list(["a.safetensors", "b.safetensors"])299        assert result == {"a.safetensors", "b.safetensors"}300 301    def test_cloud_format_dicts(self):302        result = parse_model_list([303            {"name": "a.safetensors", "pathIndex": 0},304            {"name": "b.safetensors", "pathIndex": 1},305        ])306        assert result == {"a.safetensors", "b.safetensors"}307 308    def test_empty(self):309        assert parse_model_list([]) == set()310 311    def test_garbage(self):312        assert parse_model_list("not a list") == set()313        assert parse_model_list(None) == set()314 315    def test_mixed_format(self):316        result = parse_model_list([317            "string-form.safetensors",318            {"name": "dict-form.safetensors"},319        ])320        assert result == {"string-form.safetensors", "dict-form.safetensors"}321 322 323# =============================================================================324# Folder aliases325# =============================================================================326 327class TestFolderAliases:328    def test_unet_aliases_diffusion_models(self):329        aliases = folder_aliases_for("unet")330        assert "unet" in aliases331        assert "diffusion_models" in aliases332 333    def test_clip_aliases_text_encoders(self):334        aliases = folder_aliases_for("clip")335        assert "clip" in aliases336        assert "text_encoders" in aliases337 338    def test_unknown_folder_returns_self(self):339        assert folder_aliases_for("checkpoints") == ["checkpoints"]340 341    def test_primary_first(self):342        # Order matters: primary should be first for human-friendly fix hints343        assert folder_aliases_for("unet")[0] == "unet"344        assert folder_aliases_for("diffusion_models")[0] == "diffusion_models"345 346 347# =============================================================================348# Media-type detection349# =============================================================================350 351class TestMediaType:352    def test_video_extensions(self):353        assert media_type_from_filename("vid.mp4") == "video"354        assert media_type_from_filename("foo.webm") == "video"355        assert media_type_from_filename("bar.gif") == "video"356 357    def test_audio_extensions(self):358        assert media_type_from_filename("song.wav") == "audio"359        assert media_type_from_filename("music.mp3") == "audio"360 361    def test_image_default(self):362        assert media_type_from_filename("pic.png") == "image"363        assert media_type_from_filename("image.jpg") == "image"364        assert media_type_from_filename("unknown.xyz") == "image"365 366    def test_3d(self):367        assert media_type_from_filename("model.glb") == "3d"368        assert media_type_from_filename("scene.gltf") == "3d"369 370 371# =============================================================================372# Cross-host header stripping (security)373# =============================================================================374 375class TestRedirectHeaderStripping:376    """Verify X-API-Key is dropped when redirect crosses to a different host377    (e.g. cloud /api/view → S3 signed URL). Critical to prevent leaking auth378    tokens to the storage backend.379    """380 381    def _build_session(self):382        from _common import _StripSensitiveOnRedirectSession, HAS_REQUESTS383        if not HAS_REQUESTS:384            import pytest385            pytest.skip("requests not installed")386        return _StripSensitiveOnRedirectSession()387 388    def test_strips_x_api_key_cross_host(self):389        import requests390        s = self._build_session()391        prep = requests.PreparedRequest()392        prep.prepare(method="GET", url="https://other.example.com/file",393                     headers={"X-API-Key": "leak", "Authorization": "Bearer x"})394        resp = requests.Response()395        orig = requests.PreparedRequest()396        orig.prepare(method="GET", url="https://cloud.comfy.org/api/view", headers={})397        resp.request = orig398        s.rebuild_auth(prep, resp)399        assert "X-API-Key" not in prep.headers400        assert "Authorization" not in prep.headers401 402    def test_preserves_x_api_key_same_host(self):403        import requests404        s = self._build_session()405        prep = requests.PreparedRequest()406        prep.prepare(method="GET", url="https://cloud.comfy.org/foo",407                     headers={"X-API-Key": "keep"})408        resp = requests.Response()409        orig = requests.PreparedRequest()410        orig.prepare(method="GET", url="https://cloud.comfy.org/bar", headers={})411        resp.request = orig412        s.rebuild_auth(prep, resp)413        assert prep.headers.get("X-API-Key") == "keep"414 415    def test_strips_cookie_cross_host(self):416        import requests417        s = self._build_session()418        prep = requests.PreparedRequest()419        prep.prepare(method="GET", url="https://other.example.com/x",420                     headers={"Cookie": "session=secret"})421        resp = requests.Response()422        orig = requests.PreparedRequest()423        orig.prepare(method="GET", url="https://cloud.comfy.org/foo", headers={})424        resp.request = orig425        s.rebuild_auth(prep, resp)426        assert "Cookie" not in prep.headers427 428 429# =============================================================================430# Video workflow detection431# =============================================================================432 433class TestVideoWorkflow:434    def test_image_workflow(self, sd15_workflow):435        assert looks_like_video_workflow(sd15_workflow) is False436 437    def test_animatediff_workflow(self, workflows_dir):438        import json439        wf = json.loads((workflows_dir / "animatediff_video.json").read_text(encoding="utf-8"))440        assert looks_like_video_workflow(wf) is True441 442    def test_wan_workflow(self, video_workflow):443        assert looks_like_video_workflow(video_workflow) is True444