mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-12 07:10:52 +08:00
121 lines
5.4 KiB
Python
121 lines
5.4 KiB
Python
### 🗻 This file is created through the spirit of Mount Fuji at its peak
|
|
# TODO(yoland): clean up this after I get back down
|
|
import os
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from comfy.cli_args_types import Configuration
|
|
from comfy.cmd import folder_paths
|
|
from comfy.cmd.folder_paths import init_default_paths
|
|
from comfy.component_model.folder_path_types import FolderNames, ModelPaths
|
|
from comfy.execution_context import context_folder_names_and_paths
|
|
|
|
|
|
@pytest.fixture()
|
|
def clear_folder_paths():
|
|
# Clear the global dictionary before each test to ensure isolation
|
|
with context_folder_names_and_paths(FolderNames()):
|
|
yield
|
|
|
|
|
|
@pytest.fixture
|
|
def temp_dir():
|
|
with tempfile.TemporaryDirectory() as tmpdirname:
|
|
yield tmpdirname
|
|
|
|
|
|
def test_get_directory_by_type():
|
|
test_dir = "/test/dir"
|
|
folder_paths.set_output_directory(test_dir)
|
|
assert folder_paths.get_directory_by_type("output") == test_dir
|
|
assert folder_paths.get_directory_by_type("invalid") is None
|
|
|
|
|
|
def test_annotated_filepath():
|
|
assert folder_paths.annotated_filepath("test.txt") == ("test.txt", None)
|
|
assert folder_paths.annotated_filepath("test.txt [output]") == ("test.txt", folder_paths.get_output_directory())
|
|
assert folder_paths.annotated_filepath("test.txt [input]") == ("test.txt", folder_paths.get_input_directory())
|
|
assert folder_paths.annotated_filepath("test.txt [temp]") == ("test.txt", folder_paths.get_temp_directory())
|
|
|
|
|
|
def test_get_annotated_filepath():
|
|
default_dir = "/default/dir"
|
|
assert folder_paths.get_annotated_filepath("test.txt", default_dir) == os.path.join(default_dir, "test.txt")
|
|
assert folder_paths.get_annotated_filepath("test.txt [output]") == os.path.join(folder_paths.get_output_directory(), "test.txt")
|
|
|
|
|
|
def test_add_model_folder_path_append(clear_folder_paths):
|
|
folder_paths.add_model_folder_path("test_folder", "/default/path", is_default=True)
|
|
folder_paths.add_model_folder_path("test_folder", "/test/path", is_default=False)
|
|
assert folder_paths.get_folder_paths("test_folder") == ["/default/path", "/test/path"]
|
|
|
|
|
|
def test_add_model_folder_path_insert(clear_folder_paths):
|
|
folder_paths.add_model_folder_path("test_folder", "/test/path", is_default=False)
|
|
folder_paths.add_model_folder_path("test_folder", "/default/path", is_default=True)
|
|
assert folder_paths.get_folder_paths("test_folder") == ["/default/path", "/test/path"]
|
|
|
|
|
|
def test_add_model_folder_path_re_add_existing_default(clear_folder_paths):
|
|
folder_paths.add_model_folder_path("test_folder", "/test/path", is_default=False)
|
|
folder_paths.add_model_folder_path("test_folder", "/old_default/path", is_default=True)
|
|
assert folder_paths.get_folder_paths("test_folder") == ["/old_default/path", "/test/path"]
|
|
folder_paths.add_model_folder_path("test_folder", "/test/path", is_default=True)
|
|
assert folder_paths.get_folder_paths("test_folder") == ["/test/path", "/old_default/path"]
|
|
|
|
|
|
def test_add_model_folder_path_re_add_existing_non_default(clear_folder_paths):
|
|
folder_paths.add_model_folder_path("test_folder", "/test/path", is_default=False)
|
|
folder_paths.add_model_folder_path("test_folder", "/default/path", is_default=True)
|
|
assert folder_paths.get_folder_paths("test_folder") == ["/default/path", "/test/path"]
|
|
folder_paths.add_model_folder_path("test_folder", "/test/path", is_default=False)
|
|
assert folder_paths.get_folder_paths("test_folder") == ["/default/path", "/test/path"]
|
|
|
|
|
|
def test_filter_files_extensions():
|
|
files = ["file1.txt", "file2.jpg", "file3.png", "file4.txt"]
|
|
assert folder_paths.filter_files_extensions(files, [".txt"]) == ["file1.txt", "file4.txt"]
|
|
assert folder_paths.filter_files_extensions(files, [".jpg", ".png"]) == ["file2.jpg", "file3.png"]
|
|
assert folder_paths.filter_files_extensions(files, []) == files
|
|
|
|
|
|
def test_get_filename_list(temp_dir):
|
|
base_path = Path(temp_dir)
|
|
fn = FolderNames(base_paths=[base_path])
|
|
rel_path = Path("test/path")
|
|
fn.add(ModelPaths(["test_folder"], additional_relative_directory_paths=[rel_path], supported_extensions={".txt"}))
|
|
dir_path = base_path / rel_path
|
|
Path.mkdir(dir_path, parents=True, exist_ok=True)
|
|
files = ["file1.txt", "file2.jpg"]
|
|
|
|
for file in files:
|
|
Path.touch(dir_path / file, exist_ok=True)
|
|
|
|
with context_folder_names_and_paths(fn):
|
|
assert folder_paths.get_filename_list("test_folder") == ["file1.txt"]
|
|
|
|
|
|
def test_get_save_image_path(temp_dir):
|
|
with context_folder_names_and_paths(FolderNames(base_paths=[Path(temp_dir)])):
|
|
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path("test", temp_dir, 100, 100)
|
|
assert os.path.samefile(full_output_folder, temp_dir)
|
|
assert filename == "test"
|
|
assert counter == 1
|
|
assert subfolder == ""
|
|
assert filename_prefix == "test"
|
|
|
|
|
|
def test_add_output_path_absolute(temp_dir):
|
|
names = FolderNames()
|
|
config = Configuration()
|
|
config.cwd = str(temp_dir)
|
|
init_default_paths(names, config)
|
|
with context_folder_names_and_paths(names):
|
|
folder_paths.add_model_folder_path("diffusion_models", os.path.join(folder_paths.get_output_directory(), "diffusion_models"))
|
|
mp: ModelPaths = next(names.get_paths("diffusion_models"))
|
|
assert len(mp.additional_absolute_directory_paths) == 0
|
|
assert len(mp.additional_relative_directory_paths) == 1
|
|
assert list(mp.additional_relative_directory_paths)[0] == (Path("output") / "diffusion_models")
|