mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-08-02 15:04:50 +08:00
Allow changing folder_paths.base_path via command line argument. (#6600)
* Reimpl. CLI arg directly inside folder_paths. * Update tests to use CLI arg mocking. * Revert last-minute refactor. * Fix test state polution.
This commit is contained in:
@@ -1,19 +1,23 @@
|
||||
### 🗻 This file is created through the spirit of Mount Fuji at its peak
|
||||
# TODO(yoland): clean up this after I get back down
|
||||
import sys
|
||||
import pytest
|
||||
import os
|
||||
import tempfile
|
||||
from unittest.mock import patch
|
||||
from importlib import reload
|
||||
|
||||
import folder_paths
|
||||
import comfy.cli_args
|
||||
from comfy.options import enable_args_parsing
|
||||
enable_args_parsing()
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def clear_folder_paths():
|
||||
# Clear the global dictionary before each test to ensure isolation
|
||||
original = folder_paths.folder_names_and_paths.copy()
|
||||
folder_paths.folder_names_and_paths.clear()
|
||||
# Reload the module after each test to ensure isolation
|
||||
yield
|
||||
folder_paths.folder_names_and_paths = original
|
||||
reload(folder_paths)
|
||||
|
||||
@pytest.fixture
|
||||
def temp_dir():
|
||||
@@ -21,7 +25,21 @@ def temp_dir():
|
||||
yield tmpdirname
|
||||
|
||||
|
||||
def test_get_directory_by_type():
|
||||
@pytest.fixture
|
||||
def set_base_dir():
|
||||
def _set_base_dir(base_dir):
|
||||
# Mock CLI args
|
||||
with patch.object(sys, 'argv', ["main.py", "--base-directory", base_dir]):
|
||||
reload(comfy.cli_args)
|
||||
reload(folder_paths)
|
||||
yield _set_base_dir
|
||||
# Reload the modules after each test to ensure isolation
|
||||
with patch.object(sys, 'argv', ["main.py"]):
|
||||
reload(comfy.cli_args)
|
||||
reload(folder_paths)
|
||||
|
||||
|
||||
def test_get_directory_by_type(clear_folder_paths):
|
||||
test_dir = "/test/dir"
|
||||
folder_paths.set_output_directory(test_dir)
|
||||
assert folder_paths.get_directory_by_type("output") == test_dir
|
||||
@@ -96,3 +114,49 @@ def test_get_save_image_path(temp_dir):
|
||||
assert counter == 1
|
||||
assert subfolder == ""
|
||||
assert filename_prefix == "test"
|
||||
|
||||
|
||||
def test_base_path_changes(set_base_dir):
|
||||
test_dir = os.path.abspath("/test/dir")
|
||||
set_base_dir(test_dir)
|
||||
|
||||
assert folder_paths.base_path == test_dir
|
||||
assert folder_paths.models_dir == os.path.join(test_dir, "models")
|
||||
assert folder_paths.input_directory == os.path.join(test_dir, "input")
|
||||
assert folder_paths.output_directory == os.path.join(test_dir, "output")
|
||||
assert folder_paths.temp_directory == os.path.join(test_dir, "temp")
|
||||
assert folder_paths.user_directory == os.path.join(test_dir, "user")
|
||||
|
||||
assert os.path.join(test_dir, "custom_nodes") in folder_paths.get_folder_paths("custom_nodes")
|
||||
|
||||
for name in ["checkpoints", "loras", "vae", "configs", "embeddings", "controlnet", "classifiers"]:
|
||||
assert folder_paths.get_folder_paths(name)[0] == os.path.join(test_dir, "models", name)
|
||||
|
||||
|
||||
def test_base_path_change_clears_old(set_base_dir):
|
||||
test_dir = os.path.abspath("/test/dir")
|
||||
set_base_dir(test_dir)
|
||||
|
||||
assert len(folder_paths.get_folder_paths("custom_nodes")) == 1
|
||||
|
||||
single_model_paths = [
|
||||
"checkpoints",
|
||||
"loras",
|
||||
"vae",
|
||||
"configs",
|
||||
"clip_vision",
|
||||
"style_models",
|
||||
"diffusers",
|
||||
"vae_approx",
|
||||
"gligen",
|
||||
"upscale_models",
|
||||
"embeddings",
|
||||
"hypernetworks",
|
||||
"photomaker",
|
||||
"classifiers",
|
||||
]
|
||||
for name in single_model_paths:
|
||||
assert len(folder_paths.get_folder_paths(name)) == 1
|
||||
|
||||
for name in ["controlnet", "diffusion_models", "text_encoders"]:
|
||||
assert len(folder_paths.get_folder_paths(name)) == 2
|
||||
|
Reference in New Issue
Block a user