fred-dev/comfy_ui_ali

1### 🗻 This file is created through the spirit of Mount Fuji at its peak2# TODO(yoland): clean up this after I get back down3import sys4import pytest5import os6import tempfile7from unittest.mock import patch8from importlib import reload9 10import folder_paths11import comfy.cli_args12from comfy.options import enable_args_parsing13enable_args_parsing()14 15 16@pytest.fixture()17def clear_folder_paths():18    # Reload the module after each test to ensure isolation19    yield20    reload(folder_paths)21 22@pytest.fixture23def temp_dir():24    with tempfile.TemporaryDirectory() as tmpdirname:25        yield tmpdirname26 27 28@pytest.fixture29def set_base_dir():30    def _set_base_dir(base_dir):31        # Mock CLI args32        with patch.object(sys, 'argv', ["main.py", "--base-directory", base_dir]):33            reload(comfy.cli_args)34            reload(folder_paths)35    yield _set_base_dir36    # Reload the modules after each test to ensure isolation37    with patch.object(sys, 'argv', ["main.py"]):38        reload(comfy.cli_args)39        reload(folder_paths)40 41 42def test_get_directory_by_type(clear_folder_paths):43    test_dir = "/test/dir"44    folder_paths.set_output_directory(test_dir)45    assert folder_paths.get_directory_by_type("output") == test_dir46    assert folder_paths.get_directory_by_type("invalid") is None47 48def test_annotated_filepath():49    assert folder_paths.annotated_filepath("test.txt") == ("test.txt", None)50    assert folder_paths.annotated_filepath("test.txt [output]") == ("test.txt", folder_paths.get_output_directory())51    assert folder_paths.annotated_filepath("test.txt [input]") == ("test.txt", folder_paths.get_input_directory())52    assert folder_paths.annotated_filepath("test.txt [temp]") == ("test.txt", folder_paths.get_temp_directory())53 54def test_get_annotated_filepath():55    default_dir = "/default/dir"56    assert folder_paths.get_annotated_filepath("test.txt", default_dir) == os.path.join(default_dir, "test.txt")57    assert folder_paths.get_annotated_filepath("test.txt [output]") == os.path.join(folder_paths.get_output_directory(), "test.txt")58 59def test_add_model_folder_path_append(clear_folder_paths):60    folder_paths.add_model_folder_path("test_folder", "/default/path", is_default=True)61    folder_paths.add_model_folder_path("test_folder", "/test/path", is_default=False)62    assert folder_paths.get_folder_paths("test_folder") == ["/default/path", "/test/path"]63 64 65def test_add_model_folder_path_insert(clear_folder_paths):66    folder_paths.add_model_folder_path("test_folder", "/test/path", is_default=False)67    folder_paths.add_model_folder_path("test_folder", "/default/path", is_default=True)68    assert folder_paths.get_folder_paths("test_folder") == ["/default/path", "/test/path"]69 70 71def test_add_model_folder_path_re_add_existing_default(clear_folder_paths):72    folder_paths.add_model_folder_path("test_folder", "/test/path", is_default=False)73    folder_paths.add_model_folder_path("test_folder", "/old_default/path", is_default=True)74    assert folder_paths.get_folder_paths("test_folder") == ["/old_default/path", "/test/path"]75    folder_paths.add_model_folder_path("test_folder", "/test/path", is_default=True)76    assert folder_paths.get_folder_paths("test_folder") == ["/test/path", "/old_default/path"]77 78 79def test_add_model_folder_path_re_add_existing_non_default(clear_folder_paths):80    folder_paths.add_model_folder_path("test_folder", "/test/path", is_default=False)81    folder_paths.add_model_folder_path("test_folder", "/default/path", is_default=True)82    assert folder_paths.get_folder_paths("test_folder") == ["/default/path", "/test/path"]83    folder_paths.add_model_folder_path("test_folder", "/test/path", is_default=False)84    assert folder_paths.get_folder_paths("test_folder") == ["/default/path", "/test/path"]85 86 87def test_recursive_search(temp_dir):88    os.makedirs(os.path.join(temp_dir, "subdir"))89    open(os.path.join(temp_dir, "file1.txt"), "w").close()90    open(os.path.join(temp_dir, "subdir", "file2.txt"), "w").close()91 92    files, dirs = folder_paths.recursive_search(temp_dir)93    assert set(files) == {"file1.txt", os.path.join("subdir", "file2.txt")}94    assert len(dirs) == 2  # temp_dir and subdir95 96def test_filter_files_extensions():97    files = ["file1.txt", "file2.jpg", "file3.png", "file4.txt"]98    assert folder_paths.filter_files_extensions(files, [".txt"]) == ["file1.txt", "file4.txt"]99    assert folder_paths.filter_files_extensions(files, [".jpg", ".png"]) == ["file2.jpg", "file3.png"]100    assert folder_paths.filter_files_extensions(files, []) == files101 102@patch("folder_paths.recursive_search")103@patch("folder_paths.folder_names_and_paths")104def test_get_filename_list(mock_folder_names_and_paths, mock_recursive_search):105    mock_folder_names_and_paths.__getitem__.return_value = (["/test/path"], {".txt"})106    mock_recursive_search.return_value = (["file1.txt", "file2.jpg"], {})107    assert folder_paths.get_filename_list("test_folder") == ["file1.txt"]108 109def test_get_save_image_path(temp_dir):110    with patch("folder_paths.output_directory", temp_dir):111        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path("test", temp_dir, 100, 100)112        assert os.path.samefile(full_output_folder, temp_dir)113        assert filename == "test"114        assert counter == 1115        assert subfolder == ""116        assert filename_prefix == "test"117 118 119def test_base_path_changes(set_base_dir):120    test_dir = os.path.abspath("/test/dir")121    set_base_dir(test_dir)122 123    assert folder_paths.base_path == test_dir124    assert folder_paths.models_dir == os.path.join(test_dir, "models")125    assert folder_paths.input_directory == os.path.join(test_dir, "input")126    assert folder_paths.output_directory == os.path.join(test_dir, "output")127    assert folder_paths.temp_directory == os.path.join(test_dir, "temp")128    assert folder_paths.user_directory == os.path.join(test_dir, "user")129 130    assert os.path.join(test_dir, "custom_nodes") in folder_paths.get_folder_paths("custom_nodes")131 132    for name in ["checkpoints", "loras", "vae", "configs", "embeddings", "controlnet", "classifiers"]:133        assert folder_paths.get_folder_paths(name)[0] == os.path.join(test_dir, "models", name)134 135 136def test_base_path_change_clears_old(set_base_dir):137    test_dir = os.path.abspath("/test/dir")138    set_base_dir(test_dir)139 140    assert len(folder_paths.get_folder_paths("custom_nodes")) == 1141 142    single_model_paths = [143        "checkpoints",144        "loras",145        "vae",146        "configs",147        "clip_vision",148        "style_models",149        "diffusers",150        "vae_approx",151        "gligen",152        "upscale_models",153        "embeddings",154        "hypernetworks",155        "photomaker",156        "classifiers",157    ]158    for name in single_model_paths:159        assert len(folder_paths.get_folder_paths(name)) == 1160 161    for name in ["controlnet", "diffusion_models", "text_encoders"]:162        assert len(folder_paths.get_folder_paths(name)) == 2163